sgl-project/sglang · error · TypeError

O partial tensor must match dtype_partial

Error message

O partial tensor must match dtype_partial

What it means

The flash-attention SplitKV combine kernel requires the partial output tensor mO_partial to match the dtype the operation was constructed with (dtype_partial, normally Float32, since split-KV partials accumulate in fp32). This is a static const_expr type check in __call__.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/flash_fwd_combine.py:221

    @cute.jit
    def __call__(
        self,
        mO_partial: cute.Tensor,
        mLSE_partial: cute.Tensor,
        mO: cute.Tensor,
        mLSE: Optional[cute.Tensor] = None,
        cu_seqlens: Optional[cute.Tensor] = None,
        seqused: Optional[cute.Tensor] = None,
        num_splits_dynamic_ptr: Optional[cute.Tensor] = None,
        varlen_batch_idx: Optional[cute.Tensor] = None,
        semaphore_to_reset: Optional[cute.Tensor] = None,
        # Always keep stream as the last parameter (EnvStream: obtained implicitly via TVM FFI).
        stream: cuda.CUstream = None,
    ):
        # Type checking
        if const_expr(not (mO_partial.element_type == self.dtype_partial)):
            raise TypeError("O partial tensor must match dtype_partial")
        if const_expr(not (mO.element_type == self.dtype)):
            raise TypeError("O tensor must match dtype")
        if const_expr(mLSE_partial.element_type not in [Float32]):
            raise TypeError("LSE partial tensor must be Float32")
        if const_expr(mLSE is not None and mLSE.element_type not in [Float32]):
            raise TypeError("LSE tensor must be Float32")

        # Shape validation - input tensors are in user format, need to be converted to kernel format
        if const_expr(len(mO_partial.shape) not in [4, 5]):
            raise ValueError(
                "O partial tensor must have 4 or 5 dimensions: (num_splits, batch, seqlen, nheads, headdim) or (num_splits, total_q, nheads, headdim)"
            )
        if const_expr(len(mLSE_partial.shape) not in [3, 4]):
            raise ValueError(
                "LSE partial tensor must have 3 or 4 dimensions: (num_splits, batch, seqlen, nheads) or (num_splits, total_q, nheads)"
            )
        if const_expr(len(mO.shape) not in [3, 4]):
            raise ValueError(

View on GitHub (pinned to 0132848349)

Solutions

1. Allocate mO_partial as float32 to match the default dtype_partial
2. Or construct the combine operation with dtype_partial matching your actual partial tensor's element type
3. Keep partial allocation and combine-op construction driven by one shared dtype variable

Example fix

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

assert O_partial.dtype == torch.float32, 'combine expects fp32 partials'
assert combine_op.dtype_partial == O_partial.dtype  # if op exposes dtype_partial

def partial_dtype_ok(O_partial, dtype_partial=torch.float32) -> bool:
    return O_partial.dtype == dtype_partial

When it happens

Trigger: Calling the combine operation with an mO_partial tensor whose element type differs from the op's dtype_partial configuration, e.g. fp16 partials fed to an op built for fp32 partials.

Common situations: Producing partials with a different kernel/configuration than the combine op expects; changing the split path to fp16 partials without reconfiguring the combine op's dtype_partial; buffer reuse across dtypes.

Related errors

• SplitKV partial output (mO) must be Float32
• All tensors must have the same data type
• Only Float16 or BFloat16 is supported
• LSE tensor must be Float32
• cu_seqlens_q tensor must be Int32

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