{"record":{"id":"d099ccda997b2593","repo":"sgl-project/sglang","slug":"lse-partial-tensor-must-have-3-or-4-dimensions-n","errorCode":null,"errorMessage":"LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)","messagePattern":"LSE partial tensor must have 3 or 4 dimensions: \\(num_splits, batch, seqlen, nheads\\) or \\(num_splits, total_q, nheads\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py","lineNumber":235,"sourceCode":"        stream: cuda.CUstream = None,\n    ):\n        # Type checking\n        if const_expr(not (mO_partial.element_type == self.dtype_partial)):\n            raise TypeError(\"O partial tensor must match dtype_partial\")\n        if const_expr(not (mO.element_type == self.dtype)):\n            raise TypeError(\"O tensor must match dtype\")\n        if const_expr(mLSE_partial.element_type not in [Float32]):\n            raise TypeError(\"LSE partial tensor must be Float32\")\n        if const_expr(mLSE is not None and mLSE.element_type not in [Float32]):\n            raise TypeError(\"LSE tensor must be Float32\")\n\n        # Shape validation - input tensors are in user format, need to be converted to kernel format\n        if const_expr(len(mO_partial.shape) not in [4, 5]):\n            raise ValueError(\n                \"O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)\"\n            )\n        if const_expr(len(mLSE_partial.shape) not in [3, 4]):\n            raise ValueError(\n                \"LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)\"\n            )\n        if const_expr(len(mO.shape) not in [3, 4]):\n            raise ValueError(\n                \"O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)\"\n            )\n        if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]):\n            raise ValueError(\n                \"LSE tensor must have 2 or 3 dimensions: (batch, seqlen, nheads) or (total_q, nheads)\"\n            )\n\n        mO_partial, mO = [assume_tensor_aligned(t) for t in (mO_partial, mO)]\n        # (num_splits, b, seqlen, h, d) -> (seqlen, d, num_splits, h, b)\n        # or (num_splits, total_q, h, d) -> (total_q, d, num_splits, h)\n        O_partial_layout_transpose = (\n            [2, 4, 0, 3, 1] if const_expr(cu_seqlens is None) else [1, 3, 0, 2]\n        )\n        # (b, seqlen, h, d) -> (seqlen, d, h, b) or (total_q, h, d) -> (total_q, d, h)","sourceCodeStart":217,"sourceCodeEnd":253,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py#L217-L253","documentation":"The partial LSE tensor must be 3D (num_splits, batch, nheads) batched or 4D (num_splits, total_q, nheads) varlen, matching the layout of the O partials.","triggerScenarios":"Passing mLSE_partial with rank other than 3 or 4, or a layout inconsistent with the O partials (e.g. 2D LSE for 4D O partials).","commonSituations":"Computing LSE in a transposed shape (nheads, batch) then flattening; mismatched batched vs varlen call styles between O and LSE.","solutions":["Match mLSE_partial rank to mO_partial: 4D O partials -> 3D LSE partials (num_splits,b,h); 5D O partials -> 4D (num_splits,total_q,h)","Verify seqlen/total_q axis matches the O partial tensor"],"exampleFix":"// before\nmLSE_partial = lse.view(num_splits, -1, h)  # wrong for varlen\n// after\nmLSE_partial = lse.view(num_splits, total_q, h)","handlingStrategy":"validation","validationCode":"assert mLSE_partial.dim() in (3, 4)\nassert mLSE_partial.dim() == mO_partial.dim() - 1","typeGuard":"def lse_partial_ok(lse, o_partial) -> bool:\n    return lse.dim() in (3, 4) and lse.dim() == o_partial.dim() - 1","tryCatchPattern":null,"preventionTips":["Derive LSE shape from the O partial shape programmatically instead of hand-writing it"],"tags":["flash-attention","shape-mismatch","validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}