{"record":{"id":"eb285b11f13656d0","repo":"sgl-project/sglang","slug":"hd256-forward-varlen-expects-k-rank-3-or-5-got-ra","errorCode":null,"errorMessage":"hd256 forward varlen expects k rank 3 or 5, got rank {k_rank}","messagePattern":"hd256 forward varlen expects k rank 3 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":295,"sourceCode":"                d = mQ.shape[4]\n            elif cutlass.const_expr(q_rank == 4):\n                s_q = mQ.shape[1]\n                h_q = mQ.shape[2]\n                d = mQ.shape[3]\n            else:\n                raise RuntimeError(\n                    f\"hd256 forward non-varlen expects q rank 4 or 5, got rank {q_rank}\"\n                )\n\n        if cutlass.const_expr(cum_seqlen_k is not None):\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 == 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:","sourceCodeStart":277,"sourceCodeEnd":313,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/flash_attn/cute/sm100_hd256_2cta_fmha_forward.py#L277-L313","documentation":"The hd256 (head-dim 256) 2-CTA FMHA forward varlen kernel accepts K tensors of rank 3 (token-major [total_tokens, Hk, D]) or rank 5 (batched [B, G, S, Hk, D]) only. Any other rank is rejected at compile-time of the CUTLASS JitExecutor call. It means the caller passed a K layout the varlen kernel cannot interpret.","triggerScenarios":"Calling the sm100 hd256 2cta fmha forward with cum_seqlens present and a k tensor whose ndim is not 3 or 5 (e.g. rank-4 [B, S, Hk, D] with varlen path, or a squeezed/expanded tensor).","commonSituations":"Passing a non-varlen shaped [B,S,H,D] K while still supplying cum_seqlens; upstream attention backend reshaping KV cache incorrectly; head-dim-256 model (e.g. GQA-256) misconfigured.","solutions":["Reshape k to rank 3 [total_tokens, Hk, D] when using varlen (cum_seqlen) path","Or pass rank 5 [B, G, S, Hk, D] layout","Drop cum_seqlens to use the non-varlen path which accepts rank 4/5","Check the caller in the attention backend that packs KV for hd256 models"],"exampleFix":"# before\nout = fmha(q, k, k, cu_seqlens_q, cu_seqlens_k)  # k is [B, S, Hk, D]\n# after\nk = k.reshape(-1, k.shape[-2], k.shape[-1])  # [total_tokens, Hk, D]\nout = fmha(q, k, k, cu_seqlens_q, cu_seqlens_k)","handlingStrategy":"validation","validationCode":"assert k.ndim in (3, 5), f'varlen k must be rank 3 or 5, got {k.ndim}'","typeGuard":"def is_valid_varlen_k(k: torch.Tensor, varlen: bool) -> bool:\n    return k.ndim in (3, 5) if varlen else k.ndim in (4, 5)","tryCatchPattern":null,"preventionTips":["Centralize K packing in one helper that returns rank-3 for varlen","Add shape asserts in the backend wrapper before dispatch"],"tags":["cuda","flash-attention","tensor-shape","sm100","cutlass"],"backgroundTag":"invalid-tensor-rank","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}