{"record":{"id":"225ebcdceeaa8f3f","repo":"sgl-project/sglang","slug":"initial-state-must-be-a-4d-tensor-got-ndim-ini-225ebc","errorCode":null,"errorMessage":"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim}).","messagePattern":"`initial_state` must be a 4D tensor \\(got ndim=(.+?)\\)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/helion/kda_decode.py","lineNumber":281,"sourceCode":"        )\n    ):\n        raise ValueError(\"All inputs must be on the same device.\")\n\n    B = mixed_qkv.shape[0]\n    if a.shape[0] != B or b.shape[0] != B:\n        raise ValueError(\n            \"Mismatched batch sizes: \"\n            f\"mixed_qkv.shape[0]={B}, a.shape[0]={a.shape[0]}, \"\n            f\"b.shape[0]={b.shape[0]}.\"\n        )\n    if ssm_state_indices.shape[0] != B:\n        raise ValueError(\n            f\"`ssm_state_indices` must have shape [B] \"\n            f\"(got {tuple(ssm_state_indices.shape)}; expected ({B},)).\"\n        )\n\n    if initial_state.ndim != 4:\n        raise ValueError(\n            f\"`initial_state` must be a 4D tensor (got ndim={initial_state.ndim}).\"\n        )\n    if initial_state.stride(-1) != 1:\n        raise ValueError(\"`initial_state` must be contiguous in the last dim.\")\n    HV, V, K = initial_state.shape[-3:]\n    if not _is_power_of_two(K) or not _is_power_of_two(V):\n        raise ValueError(\n            \"Helion KDA decode requires power-of-two key and value head \"\n            f\"dimensions (got K={K}, V={V}).\"\n        )\n    if a.shape[1] != HV * K:\n        raise ValueError(\n            f\"`a` must have shape [B, HV*K] with HV={HV}, K={K} \"\n            f\"(got a.shape={tuple(a.shape)}).\"\n        )\n    if b.shape[1] != HV:\n        raise ValueError(\n            f\"`b` must have shape [B, HV] with HV={HV} (got b.shape={tuple(b.shape)}).\"","sourceCodeStart":263,"sourceCodeEnd":299,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/helion/kda_decode.py#L263-L299","documentation":"Helion KDA decode validates that the recurrent SSM state tensor has 4 dimensions. initial_state is the [B, HV, V, K] (or pool-shaped) state pool read at decode time; a tensor with any other rank cannot be indexed by the fused recurrent kernel, so validate_packed_decode_inputs rejects it up front before launching.","triggerScenarios":"Calling helion_fused_recurrent_kda_packed_decode or helion_fused_recurrent_kda_replayssm_decode with an initial_state that is 2D/3D/5D — e.g. passing a per-head state [HV, V, K] instead of the batched 4D pool, or a flattened vector.","commonSituations":"Migrating a single-request decode loop to the packed batch API and forgetting to add the batch/pool dim; passing a state pool reshaped by .view() with the wrong rank; wiring a ReplaySSM checkpoint tensor of the wrong shape into the decode call.","solutions":["Reshape initial_state to 4D: [batch_or_slots, HV, V, K] matching the state pool layout","Check where initial_state is produced (state cache / mamba-style pool) and verify its rank before the decode call","Add an assert initial_state.ndim == 4 in your wrapper so failures surface at the producer, not the kernel"],"exampleFix":"// before\nstate = pool.get_state(idx)  # accidentally [HV, V, K]\nout = helion_fused_recurrent_kda_packed_decode(qkv, a, b, ..., state, ...)\n// after\nstate = pool.get_state(idx)\nassert state.ndim == 4, state.shape  # [slots, HV, V, K]\nout = helion_fused_recurrent_kda_packed_decode(qkv, a, b, ..., state, ...)","handlingStrategy":"validation","validationCode":"assert initial_state.ndim == 4, f\"expected 4D state pool, got {initial_state.shape}\"","typeGuard":"def is_valid_state_pool(t: torch.Tensor) -> bool:\n    return isinstance(t, torch.Tensor) and t.ndim == 4","tryCatchPattern":null,"preventionTips":["Always build initial_state via torch.zeros/zeros_like of shape [slots, HV, V, K]","Add ndim asserts in wrapper code near the state-pool producer"],"tags":["kda","helion","shape-validation","tensor-ndim"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}