sgl-project/sglang · error · ValueError

`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim=

Error message

`a` and `b` must be 2D tensors (got a.ndim={a.ndim}, b.ndim={b.ndim}).

What it means

The KDA packed decode kernel requires the gated delta-rule a and b tensors as 2D [B, dim] tensors. Passing tensors with any other rank fails validation before kernel launch.

Source

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

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})."
        )
    if not out.is_contiguous():
        raise ValueError("`out` must be contiguous.")

    device = mixed_qkv.device
    if any(

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape a and b to 2D matching mixed_qkv's batch size
2. Verify the host wrapper producing a/b emits [B, D]

Example fix

# before
out = decode(qkv, a_3d, b_3d, ...)
# after
out = decode(qkv, a_3d.reshape(a_3d.shape[0], -1), b_3d.reshape(b_3d.shape[0], -1), ...)

assert a.ndim == 2 and b.ndim == 2, f'a/b must be 2D, got {a.ndim}, {b.ndim}'

def are_2d(*ts: torch.Tensor) -> bool:
    return all(t.ndim == 2 for t in ts)

When it happens

Trigger: Calling packed decode with a/b of ndim != 2, e.g. [B, seq, d] or [B, n_groups, d].

Common situations: Using prefill-shaped a/b tensors during decode; forgetting to squeeze a seq or head dim after scheduling.

Related errors

• `mixed_qkv` must be a 2D tensor (got ndim={mixed_qkv.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/d7d025016b173fc0. Report an issue: GitHub.