sgl-project/sglang · error · ValueError

`g_cache` must have dtype torch.float32.

Error message

`g_cache` must have dtype torch.float32.

What it means

The ReplaySSM decode kernel requires the gate cache g_cache to be float32. Unlike d/k caches which follow the model dtype, the decay/gate cache must be float32 because the kernel accumulates exponential decay products that overflow or lose precision in bf16/fp16. The wrapper hard-checks g_cache.dtype is torch.float32 before launching.

Source

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

        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,
    )
    result = kernel(
        mixed_qkv,
        flat_a,

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate g_cache with dtype=torch.float32 explicitly
2. If loading from a bf16 checkpoint, call .to(torch.float32) on the gate cache before the first decode step
3. Centralize the dtype rule in the pool allocator: d/k follow model dtype, g is always fp32

Example fix

# before
g_cache = torch.empty(slots, num_v_heads, cache_len, key_dim, dtype=torch.bfloat16, device='cuda')
# after
g_cache = torch.empty(slots, num_v_heads, cache_len, key_dim, dtype=torch.float32, device='cuda')

if g_cache.dtype is not torch.float32:\n    g_cache = g_cache.float()

def fp32_gate_cache(t: torch.Tensor) -> torch.Tensor:\n    return t if t.dtype is torch.float32 else t.to(torch.float32)

When it happens

Trigger: Calling helion_fused_recurrent_kda_replayssm_decode with g_cache in bfloat16/float16 (e.g. allocated with the model dtype for uniformity, or read from a checkpoint stored in bf16).

Common situations: Allocating the whole hybrid state pool in the model dtype (bf16) to save memory; loading g_cache weights from a bf16 checkpoint without casting; refactoring a pool allocator that used a single dtype parameter.

Related errors

• `g_cache` must have shape [slots, HV, L, K].
• `write_pos` must be a 1D int32 tensor.
• `force_flush` must be a length-B int32 tensor or None.
• `d_cache` must have shape [slots, HV, L, V].
• `k_cache` must have shape [slots, H, L, K].

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