{"record":{"id":"3d1aa357bd5e04ed","repo":"sgl-project/sglang","slug":"k-cache-must-have-shape-slots-h-l-k","errorCode":null,"errorMessage":"`k_cache` must have shape [slots, H, L, K].","messagePattern":"`k_cache` must have shape \\[slots, H, L, K\\]\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_replayssm.py","lineNumber":740,"sourceCode":"\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(\n        is_bf16_state=initial_state.dtype is torch.bfloat16,\n        num_v_heads=num_v_heads,","sourceCodeStart":722,"sourceCodeEnd":758,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_replayssm.py#L722-L758","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Allocate k_cache as [slots, num_q_heads, cache_length, key_dim] with cache_length == d_cache.size(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","If the cache came from a pool, check the pool's per-layer shape metadata against the model config"],"exampleFix":"# before\nk_cache = torch.empty(slots, num_v_heads, cache_len, key_dim, dtype=dtype, device='cuda')\n# after\nk_cache = torch.empty(slots, num_q_heads, cache_len, key_dim, dtype=dtype, device='cuda')","handlingStrategy":"validation","validationCode":"assert k_cache.shape[1:] == (num_q_heads, d_cache.size(2), key_dim), k_cache.shape","typeGuard":"def valid_k_cache(t: torch.Tensor, h: int, k: int, l: int) -> bool:\n    return t.ndim == 4 and t.shape[1:] == (h, l, k)","tryCatchPattern":null,"preventionTips":["Remember k_cache uses H (query heads) while d/g use HV (value heads)","Keep a single config dataclass with num_q_heads/num_v_heads/key_dim/value_dim used by every allocation site"],"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"}