sgl-project/sglang · error · NotImplementedError

KDA cutedsl: safe_gate (lower_bound) not yet supported

Error message

KDA cutedsl: safe_gate (lower_bound) not yet supported

What it means

The CUTLASS-DSL chunked KDA prefill kernel (chunk_kda_cutedsl) computes the gate as -exp(A_log)*softplus(g+dt_bias) in a PyTorch prologue pass, and that formulation has no support for the numerically-safe gate variant (lower_bound / safe_gate, which clamps the gate to avoid -inf decay). Passing both A_log and lower_bound is therefore explicitly rejected with NotImplementedError until the gate fusion TODO is done.

Source

Thrown at python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py:147

    Pool mode: ``h0`` is the state pool [num_slots, Hv, V, K] and ``h0_indices``
    maps each sequence to its slot; the h kernel reads AND writes the pool rows
    in place (fused state gather/scatter — no [N, Hv, V, K] intermediates), and
    the returned ``ht`` is the pool tensor itself.
    """
    import torch.nn.functional as F

    T, Hv, K = q.shape
    V = v.shape[-1]
    if scale is None:
        scale = K**-0.5
    if num_sms is None:
        num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count

    # Gate activation (standard KDA gate). Fused into the prologue is a B2 TODO;
    # for now a small PyTorch pass, matching chunk_kda's kda_gate_chunk_cumsum.
    if A_log is not None:
        if lower_bound is not None:
            raise NotImplementedError(
                "KDA cutedsl: safe_gate (lower_bound) not yet supported"
            )
        x = g.float()
        if dt_bias is not None:
            x = x + dt_bias.float().view(1, Hv, K)
        g_act = -torch.exp(A_log.float()).view(1, Hv, 1) * F.softplus(x)
    else:
        g_act = g.float()

    # Reusable scratch (eye/pack/U/W/V_new/h_chunks) + cached metadata; only the
    # returned o/ht are freshly allocated. This removes the ~0.2-0.6ms/call host
    # overhead (re-alloc + re-zero of ~200MB + metadata sync) that otherwise drags
    # the (fast) cutedsl kernels below Triton.
    ws, chunk_indices, chunk_offsets, total_chunks, total, pad_t = _kda_workspace(
        q, T, Hv, K, V, cu_seqlens
    )

    # KL/qg2 from the prologue fold the decay with a chunk-global g_last reference

View on GitHub (pinned to 0132848349)

Solutions

1. Drop lower_bound (pass None) so the standard gate path is used — only valid if the checkpoint tolerates it
2. Use the non-cutedsl chunk_kda implementation (Triton path) which supports safe_gate, by disabling the cutedsl backend/flag
3. Wait for or implement safe_gate support in the cutedsl gate prologue (the B2 TODO noted in the code)

Example fix

# before
o, final_state = chunk_kda_cutedsl(q, k, v, g, A_log=A_log, lower_bound=0.1, dt_bias=dt_bias)
# after
o, final_state = chunk_kda(q, k, v, g, A_log=A_log, lower_bound=0.1, dt_bias=dt_bias)  # Triton path supports safe_gate

if A_log is not None and lower_bound is not None:\n    # cutedsl path lacks safe_gate support; use the Triton chunk_kda path
    out = chunk_kda(q, k, v, g, A_log=A_log, lower_bound=lower_bound, ...)
else:\n    out = chunk_kda_cutedsl(...)

try:\n    out = chunk_kda_cutedsl(q, k, v, g, A_log=A_log, lower_bound=lb)\nexcept NotImplementedError:\n    out = chunk_kda(q, k, v, g, A_log=A_log, lower_bound=lb)

When it happens

Trigger: Calling chunk_kda_cutedsl(q, k, v, g, A_log=..., lower_bound=0.1, ...) — i.e. a KDA model configured with a safe-gate lower bound (e.g. KDA-2 style checkpoints) on the Blackwell CUTLASS path.

Common situations: Serving a KDA model whose config sets a nonzero lower_bound while the runtime selects the cutedsl kernel; tests like test_kda_chunk_cutedsl_realistic_gate that probe safe-gate support; Blackwell GPU routing that prefers the CUTLASS implementation over the Triton one.

Related errors

• Unexpected A_log shape: {A_log.shape}; expected numel={HV}
• Unexpected dt_bias shape: {dt_bias.shape}; expected numel={H
• eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T
• CuteDSLKDAKernel does not support target_verify
• Unexpected initial_state_source shape: {initial_state_source

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/39247e10efc2db11. Report an issue: GitHub.