sgl-project/sglang · error · TypeError

SplitKV partial output (mO) must be Float32

Error message

SplitKV partial output (mO) must be Float32

What it means

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.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd.py:209

    def _check_type(
        self,
        mQ_type: Type[cutlass.Numeric],
        mK_type: Type[cutlass.Numeric],
        mV_type: Type[cutlass.Numeric],
        mO_type: Type[cutlass.Numeric],
        mLSE_type: Type[cutlass.Numeric] | None,
        mCuSeqlensQ_type: Type[cutlass.Numeric] | None,
        mCuSeqlensK_type: Type[cutlass.Numeric] | None,
        mSeqUsedQ_type: Type[cutlass.Numeric] | None,
        mSeqUsedK_type: Type[cutlass.Numeric] | None,
    ):
        # Get the data type and check if it is fp16 or bf16
        if const_expr(self.is_split_kv):
            # SplitKV writes float32 partial outputs; Q/K/V still fp16/bf16.
            if const_expr(not (mQ_type == mK_type == mV_type)):
                raise TypeError("Q/K/V must have the same data type")
            if const_expr(mO_type != Float32):
                raise TypeError("SplitKV partial output (mO) must be Float32")
        elif const_expr(not (mQ_type == mK_type == mV_type == mO_type)):
            raise TypeError("All tensors must have the same data type")
        if const_expr(mQ_type not in [cutlass.Float16, cutlass.BFloat16]):
            raise TypeError("Only Float16 or BFloat16 is supported")
        if const_expr(mLSE_type not in [None, Float32]):
            raise TypeError("LSE tensor must be Float32")
        if const_expr(mCuSeqlensQ_type not in [None, Int32]):
            raise TypeError("cu_seqlens_q tensor must be Int32")
        if const_expr(mCuSeqlensK_type not in [None, Int32]):
            raise TypeError("cu_seqlens_k tensor must be Int32")
        if const_expr(mSeqUsedQ_type not in [None, Int32]):
            raise TypeError("seqused_q tensor must be Int32")
        if const_expr(mSeqUsedK_type not in [None, Int32]):
            raise TypeError("seqused_k tensor must be Int32")
        assert mQ_type == self.dtype

    def _setup_attributes(self):
        # ///////////////////////////////////////////////////////////////////////////////

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate the partial output tensor as float32: torch.empty(..., dtype=torch.float32, device=...)
2. Keep Q/K/V in fp16/bf16 but only the partial O in fp32 — do not unify all dtypes
3. Check the combine kernel (flash_fwd_combine) expectations: it consumes Float32 partial O and LSE

Example fix

// before
O_partial = torch.empty(shape, dtype=torch.float16, device='cuda')
// after
O_partial = torch.empty(shape, dtype=torch.float32, device='cuda')

assert O_partial.dtype == torch.float32, f'SplitKV partial O must be fp32, got {O_partial.dtype}'
assert Q.dtype in (torch.float16, torch.bfloat16)

def is_valid_splitkv_args(Q, K, V, O_partial) -> bool:
    return Q.dtype == K.dtype == V.dtype in (torch.float16, torch.bfloat16) and O_partial.dtype == torch.float32

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• All tensors must have the same data type
• Only Float16 or BFloat16 is supported
• LSE tensor must be Float32
• O partial tensor must match dtype_partial
• cu_seqlens_q tensor must be Int32

AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/07918d2aed8661cd. Report an issue: GitHub.