sgl-project/sglang · error · ValueError

DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens pe

Error message

DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens per request; got T={T}, N={N}

What it means

The MTP variant of the kernel requires every request to carry exactly the same dense token count: T (flattened tokens) must be divisible by N (number of requests from cu_seqlens), giving 1 + num_spec tokens each. Violations (T % N != 0, N <= 0, or ragged requests) raise this error.

Source

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

    """
    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
        or tuple(recurrent_state.stride()[-3:]) != (TILE_K * TILE_K, TILE_K, 1)
        or recurrent_state.stride(0) % 4 != 0
        or recurrent_state.storage_offset() % 4 != 0
    ):
        raise ValueError(
            "cp.async recurrent state requires fp32 contiguous [H, V, K] "
            "inner layout and 16-byte-aligned slot offsets"
        )
    rings = (replayssm_rawv, replayssm_rawk, replayssm_g, replayssm_beta)
    cache_ring = all(ring is not None for ring in rings)

View on GitHub (pinned to 0132848349)

Solutions

1. Ensure all requests have identical token counts (1 + num_spec) and T == N * (1 + num_spec)
2. Pad requests with fewer draft tokens or route ragged batches to a non-MTP KDA decode path
3. Validate cu_seqlens: monotonically increasing, N >= 1, differences all equal

Example fix

# before
fused_kda_decode_mtp_dspark(q, k, v, g, beta, cu_seqlens, state)  # ragged lengths
# after
assert cu_seqlens.numel() >= 2 and T % (cu_seqlens.numel() - 1) == 0
# or route ragged batches elsewhere:
if T % N != 0:
    return kda_decode_fallback(...)

N = cu_seqlens.numel() - 1
assert N > 0 and T % N == 0 and T // N >= 1

def uniform_mtp_batch(cu_seqlens: torch.Tensor) -> bool:
    lens = cu_seqlens[1:] - cu_seqlens[:-1]
    return lens.numel() > 0 and bool((lens == lens[0]).all())

When it happens

Trigger: Calling fused_kda_decode_mtp_dspark where cu_seqlens segments have unequal lengths (ragged speculative draft counts), N=0 (empty cu_seqlens), or T not a multiple of N.

Common situations: Mixed num_spec per request after dynamic speculative-length scheduling; a cu_seqlens built for plain decode concatenated with MTP drafts; off-by-one in cu_seqlens producing an extra empty segment.

Related errors

• expected q/k shape {expected_shape}
• expected v/g shape {expected_shape}
• expected beta shape {(1, T, H)}
• expected recurrent state layout [pool, H, V=128, K=128]
• Unexpected A_log shape: {A_log.shape}; expected numel={HV}

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