{"record":{"id":"61439f2963f5095a","repo":"sgl-project/sglang","slug":"o-tensor-must-match-dtype","errorCode":null,"errorMessage":"O tensor must match dtype","messagePattern":"O tensor must match dtype","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py","lineNumber":223,"sourceCode":"    def __call__(\n        self,\n        mO_partial: cute.Tensor,\n        mLSE_partial: cute.Tensor,\n        mO: cute.Tensor,\n        mLSE: Optional[cute.Tensor] = None,\n        cu_seqlens: Optional[cute.Tensor] = None,\n        seqused: Optional[cute.Tensor] = None,\n        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,\n        varlen_batch_idx: Optional[cute.Tensor] = None,\n        semaphore_to_reset: Optional[cute.Tensor] = None,\n        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).\n        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            )","sourceCodeStart":205,"sourceCodeEnd":241,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py#L205-L241","documentation":"The FlashAttention-4 CuTe split-K combine kernel requires the output tensor mO to have the exact element dtype the combine object was constructed with (self.dtype). A dtype mismatch means the kernel would write garbage or fail to compile, so it is rejected up front.","triggerScenarios":"Calling FlashAttnCombineCombine (or the combine stage of a split-K FA4 run) with an mO tensor whose dtype differs from the dtype the combine class was instantiated with, e.g. bf16 combine object but fp16 output tensor.","commonSituations":"Mixed-precision attention setups: partials accumulated in one dtype (e.g. fp16) but the final output buffer allocated as bf16, or reusing a combine object across models with different dtypes.","solutions":["Allocate/convert mO to match the combine object's dtype: mO = torch.empty(..., dtype=combine.dtype, device=...)","Check where the combine object is constructed and align its dtype argument with your output tensor","If bridging precisions, do an explicit .to(dtype) after the kernel instead of passing a mismatched buffer"],"exampleFix":"// before\nmO = torch.empty(shape, dtype=torch.float16)\ncombine(mO_partial, mLSE_partial, mO, mLSE)\n// after\nmO = torch.empty(shape, dtype=combine.dtype)\ncombine(mO_partial, mLSE_partial, mO, mLSE)","handlingStrategy":"type-guard","validationCode":"assert mO.dtype == torch_dtype_used_for_combine, f'expected {combine.dtype}, got {mO.dtype}'","typeGuard":"def out_matches_combine(mO, combine) -> bool:\n    return mO.dtype == combine.dtype","tryCatchPattern":null,"preventionTips":["Allocate the output tensor from the combine object's dtype, not a hardcoded torch dtype","Centralize dtype selection in one config constant used by both quantizer and buffers"],"tags":["flash-attention","dtype-mismatch","cuda-kernel","validation"],"backgroundTag":"dtype-mismatch","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}