{"record":{"id":"066bd06e9ab538bd","repo":"sgl-project/sglang","slug":"g-cache-must-have-dtype-torch-float32","errorCode":null,"errorMessage":"`g_cache` must have dtype torch.float32.","messagePattern":"`g_cache` must have dtype torch\\.float32\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_replayssm.py","lineNumber":744,"sourceCode":"        raise ValueError(f\"`write_pos` must have shape {(batch,)}.\")\n    if force_flush is not None and (\n        force_flush.ndim != 1\n        or force_flush.dtype is not torch.int32\n        or force_flush.shape != (batch,)\n    ):\n        raise ValueError(\"`force_flush` must be a length-B int32 tensor or None.\")\n\n    cache_length = d_cache.size(2)\n    if cache_length < 1:\n        raise ValueError(\"ReplaySSM cache length must be at least 1.\")\n    if d_cache.shape[1:] != (num_v_heads, cache_length, value_dim):\n        raise ValueError(\"`d_cache` must have shape [slots, HV, L, V].\")\n    if k_cache.shape[1:] != (num_q_heads, cache_length, key_dim):\n        raise ValueError(\"`k_cache` must have shape [slots, H, L, K].\")\n    if g_cache.shape[1:] != (num_v_heads, cache_length, key_dim):\n        raise ValueError(\"`g_cache` must have shape [slots, HV, L, K].\")\n    if g_cache.dtype is not torch.float32:\n        raise ValueError(\"`g_cache` must have dtype torch.float32.\")\n\n    device = mixed_qkv.device\n    if any(\n        tensor.device != device for tensor in (d_cache, k_cache, g_cache, write_pos)\n    ):\n        raise ValueError(\"ReplaySSM inputs must be on the same device.\")\n    if force_flush is not None and force_flush.device != device:\n        raise ValueError(\"`force_flush` must be on the same device as the inputs.\")\n\n    cache_block = helion.next_power_of_2(max(16, cache_length))\n    use_lower_bound = lower_bound is not None\n    kernel = _select_replayssm_decode_kernel(\n        is_bf16_state=initial_state.dtype is torch.bfloat16,\n        num_v_heads=num_v_heads,\n    )\n    result = kernel(\n        mixed_qkv,\n        flat_a,","sourceCodeStart":726,"sourceCodeEnd":762,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py#L726-L762","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Allocate g_cache with dtype=torch.float32 explicitly","If loading from a bf16 checkpoint, call .to(torch.float32) on the gate cache before the first decode step","Centralize the dtype rule in the pool allocator: d/k follow model dtype, g is always fp32"],"exampleFix":"# before\ng_cache = torch.empty(slots, num_v_heads, cache_len, key_dim, dtype=torch.bfloat16, device='cuda')\n# after\ng_cache = torch.empty(slots, num_v_heads, cache_len, key_dim, dtype=torch.float32, device='cuda')","handlingStrategy":"type-guard","validationCode":"if g_cache.dtype is not torch.float32:\\n    g_cache = g_cache.float()","typeGuard":"def fp32_gate_cache(t: torch.Tensor) -> torch.Tensor:\\n    return t if t.dtype is torch.float32 else t.to(torch.float32)","tryCatchPattern":null,"preventionTips":["Never allocate g_cache with the model dtype; hardcode torch.float32","Check dtypes when loading checkpoints: cast gate caches to fp32 on load"],"tags":["helion","kda","replayssm","dtype","gate-cache"],"backgroundTag":"tensor-dtype-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}