sgl-project/sglang · error · ValueError

`k_cache` must have shape [slots, H, L, K].

Error message

`k_cache` must have shape [slots, H, L, K].

What it means

The ReplaySSM decode kernel requires k_cache shaped [slots, H, L, K] with H=num_q_heads (note: the key cache is grouped by query heads, not value heads) and K=key_dim, with L matching d_cache.size(2). The wrapper checks k_cache.shape[1:] != (num_q_heads, cache_length, key_dim) because the Triton kernel loads K with those strides for the chunked recurrence. A mismatch means the key cache was allocated or grouped inconsistently with the other ReplaySSM caches.

Source

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

    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(
        is_bf16_state=initial_state.dtype is torch.bfloat16,
        num_v_heads=num_v_heads,

View on GitHub (pinned to 0132848349)

Solutions

  1. Allocate k_cache as [slots, num_q_heads, cache_length, key_dim] with cache_length == d_cache.size(2)
  2. Confirm num_q_heads (H) is used for k_cache while num_v_heads (HV) is used for d_cache/g_cache — they differ in GQA-style KDA models
  3. If the cache came from a pool, check the pool's per-layer shape metadata against the model config

Example fix

# before
k_cache = torch.empty(slots, num_v_heads, cache_len, key_dim, dtype=dtype, device='cuda')
# after
k_cache = torch.empty(slots, num_q_heads, cache_len, key_dim, dtype=dtype, device='cuda')
Defensive patterns

Strategy: validation

Validate before calling

assert k_cache.shape[1:] == (num_q_heads, d_cache.size(2), key_dim), k_cache.shape

Type guard

def valid_k_cache(t: torch.Tensor, h: int, k: int, l: int) -> bool:
    return t.ndim == 4 and t.shape[1:] == (h, l, k)

Prevention

When it happens

Trigger: Calling helion_fused_recurrent_kda_replayssm_decode where k_cache.shape[1:] != (num_q_heads, d_cache.size(2), key_dim) — e.g. passing a k_cache with HV head grouping (same as g_cache/d_cache) instead of H grouping, or a cache whose length differs from d_cache's.

Common situations: Reusing one allocation loop for all three caches (d/k/g) and accidentally using num_v_heads for k_cache; mixed-dimming hybrid attention pools; TP runs where the k cache was sharded on a different head axis than expected.

Related errors

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