sgl-project/sglang · error · ValueError

expected recurrent state layout [pool, H, V=128, K=128]

Error message

expected recurrent state layout [pool, H, V=128, K=128]

What it means

The recurrent state pool must be a 4D tensor [pool, H, 128, 128] with float32 dtype, contiguous inner strides (TILE_K*TILE_K, TILE_K, 1), pool stride 4-aligned, and 4-aligned storage offset — requirements of the kernel's vectorized state updates. The shape check [pool, H, V=128, K=128] fails otherwise.

Source

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

    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)
    if any(ring is not None for ring in rings) and not cache_ring:
        raise ValueError("ReplaySSM requires all four replayssm_* rings")
    if cache_ring:
        ring_len = replayssm_rawv.shape[2]
        if (
            ring_len < 1 + num_spec

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate recurrent_state as torch.float32 with shape [pool_size, H, 128, 128], contiguous
2. If it is a view/slice, call .contiguous() and ensure the storage offset is 4-element aligned (allocate fresh rather than slicing)
3. For head dims != 128, use the generic KDA decode backend instead of this fused MTP kernel

Example fix

# before
state = torch.zeros(pool, H, D, D, dtype=x.dtype, device='cuda')  # bf16, D!=128
# after
state = torch.zeros(pool, H, 128, 128, dtype=torch.float32, device='cuda')

assert recurrent_state.shape[1:] == (H, 128, 128)
assert recurrent_state.dtype == torch.float32 and recurrent_state.is_contiguous()

def is_valid_state(s: torch.Tensor, H: int) -> bool:
    return (s.dtype == torch.float32 and s.shape[1:] == (H, 128, 128)
            and tuple(s.stride()[-3:]) == (128*128, 128, 1))

When it happens

Trigger: Passing recurrent_state with shape [pool, H, 64, 64] (head_dim 64 model), a bf16/fp16 state, a transposed or non-contiguous view, or a slice with unaligned storage offset.

Common situations: Allocating the Mamba/KDA state pool in model dtype instead of fp32; sharing a state cache across models with different head dims; passing state.transpose(-1, -2) views.

Related errors

• expected q/k shape {expected_shape}
• expected v/g shape {expected_shape}
• expected beta shape {(1, T, H)}
• DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens pe
• 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/eaa2dc5f3ffb7390. Report an issue: GitHub.