{"record":{"id":"47bf1c97166f5598","repo":"sgl-project/sglang","slug":"d-cache-must-have-shape-slots-hv-l-v","errorCode":null,"errorMessage":"`d_cache` must have shape [slots, HV, L, V].","messagePattern":"`d_cache` must have shape \\[slots, HV, L, V\\]\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_replayssm.py","lineNumber":738,"sourceCode":"        ssm_state_indices,\n    )\n\n    if write_pos.ndim != 1 or write_pos.dtype is not torch.int32:\n        raise ValueError(\"`write_pos` must be a 1D int32 tensor.\")\n    if write_pos.shape != (batch,):\n        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(","sourceCodeStart":720,"sourceCodeEnd":756,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py#L720-L756","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["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, ...)","If you transposed the cache for a custom kernel, call .permute/.contiguous back to [slots, HV, L, V] before invoking the op","Verify num_v_heads/value_dim passed to the call match the model config used to size the state pool","In multi-GPU runs, make sure the cache was sharded on the head dimension consistently with the other caches"],"exampleFix":"# before\nd_cache = torch.empty(slots, cache_len, num_v_heads, value_dim, dtype=torch.float32, device='cuda')\n# after\nd_cache = torch.empty(slots, num_v_heads, cache_len, value_dim, dtype=torch.float32, device='cuda')","handlingStrategy":"validation","validationCode":"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)","typeGuard":"def valid_d_cache(t: torch.Tensor, hv: int, v: int) -> bool:\n    return t.ndim == 4 and t.shape[1] == hv and t.shape[3] == v","tryCatchPattern":"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","preventionTips":["Allocate all ReplaySSM caches from one helper that takes (num_q_heads, num_v_heads, key_dim, value_dim, L)","Assert cache shapes in unit tests (as test_replayssm_decode_contract does)"],"tags":["helion","kda","replayssm","tensor-shape","cache-layout"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}