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

The packed KDA (Kimi Delta Attention) decode validator requires mixed_qkv as a 2D [B, qkv_len] tensor because the Triton/Helion kernel iterates over a flattened batch of packed projections. Any other ndim is rejected before launching.

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_decode.py:229

def _is_power_of_two(value: int) -> bool:
    return value > 0 and value & (value - 1) == 0


def validate_packed_decode_inputs(
    mixed_qkv: torch.Tensor,
    a: torch.Tensor,
    b: torch.Tensor,
    A_log: torch.Tensor,
    dt_bias: torch.Tensor,
    initial_state: torch.Tensor,
    out: torch.Tensor,
    ssm_state_indices: torch.Tensor,
) -> tuple[int, int, int, int, int]:
    """Apply the shape and layout checks from SGLang's packed wrapper."""
    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(
            "`ssm_state_indices` must be 1D for packed decode "
            f"(got ndim={ssm_state_indices.ndim})."

View on GitHub (pinned to 0132848349)

Solutions

1. Flatten mixed_qkv to 2D: mixed_qkv.view(-1, mixed_qkv.shape[-1]) or ensure decode path packs [B, total_qkv_dim]
2. Check the caller's qkv packing step matches the kernel's expected layout

Example fix

# before
out = helion_fused_recurrent_kda_packed_decode(qkv_3d, ...)  # [B, seq, D]
# after
qkv = qkv_3d.reshape(qkv_3d.shape[0], -1)
out = helion_fused_recurrent_kda_packed_decode(qkv, ...)

assert mixed_qkv.ndim == 2, f'mixed_qkv must be 2D, got {mixed_qkv.ndim}'
mixed_qkv = mixed_qkv.view(mixed_qkv.shape[0], -1) if mixed_qkv.ndim > 2 else mixed_qkv

def is_packed_qkv(t: torch.Tensor) -> bool:
    return t.ndim == 2 and t.stride(-1) == 1

When it happens

Trigger: Calling helion_fused_recurrent_kda_packed_decode (or replayssm variant) with a 3D/4D qkv tensor (e.g. [B, seq, D] or [B, H, D, ...]).

Common situations: Passing a decode-shaped tensor that still has a sequence dim; upstream reorganize_qkv not flattening before decode; reusing prefill-style tensors.

Related errors

• `a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim=
• `A_log`/`dt_bias` must be 1D tensors.
• `ssm_state_indices` must be 1D for packed decode (got ndim={
• `mixed_qkv` must be contiguous in the last dim.
• `a`/`b` must be contiguous in the last dim.

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