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
- Allocate mLSE_partial as torch.float32
- 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
- Remember LSE is always fp32 by contract in flash attention kernels
- Add an assert on LSE dtype in test fixtures for attention paths
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
- O tensor must match dtype
- LSE tensor must be Float32
- v_cache must be provided
- q can only be None when only_qv=True
- q must be provided unless qv is provided with only_qv=True
AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28).
Data as JSON: /api/errors/1c63354bfa355a32.
Report an issue: GitHub.