{"record":{"id":"d638a11998d645c2","repo":"sgl-project/sglang","slug":"hd256-forward-non-varlen-expects-k-rank-4-or-5-go","errorCode":null,"errorMessage":"hd256 forward non-varlen expects k rank 4 or 5, got rank {k_rank}","messagePattern":"hd256 forward non-varlen expects k rank 4 or 5, got rank (.+?)","errorType":"exception","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py","lineNumber":306,"sourceCode":"            if cutlass.const_expr(k_rank == 5):\n                s_k = mK.shape[1]\n                h_k = mK.shape[2]\n            elif cutlass.const_expr(k_rank == 3):\n                s_k = mK.shape[0]\n                h_k = mK.shape[1]\n            else:\n                raise RuntimeError(\n                    f\"hd256 forward varlen expects k rank 3 or 5, got rank {k_rank}\"\n                )\n        else:\n            if cutlass.const_expr(k_rank == 5):\n                s_k = mK.shape[1]\n                h_k = mK.shape[2]\n            elif cutlass.const_expr(k_rank == 4):\n                s_k = mK.shape[1]\n                h_k = mK.shape[2]\n            else:\n                raise RuntimeError(\n                    f\"hd256 forward non-varlen expects k rank 4 or 5, got rank {k_rank}\"\n                )\n        if cutlass.const_expr(cum_seqlen_q is not None):\n            b = mCuSeqlensQ.shape[0] - 1\n        elif cutlass.const_expr(cum_seqlen_k is not None):\n            b = mCuSeqlensK.shape[0] - 1\n        else:\n            b = mQ.shape[0]\n\n        scale_softmax = softmax_scale\n        scale_softmax_log2 = softmax_scale * math.log2(math.exp(1.0))\n        scale_output = 1.0\n        s_lse = s_q\n        h_r = h_q // h_k\n        s_q64 = Int64(s_q)\n        s_k64 = Int64(s_k)\n        s_lse64 = Int64(s_lse)\n        d64 = cute.assume(Int64(d), divby=128)","sourceCodeStart":288,"sourceCodeEnd":324,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py#L288-L324","documentation":"The hd256 2-CTA FMHA forward non-varlen (no cumulative-seqlen) path accepts K of rank 4 [B, S, Hk, D] or rank 5 [B, G, S, Hk, D]. Other ranks (e.g. rank-3 packed tokens without cum_seqlens) are rejected because batch/seq dims cannot be inferred.","triggerScenarios":"Calling the forward without cum_seqlens but with a rank-3 token-packed K tensor, or any rank other than 4/5.","commonSituations":"Migrating from varlen to batched path and forgetting to expand the packed tensor; passing a debug/random tensor with wrong ndim in unit tests.","solutions":["Pass rank 4 [B, S, Hk, D] or rank 5 [B, G, S, Hk, D] K","Or supply cu_seqlens to take the varlen path with rank-3 K","Verify upstream KV cache reshape logic"],"exampleFix":"# before\nout = fmha(q, k_packed, v_packed)  # k_packed: [total_tokens, Hk, D], no cu_seqlens\n# after\nk = k_packed.view(B, S, Hk, D)\nout = fmha(q, k, v)","handlingStrategy":"validation","validationCode":"assert k.ndim in (4, 5), f'non-varlen k must be rank 4 or 5, got {k.ndim}'","typeGuard":"def is_valid_batched_k(k: torch.Tensor) -> bool:\n    return k.ndim in (4, 5)","tryCatchPattern":null,"preventionTips":["Decide varlen vs batched path once and shape tensors accordingly","Unit-test kernel wrappers with all supported ranks"],"tags":["cuda","flash-attention","tensor-shape","sm100"],"backgroundTag":"invalid-tensor-rank","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}