{"record":{"id":"eaa2dc5f3ffb7390","repo":"sgl-project/sglang","slug":"expected-recurrent-state-layout-pool-h-v-128-k","errorCode":null,"errorMessage":"expected recurrent state layout [pool, H, V=128, K=128]","messagePattern":"expected recurrent state layout \\[pool, H, V=128, K=128\\]","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py","lineNumber":1003,"sourceCode":"    expected_shape = (1, T, H, TILE_K)\n    if tuple(x_q.shape) != expected_shape or tuple(x_k.shape) != expected_shape:\n        raise ValueError(f\"expected q/k shape {expected_shape}\")\n    if tuple(x_v.shape) != expected_shape or tuple(g.shape) != expected_shape:\n        raise ValueError(f\"expected v/g shape {expected_shape}\")\n    if tuple(beta.shape) != (1, T, H):\n        raise ValueError(f\"expected beta shape {(1, T, H)}\")\n    # T // N == 1 is num_spec == 0: one token per request, i.e. a plain decode\n    # step. The backend never dispatches here for it (that is the dedicated\n    # decode kernel's job), but the layout is legal and benchmarks compare the\n    # two at this point, so the wrapper accepts it.\n    if N <= 0 or T % N != 0 or T // N < 1:\n        raise ValueError(\n            f\"DSpARK KDA MTP requires a fixed 1 + num_spec dense tokens per \"\n            f\"request; got T={T}, N={N}\"\n        )\n    num_spec = T // N - 1\n    if recurrent_state.shape[1:] != (H, TILE_K, TILE_K):\n        raise ValueError(\"expected recurrent state layout [pool, H, V=128, K=128]\")\n    if (\n        recurrent_state.dtype != torch.float32\n        or tuple(recurrent_state.stride()[-3:]) != (TILE_K * TILE_K, TILE_K, 1)\n        or recurrent_state.stride(0) % 4 != 0\n        or recurrent_state.storage_offset() % 4 != 0\n    ):\n        raise ValueError(\n            \"cp.async recurrent state requires fp32 contiguous [H, V, K] \"\n            \"inner layout and 16-byte-aligned slot offsets\"\n        )\n    rings = (replayssm_rawv, replayssm_rawk, replayssm_g, replayssm_beta)\n    cache_ring = all(ring is not None for ring in rings)\n    if any(ring is not None for ring in rings) and not cache_ring:\n        raise ValueError(\"ReplaySSM requires all four replayssm_* rings\")\n    if cache_ring:\n        ring_len = replayssm_rawv.shape[2]\n        if (\n            ring_len < 1 + num_spec","sourceCodeStart":985,"sourceCodeEnd":1021,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/kimi_k3/kda_decode_mtp.py#L985-L1021","documentation":"The recurrent state pool must be a 4D tensor [pool, H, 128, 128] with float32 dtype, contiguous inner strides (TILE_K*TILE_K, TILE_K, 1), pool stride 4-aligned, and 4-aligned storage offset — requirements of the kernel's vectorized state updates. The shape check [pool, H, V=128, K=128] fails otherwise.","triggerScenarios":"Passing recurrent_state with shape [pool, H, 64, 64] (head_dim 64 model), a bf16/fp16 state, a transposed or non-contiguous view, or a slice with unaligned storage offset.","commonSituations":"Allocating the Mamba/KDA state pool in model dtype instead of fp32; sharing a state cache across models with different head dims; passing state.transpose(-1, -2) views.","solutions":["Allocate recurrent_state as torch.float32 with shape [pool_size, H, 128, 128], contiguous","If it is a view/slice, call .contiguous() and ensure the storage offset is 4-element aligned (allocate fresh rather than slicing)","For head dims != 128, use the generic KDA decode backend instead of this fused MTP kernel"],"exampleFix":"# before\nstate = torch.zeros(pool, H, D, D, dtype=x.dtype, device='cuda')  # bf16, D!=128\n# after\nstate = torch.zeros(pool, H, 128, 128, dtype=torch.float32, device='cuda')","handlingStrategy":"validation","validationCode":"assert recurrent_state.shape[1:] == (H, 128, 128)\nassert recurrent_state.dtype == torch.float32 and recurrent_state.is_contiguous()","typeGuard":"def is_valid_state(s: torch.Tensor, H: int) -> bool:\n    return (s.dtype == torch.float32 and s.shape[1:] == (H, 128, 128)\n            and tuple(s.stride()[-3:]) == (128*128, 128, 1))","tryCatchPattern":null,"preventionTips":["Allocate the state pool once as fp32 [pool, H, 128, 128] contiguous","Avoid transposed or sliced views into the pool; index whole rows","Assert head_dim == 128 in the model config when enabling this fused path"],"tags":["kda","mtp","recurrent-state","dtype-stride"],"backgroundTag":"invalid-tensor-layout","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}