sgl-project/sglang · error · ValueError

`initial_state` must be a 4D tensor (got ndim={initial_state

Error message

`initial_state` must be a 4D tensor (got ndim={initial_state.ndim}).

What it means

Helion KDA decode validates that the recurrent SSM state tensor has 4 dimensions. initial_state is the [B, HV, V, K] (or pool-shaped) state pool read at decode time; a tensor with any other rank cannot be indexed by the fused recurrent kernel, so validate_packed_decode_inputs rejects it up front before launching.

Source

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

        )
    ):
        raise ValueError("All inputs must be on the same device.")

    B = mixed_qkv.shape[0]
    if a.shape[0] != B or b.shape[0] != B:
        raise ValueError(
            "Mismatched batch sizes: "
            f"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, "
            f"b.shape[0]={b.shape[0]}."
        )
    if ssm_state_indices.shape[0] != B:
        raise ValueError(
            f"`ssm_state_indices` must have shape [B] "
            f"(got {tuple(ssm_state_indices.shape)}; expected ({B},))."
        )

    if initial_state.ndim != 4:
        raise ValueError(
            f"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim})."
        )
    if initial_state.stride(-1) != 1:
        raise ValueError("`initial_state` must be contiguous in the last dim.")
    HV, V, K = initial_state.shape[-3:]
    if not _is_power_of_two(K) or not _is_power_of_two(V):
        raise ValueError(
            "Helion KDA decode requires power-of-two key and value head "
            f"dimensions (got K={K}, V={V})."
        )
    if a.shape[1] != HV * K:
        raise ValueError(
            f"`a` must have shape [B, HV*K] with HV={HV}, K={K} "
            f"(got a.shape={tuple(a.shape)})."
        )
    if b.shape[1] != HV:
        raise ValueError(
            f"`b` must have shape [B, HV] with HV={HV} (got b.shape={tuple(b.shape)})."

View on GitHub (pinned to 0132848349)

Solutions

Reshape initial_state to 4D: [batch_or_slots, HV, V, K] matching the state pool layout
Check where initial_state is produced (state cache / mamba-style pool) and verify its rank before the decode call
Add an assert initial_state.ndim == 4 in your wrapper so failures surface at the producer, not the kernel

Example fix

// before
state = pool.get_state(idx)  # accidentally [HV, V, K]
out = helion_fused_recurrent_kda_packed_decode(qkv, a, b, ..., state, ...)
// after
state = pool.get_state(idx)
assert state.ndim == 4, state.shape  # [slots, HV, V, K]
out = helion_fused_recurrent_kda_packed_decode(qkv, a, b, ..., state, ...)

Defensive patterns

Strategy: validation

Validate before calling

assert initial_state.ndim == 4, f"expected 4D state pool, got {initial_state.shape}"

Type guard

def is_valid_state_pool(t: torch.Tensor) -> bool:
    return isinstance(t, torch.Tensor) and t.ndim == 4

Prevention

Always build initial_state via torch.zeros/zeros_like of shape [slots, HV, V, K]
Add ndim asserts in wrapper code near the state-pool producer

When it happens

Trigger: Calling helion_fused_recurrent_kda_packed_decode or helion_fused_recurrent_kda_replayssm_decode with an initial_state that is 2D/3D/5D — e.g. passing a per-head state [HV, V, K] instead of the batched 4D pool, or a flattened vector.

Common situations: Migrating a single-request decode loop to the packed batch API and forgetting to add the batch/pool dim; passing a state pool reshaped by .view() with the wrong rank; wiring a ReplaySSM checkpoint tensor of the wrong shape into the decode call.

Related errors

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