sgl-project/sglang · error · ValueError

f"recurrent_kda needs a [N, HV, V, K] state pool; got shape

Error message

f"recurrent_kda needs a [N, HV, V, K] state pool; got shape {tuple(ssm_states.shape)}"

What it means

recurrent_kda expects the SSM state pool shaped [N, HV, V, K]; this ValueError fires when the pool passed to decode is not 4-D. The check runs once per tensor identity before the first decode call.

Source

Thrown at python/sglang/srt/layers/attention/linear/kernels/kda_flashinfer.py:105

    def _check_state_stride_contract(self, ssm_states: torch.Tensor) -> None:
        """One-time (per pool view) check that ``ssm_states`` matches the
        layout ``recurrent_kda`` was compiled for.

        The kernel's state argument is a CuTe fake tensor of shape
        ``[N, HV, V, K]`` with stride ``(sym_int64(divisibility=16), V*K, K, 1)``
        and ``assumed_align=32`` (flashinfer ``kda_kernels/recurrent_kda.py``):
        the slot stride is free — which is what lets the envelope-strided pools
        (unified memory / page-major layout, slot stride = per-slot envelope
        pitch) be passed in and updated IN PLACE on the cu_seqlens path — but
        the inner strides are compiled-in constants and the divisibility /
        alignment are hard assumptions. A pool violating them would mis-address
        state in-kernel without any error; fail loudly here instead.
        """
        key = id(ssm_states)
        if key in self._state_contract_ok:
            return
        if ssm_states.dim() != 4:
            raise ValueError(
                f"recurrent_kda needs a [N, HV, V, K] state pool; got "
                f"shape {tuple(ssm_states.shape)}"
            )
        _, hv, v, k = ssm_states.shape
        if ssm_states.stride()[1:] != (v * k, k, 1):
            raise ValueError(
                "recurrent_kda state inner strides must be compact "
                f"(V*K, K, 1)=({v * k}, {k}, 1); got {ssm_states.stride()[1:]} "
                "(only the slot stride may be non-compact)"
            )
        base_bytes = ssm_states.storage_offset() * ssm_states.element_size()
        if ssm_states.stride(0) % 16 != 0 or base_bytes % 32 != 0:
            raise ValueError(
                "recurrent_kda state pool breaks the compiled stride contract: "
                f"slot stride {ssm_states.stride(0)} elements must be a multiple "
                f"of 16 and the base byte offset {base_bytes} a multiple of 32 "
                "(sym_int64(divisibility=16) / assumed_align=32)"
            )

View on GitHub (pinned to 0132848349)

Solutions

1. Reshape/allocate the state pool as [num_slots, num_v_heads, head_v_dim, head_k_dim]
2. If the pool came from a cache refactor, verify the mamba cache allocation code path
3. Add a unit assert on ssm_states.dim()==4 before dispatching to decode

Example fix

# before
ssm_states = pool.view(-1, hv * v * k)  # flat
# after
ssm_states = pool.view(num_slots, hv, v, k)

assert ssm_states.dim() == 4, f'expected [N,HV,V,K], got {tuple(ssm_states.shape)}'
N, HV, V, K = ssm_states.shape
assert (HV, V, K) == (num_v_heads, head_v_dim, head_k_dim)

def is_valid_kda_state_pool(t: torch.Tensor) -> bool:
    return t.dim() == 4

try:
    kernel.decode(...)
except ValueError as e:
    if 'state pool' in str(e):
        ssm_states = ssm_states.reshape(num_slots, hv, v, k)
    else:
        raise

When it happens

Trigger: Calling FlashInferKDAKernel.decode (or a test harness) with an ssm_states tensor whose dim() != 4, e.g. a flat pool or a [B, T, H, V, K] 5-D buffer.

Common situations: Custom MambaCache / state-pool layouts that don't match SGLang's [N, HV, V, K] contract; refactors of the hybrid-attention cache that reshape the state pool.

Related errors

• `dt_bias` must have {HV * K} elements (got {dt_bias.numel()}
• f"recurrent_kda state inner strides must be compact (V*K, K,
• f"SGLang KDA state must be [B,H,V,K] with (V,K)={expected},
• NVIDIA KDA state must be [B,H,K,V] with (K,V)={expected}, go
• v_cache must be provided

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