sgl-project/sglang · error · TypeError

LSE partial tensor must be Float32

Error message

LSE partial tensor must be Float32

What it means

The FA4 CuTe combine kernel requires the partial log-sum-exp tensor (mLSE_partial) to be float32. LSE values are used to renormalize attention across splits and must be accumulated in fp32 for numerical correctness, so any other dtype is rejected.

Source

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

        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(
                "O tensor must have 3 or 4 dimensions: (batch, seqlen, nheads, headdim) or (total_q, nheads, headdim)"
            )
        if const_expr(mLSE is not None and len(mLSE.shape) not in [2, 3]):
            raise ValueError(

View on GitHub (pinned to 0132848349)

Solutions

  1. Allocate mLSE_partial as torch.float32
  2. Fix whatever produced the partials if it emitted non-fp32 LSE — the upstream kernel contract is fp32 LSE

Example fix

// before
mLSE_partial = torch.empty(sp.shape, dtype= torch.bfloat16)
// after
mLSE_partial = torch.empty(lse_shape, dtype=torch.float32)
Defensive patterns

Strategy: validation

Validate before calling

assert mLSE_partial.dtype == torch.float32

Type guard

def is_fp32(t) -> bool: return t.dtype == torch.float32

Prevention

When it happens

Trigger: Calling the combine __call__ with mLSE_partial of dtype float16/bfloat16 instead of float32.

Common situations: Users allocating all attention buffers with the model dtype (bf16/fp16) for uniformity; or casting LSE tensors when moving between devices.

Related errors


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