{"record":{"id":"806b7bc6d1a8b594","repo":"sgl-project/sglang","slug":"initial-state-must-be-contiguous-in-the-last-dim","errorCode":null,"errorMessage":"`initial_state` must be contiguous in the last dim.","messagePattern":"`initial_state` must be contiguous in the last dim\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/fla/fused_recurrent.py","lineNumber":326,"sourceCode":"        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]}, 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] (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 a.shape[1] != HV or b.shape[1] != HV:\n        raise ValueError(\n            f\"`a`/`b` must have shape [B, HV] with HV={HV} (got a.shape={tuple(a.shape)}, b.shape={tuple(b.shape)}).\"\n        )\n    if A_log.numel() != HV or dt_bias.numel() != HV:\n        raise ValueError(\n            f\"`A_log` and `dt_bias` must have {HV} elements (got A_log.numel()={A_log.numel()}, dt_bias.numel()={dt_bias.numel()}).\"\n        )\n    if out.shape != (B, 1, HV, V):\n        raise ValueError(\n            f\"`out` must have shape {(B, 1, HV, V)} (got out.shape={tuple(out.shape)}).\"\n        )\n\n    qkv_dim = mixed_qkv.shape[1]\n    qk_dim = qkv_dim - HV * V\n    if qk_dim <= 0 or qk_dim % 2 != 0:\n        raise ValueError(","sourceCodeStart":308,"sourceCodeEnd":344,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py#L308-L344","documentation":"Like the other operand checks, the recurrent state tensor must have unit stride in its last dim (K) so the kernel can vectorize loads/stores of state rows. The wrapper raises when initial_state.stride(-1) != 1, e.g. a state cache stored transposed or a strided slice over the K dimension.","triggerScenarios":"State caches laid out (num_states, K, V, HV) passed without transpose; slicing a bigger state buffer over the last dim; states materialized from flattened views with non-compact K.","commonSituations":"Custom mamba cache formats migrated from other frameworks (e.g. (K,V) transposed layouts); TP sharding that slices the K dim producing strided rows.","solutions":["Transpose to (N, HV, V, K) with .permute(...) followed by .contiguous() once at cache-allocation time","Keep the canonical layout when allocating the cache: torch.empty(num_states, HV, V, K)"],"exampleFix":"# before\nstate = cache[layer]  # (N, HV, K, V) stored transposed\n# after\nstate = cache[layer].permute(0, 1, 3, 2).contiguous()  # (N, HV, V, K)","handlingStrategy":"validation","validationCode":"if initial_state.stride(-1) != 1:\n    initial_state = initial_state.contiguous()","typeGuard":"def state_last_contig(s: torch.Tensor) -> bool:\n    return s.ndim == 4 and s.stride(-1) == 1","tryCatchPattern":null,"preventionTips":["Allocate mamba state caches directly in (N, HV, V, K) layout","Never transpose state caches lazily inside the decode loop"],"tags":["fla","fused-recurrent","state-cache","contiguity"],"backgroundTag":"tensor-not-contiguous","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}