sgl-project/sglang · error · TypeError

All tensors must have the same data type

Error message

All tensors must have the same data type

What it means

Raised by the flash-attention forward op type checker in the non-SplitKV path: Q, K, V, and O must all share the same data type for the kernel to compile (mixed-dtype tensors are unsupported). This is a static const_expr check performed on tensor element types during __call__.

Source

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

        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):
        # ///////////////////////////////////////////////////////////////////////////////
        # Shared memory layout: Q/K/V
        # ///////////////////////////////////////////////////////////////////////////////

View on GitHub (pinned to 0132848349)

Solutions

1. Unify all four tensors: cast O to match Q/K/V (o = torch.empty_like(q) or dtype=q.dtype)
2. Verify no tensor was created under torch.get_default_dtype() == float32
3. If you intentionally want fp32 partials, use the SplitKV path where O is fp32 by design

Example fix

// before
O = torch.empty((b,h,s,d), dtype=torch.float32, device='cuda')
// after
O = torch.empty((b,h,s,d), dtype=Q.dtype, device='cuda')

assert Q.dtype == K.dtype == V.dtype == O.dtype, 'Q/K/V/O dtypes must match for non-split flash attention'

def dtypes_match(*ts) -> bool:
    return len({t.dtype for t in ts}) == 1

When it happens

Trigger: Calling FlashAttentionForward without split-KV where any of Q/K/V/O differs in dtype, e.g. Q,K,V in bfloat16 but O preallocated as float16, or one input cast differently.

Common situations: Preallocating the output tensor with a default dtype (torch.float32) while inputs are fp16; mixing autocast-produced tensors with manually cast ones; partial refactors that cast only some inputs.

Related errors

• SplitKV partial output (mO) must be Float32
• Only Float16 or BFloat16 is supported
• LSE tensor must be Float32
• cu_seqlens_q tensor must be Int32
• cu_seqlens_k tensor must be Int32

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