sgl-project/sglang · error · ValueError

expected q/k shape {expected_shape}

Error message

expected q/k shape {expected_shape}

What it means

fused_kda_decode_mtp_dspark validates that q and k are dense [1, T, H, TILE_K] (TILE_K=128) tensors — batch 1, all tokens flattened into dim 1, heads H, head dim 128. A ragged/varlen layout or wrong head dim fails this check before the DSlab kernel launches.

Source

Thrown at python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py:987

    --speculative-dspark-block-size), inferred here from T // N - 1.

    ReplaySSM: passing the four replayssm_* rings switches the kernel to
    CACHE_RING mode — per-step raw inputs go to the rings (consumed by the
    commit-time exact fold, see kda_replayssm_spec_decode.py) and the per-step
    intermediate_ssm state snapshots are skipped, so intermediate_ssm may be
    None.

    Passing all three onorm_* arguments fuses gated RMSNorm into the recurrence
    kernel.
    """
    import torch

    H = x_q.shape[2]
    N = cu_seqlens.numel() - 1
    T = x_q.shape[1]
    expected_shape = (1, T, H, TILE_K)
    if tuple(x_q.shape) != expected_shape or tuple(x_k.shape) != expected_shape:
        raise ValueError(f"expected q/k shape {expected_shape}")
    if tuple(x_v.shape) != expected_shape or tuple(g.shape) != expected_shape:
        raise ValueError(f"expected v/g shape {expected_shape}")
    if tuple(beta.shape) != (1, T, H):
        raise ValueError(f"expected beta shape {(1, T, H)}")
    # T // N == 1 is num_spec == 0: one token per request, i.e. a plain decode
    # step. The backend never dispatches here for it (that is the dedicated
    # decode kernel's job), but the layout is legal and benchmarks compare the
    # two at this point, so the wrapper accepts it.
    if N <= 0 or T % N != 0 or T // N < 1:
        raise ValueError(
            f"DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens per "
            f"request; got T={T}, N={N}"
        )
    num_spec = T // N - 1
    if recurrent_state.shape[1:] != (H, TILE_K, TILE_K):
        raise ValueError("expected recurrent state layout [pool, H, V=128, K=128]")
    if (
        recurrent_state.dtype != torch.float32

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape q/k to [1, T, H, 128] where T is the flattened token count
2. Verify head dim is exactly TILE_K=128; for other head dims use the generic KDA decode backend
3. Ensure batch size is 1 (all requests concatenated along dim 1), matching cu_seqlens

Example fix

# before
q = q.reshape(num_tokens, H, 128)  # 3D -> fails
fused_kda_decode_mtp_dspark(q, k, ...)
# after
q = q.reshape(1, num_tokens, H, 128)
k = k.reshape(1, num_tokens, H, 128)
fused_kda_decode_mtp_dspark(q, k, ...)

T = cu_seqlens[-1].item()
assert x_q.shape == (1, T, H, 128) and x_k.shape == (1, T, H, 128)

def is_dspark_qk(x: torch.Tensor, T: int, H: int) -> bool:
    return tuple(x.shape) == (1, T, H, 128)

When it happens

Trigger: Passing x_q or x_k with a batch dim != 1, a 3D layout, a head dim != TILE_K (128), or varlen cu_seqlens-packed layout instead of the dense flattened one the wrapper requires.

Common situations: Reusing a varlen [1, total_tokens, H, D] path with per-request 4D tensors; models with head_dim != 128 (e.g. 64) routed to this kernel; accidental transpose.

Related errors

• expected v/g shape {expected_shape}
• expected beta shape {(1, T, H)}
• Unexpected A_log shape: {A_log.shape}; expected numel={HV}
• Unexpected dt_bias shape: {dt_bias.shape}; expected numel={H
• Unexpected a shape for varlen: {a.shape}

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