sgl-project/sglang · error · ValueError

`mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).

Error message

`mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).

What it means

fused_recurrent_linear_replayssm_decode validates its inputs before launching: mixed_qkv must be a 2D [num_tokens, qkv_dim] tensor (one row per decode token). Passing 3D/4D tensors (e.g. [B, T, D] or [B, 1, T, D]) trips this check.

Source

Thrown at python/sglang/kernels/ops/attention/fla/fused_recurrent_linear_replayssm.py:478

    One kernel for both gate granularities, selected by ``is_kda``:
      * ``is_kda=False`` (GDN): per-head SCALAR gate.  ``a``=[B, HV],
        ``dt_bias``=[HV], ``g_cache``=[num_slots, HV, L].
      * ``is_kda=True`` (KDA): per-K-channel gate.  ``a``=[B, HV, K],
        ``dt_bias``=[HV, K], ``g_cache``=[num_slots, HV, L, K].
    ``A_log`` is [HV] (per-head scalar) for both.

    Same call surface as the packed decode plus the three ring caches
    (``d_cache`` / ``k_cache`` / ``g_cache``) and the per-decode-row
    ``write_pos`` cursor.  ``initial_state`` is both the checkpoint read (h0)
    and the (flush-only) checkpoint write (ht), in place.

    Allocates nothing persistent: the caller owns the ring tensors and is
    responsible for advancing / resetting ``write_pos`` (e.g. ``(write_pos+1) %
    L`` after each step).  This is a STANDALONE kernel; the memory-pool / cache
    integration is a later phase.
    """
    if mixed_qkv.ndim != 2:
        raise ValueError(f"`mixed_qkv` must be 2D (got ndim={mixed_qkv.ndim}).")
    if mixed_qkv.stride(-1) != 1:
        raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
    if b.ndim != 2:
        raise ValueError(f"`b` must be 2D (got b.ndim={b.ndim}).")
    if A_log.ndim != 1:
        raise ValueError("`A_log` must be a 1D tensor.")
    if initial_state.ndim != 4:
        raise ValueError(f"`initial_state` must be 4D (got ndim={initial_state.ndim}).")
    if not out.is_contiguous():
        raise ValueError("`out` must be contiguous.")
    if write_pos.ndim != 1 or write_pos.dtype != torch.int32:
        raise ValueError("`write_pos` must be a 1D int32 tensor.")
    if force_flush is not None and (
        force_flush.ndim != 1 or force_flush.dtype != torch.int32
    ):
        raise ValueError("`force_flush` must be a 1D int32 tensor or None.")

    B = mixed_qkv.shape[0]

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape mixed_qkv to 2D: mixed_qkv.reshape(-1, qkv_dim) before calling
2. Keep all per-token tensors 2D [num_tokens, dim] on this path
3. Check b is 2D and initial_state is 4D to satisfy the sibling checks

Example fix

// before
out = fused_recurrent_linear_replayssm_decode(mixed_qkv[B,1,D], ...)
// after
out = fused_recurrent_linear_replayssm_decode(mixed_qkv.reshape(-1, D), ...)

mixed_qkv = mixed_qkv.reshape(-1, mixed_qkv.shape[-1])
assert mixed_qkv.ndim == 2

def is_2d_tokens(t: torch.Tensor) -> bool:
    return t.ndim == 2

When it happens

Trigger: Feeding an un-squeezed decode projection of shape [B, 1, D] or a prefill-shaped [B, T, D] tensor into the standalone replaySSM decode kernel.

Common situations: Adapting prefill batch code to the decode path; forgetting to reshape after a fused QKV projection in a new model integration; CUDA-graph capture tests passing batched dims.

Related errors

• `dt_bias` must have {HV * K} elements (got {dt_bias.numel()}
• `b` must be 2D (got b.ndim={b.ndim}).
• `A_log` must be a 1D tensor.
• `initial_state` must be 4D (got ndim={initial_state.ndim}).
• pairs must be a torch.Tensor of shape [N, 2]

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