{"record":{"id":"07918d2aed8661cd","repo":"sgl-project/sglang","slug":"splitkv-partial-output-mo-must-be-float32","errorCode":null,"errorMessage":"SplitKV partial output (mO) must be Float32","messagePattern":"SplitKV partial output \\(mO\\) must be Float32","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd.py","lineNumber":209,"sourceCode":"    def _check_type(\n        self,\n        mQ_type: Type[cutlass.Numeric],\n        mK_type: Type[cutlass.Numeric],\n        mV_type: Type[cutlass.Numeric],\n        mO_type: Type[cutlass.Numeric],\n        mLSE_type: Type[cutlass.Numeric] | None,\n        mCuSeqlensQ_type: Type[cutlass.Numeric] | None,\n        mCuSeqlensK_type: Type[cutlass.Numeric] | None,\n        mSeqUsedQ_type: Type[cutlass.Numeric] | None,\n        mSeqUsedK_type: Type[cutlass.Numeric] | None,\n    ):\n        # Get the data type and check if it is fp16 or bf16\n        if const_expr(self.is_split_kv):\n            # SplitKV writes float32 partial outputs; Q/K/V still fp16/bf16.\n            if const_expr(not (mQ_type == mK_type == mV_type)):\n                raise TypeError(\"Q/K/V must have the same data type\")\n            if const_expr(mO_type != Float32):\n                raise TypeError(\"SplitKV partial output (mO) must be Float32\")\n        elif const_expr(not (mQ_type == mK_type == mV_type == mO_type)):\n            raise TypeError(\"All tensors must have the same data type\")\n        if const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]):\n            raise TypeError(\"Only Float16 or BFloat16 is supported\")\n        if const_expr(mLSE_type not in [None, Float32]):\n            raise TypeError(\"LSE tensor must be Float32\")\n        if const_expr(mCuSeqlensQ_type not in [None, Int32]):\n            raise TypeError(\"cu_seqlens_q tensor must be Int32\")\n        if const_expr(mCuSeqlensK_type not in [None, Int32]):\n            raise TypeError(\"cu_seqlens_k tensor must be Int32\")\n        if const_expr(mSeqUsedQ_type not in [None, Int32]):\n            raise TypeError(\"seqused_q tensor must be Int32\")\n        if const_expr(mSeqUsedK_type not in [None, Int32]):\n            raise TypeError(\"seqused_k tensor must be Int32\")\n        assert mQ_type == self.dtype\n\n    def _setup_attributes(self):\n        # ///////////////////////////////////////////////////////////////////////////////","sourceCodeStart":191,"sourceCodeEnd":227,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd.py#L191-L227","documentation":"Raised by the CUTLASS DSL flash-attention forward op's type checker when running in SplitKV mode: the partial output tensor mO must be Float32 because the split-KV kernel accumulates and writes partial results in fp32 (Q/K/V remain fp16/bf16). It is a compile-time const_expr check executed when __call__ builds the kernel.","triggerScenarios":"Calling FlashAttentionForward with is_split_kv=True (or a split-KV variant) where the output partial tensor passed as mO has element type Float16 or BFloat16 instead of Float32.","commonSituations":"Allocating the partial O tensor with the same dtype as Q/K/V (torch.float16) when preparing split-KV flash attention; porting code from a non-split path that required all tensors to share dtype; dtype changes after autocast or when combining partial outputs.","solutions":["Allocate the partial output tensor as float32: torch.empty(..., dtype=torch.float32, device=...)","Keep Q/K/V in fp16/bf16 but only the partial O in fp32 — do not unify all dtypes","Check the combine kernel (flash_fwd_combine) expectations: it consumes Float32 partial O and LSE"],"exampleFix":"// before\nO_partial = torch.empty(shape, dtype=torch.float16, device='cuda')\n// after\nO_partial = torch.empty(shape, dtype=torch.float32, device='cuda')","handlingStrategy":"type-guard","validationCode":"assert O_partial.dtype == torch.float32, f'SplitKV partial O must be fp32, got {O_partial.dtype}'\nassert Q.dtype in (torch.float16, torch.bfloat16)","typeGuard":"def is_valid_splitkv_args(Q, K, V, O_partial) -> bool:\n    return Q.dtype == K.dtype == V.dtype in (torch.float16, torch.bfloat16) and O_partial.dtype == torch.float32","tryCatchPattern":null,"preventionTips":["Centralize partial-output allocation in one helper that always uses torch.float32","Add dtype assertions in test fixtures for split-KV paths"],"tags":["cuda","dtype","flash-attention","cutlass","split-kv"],"backgroundTag":null,"analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}