sgl-project/sglang · error · ValueError

`d_cache` must have shape [slots, HV, L, V].

Error message

`d_cache` must have shape [slots, HV, L, V].

What it means

The helion_fused_recurrent_kda_replayssm_decode kernel requires the value-state cache d_cache to be laid out [slots, HV, L, V] where HV=num_v_heads, L=d_cache.size(2) and V=value_dim. The wrapper validates d_cache.shape[1:] against (num_v_heads, cache_length, value_dim) before launching the Helion kernel, because the generated Triton kernel indexes the cache with those exact strides. Any mismatch (wrong head grouping, transposed dims, or stale cache built for a different value_dim) aborts with this ValueError.

Source

Thrown at python/sglang/kernels/ops/attention/helion/kda_replayssm.py:738

        ssm_state_indices,
    )

    if write_pos.ndim != 1 or write_pos.dtype is not torch.int32:
        raise ValueError("`write_pos` must be a 1D int32 tensor.")
    if write_pos.shape != (batch,):
        raise ValueError(f"`write_pos` must have shape {(batch,)}.")
    if force_flush is not None and (
        force_flush.ndim != 1
        or force_flush.dtype is not torch.int32
        or force_flush.shape != (batch,)
    ):
        raise ValueError("`force_flush` must be a length-B int32 tensor or None.")

    cache_length = d_cache.size(2)
    if cache_length < 1:
        raise ValueError("ReplaySSM cache length must be at least 1.")
    if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):
        raise ValueError("`d_cache` must have shape [slots, HV, L, V].")
    if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):
        raise ValueError("`k_cache` must have shape [slots, H, L, K].")
    if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):
        raise ValueError("`g_cache` must have shape [slots, HV, L, K].")
    if g_cache.dtype is not torch.float32:
        raise ValueError("`g_cache` must have dtype torch.float32.")

    device = mixed_qkv.device
    if any(
        tensor.device != device for tensor in (d_cache, k_cache, g_cache, write_pos)
    ):
        raise ValueError("ReplaySSM inputs must be on the same device.")
    if force_flush is not None and force_flush.device != device:
        raise ValueError("`force_flush` must be on the same device as the inputs.")

    cache_block = helion.next_power_of_2(max(16, cache_length))
    use_lower_bound = lower_bound is not None
    kernel = _select_replayssm_decode_kernel(

View on GitHub (pinned to 0132848349)

Solutions

  1. Check d_cache.shape[1:] vs (num_v_heads, d_cache.size(2), value_dim) and reallocate the cache with torch.empty(slots, num_v_heads, cache_length, value_dim, ...)
  2. If you transposed the cache for a custom kernel, call .permute/.contiguous back to [slots, HV, L, V] before invoking the op
  3. Verify num_v_heads/value_dim passed to the call match the model config used to size the state pool
  4. In multi-GPU runs, make sure the cache was sharded on the head dimension consistently with the other caches

Example fix

# before
d_cache = torch.empty(slots, cache_len, num_v_heads, value_dim, dtype=torch.float32, device='cuda')
# after
d_cache = torch.empty(slots, num_v_heads, cache_len, value_dim, dtype=torch.float32, device='cuda')
Defensive patterns

Strategy: validation

Validate before calling

assert d_cache.ndim == 4 and d_cache.shape[1:] == (num_v_heads, d_cache.size(2), value_dim), (d_cache.shape, num_v_heads, value_dim)

Type guard

def valid_d_cache(t: torch.Tensor, hv: int, v: int) -> bool:
    return t.ndim == 4 and t.shape[1] == hv and t.shape[3] == v

Try / catch

try:\n    helion_fused_recurrent_kda_replayssm_decode(...)\nexcept ValueError as e:\n    if 'd_cache' in str(e):\n        raise RuntimeError(f'state pool layout bug: {e}') from e\n    raise

Prevention

When it happens

Trigger: Calling helion_fused_recurrent_kda_replayssm_decode with a d_cache tensor whose trailing dims are not exactly (num_v_heads, d_cache.size(2), value_dim) — e.g. passing a [slots, L, HV, V] cache, a cache allocated for a different num_v_heads, or a cache created with head_dim=value_dim*2 by mistake.

Common situations: Hybrid KDA models where the ReplaySSM state pool is allocated separately from the KV pool and its shape constants (num_v_heads, key_dim/value_dim) drift from the model config; tests (test_replayssm_decode_contract) that build caches with ad-hoc shapes; TP sharding that splits HV heads but reuses a single-GPU cache layout.

Related errors

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