sgl-project/sglang · error · ValueError

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

Error message

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

What it means

fused_recurrent_gated_delta_rule_packed_decode requires mixed_qkv as a 2D packed tensor of shape (num_tokens, qk_dim + hv*v) — one row per decode token, Q/K/V fused along dim 1. It raises when the tensor has any other rank, e.g. the unpacked (B, T, ...) layout or separate q/k/v tensors.

Source

Thrown at python/sglang/kernels/ops/attention/fla/fused_recurrent.py:281

    p_ht = ht + state_idx * stride_final_state_token
    p_ht = p_ht + i_hv * V * K + o_v[:, None] * K + o_k[None, :]
    tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)


def fused_recurrent_gated_delta_rule_packed_decode(
    mixed_qkv: torch.Tensor,
    a: torch.Tensor,
    b: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    scale: float,
    initial_state: torch.Tensor,
    out: torch.Tensor,
    ssm_state_indices: torch.Tensor,
    use_qk_l2norm_in_kernel: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    if mixed_qkv.ndim != 2:
        raise ValueError(
            f"`mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.ndim})."
        )
    if mixed_qkv.stride(-1) != 1:
        raise ValueError("`mixed_qkv` must be contiguous in the last dim.")
    if a.ndim != 2 or b.ndim != 2:
        raise ValueError(
            f"`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim})."
        )
    if a.stride(-1) != 1 or b.stride(-1) != 1:
        raise ValueError("`a`/`b` must be contiguous in the last dim.")
    if A_log.ndim != 1 or dt_bias.ndim != 1:
        raise ValueError("`A_log`/`dt_bias` must be 1D tensors.")
    if A_log.stride(0) != 1 or dt_bias.stride(0) != 1:
        raise ValueError("`A_log`/`dt_bias` must be contiguous.")
    if ssm_state_indices.ndim != 1:
        raise ValueError(
            f"`ssm_state_indices` must be 1D for packed decode (got ndim={ssm_state_indices.ndim})."
        )

View on GitHub (pinned to 0132848349)

Solutions

  1. Reshape the fused projection to 2D: mixed_qkv.view(num_tokens, -1) or .reshape(-1, qkv_dim)
  2. Ensure you are calling the packed decode entry point with the packed argument convention (mixed_qkv, a, b, A_log, dt_bias, initial_state, out, ssm_state_indices), not separate q/k/v

Example fix

# before
out, state = fused_recurrent_gated_delta_rule_packed_decode(qkv, ...)  # qkv is (B, 1, D)
# after
out, state = fused_recurrent_gated_delta_rule_packed_decode(qkv.view(-1, qkv.shape[-1]), ...)
Defensive patterns

Strategy: validation

Validate before calling

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

Type guard

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

Prevention

When it happens

Trigger: Passing a 3D/4D projection output (B, 1, qkv_dim) or (B, T, H, D) instead of a squeezed (tokens, qkv_dim) 2D tensor; passing q, k, v separately instead of the fused mixed_qkv the packed decode API expects.

Common situations: Migrating from the non-packed fused_recurrent API (which takes separate q/k/v with (B, H, T, D) shapes) to the packed decode variant used by sglang's hybrid attention; forgetting to .view(-1, qkv_dim) after a qkv Linear projection; benchmark/test code reusing 4D tensors.

Related errors

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