What causes "LSE tensor must be Float32" in sgl-project/sglang?

Allocate LSE as float32: torch.empty((b,h,s), dtype=torch.float32, device=...)

sgl-project/sglang · error · TypeError

LSE tensor must be Float32

Error message

LSE tensor must be Float32

What it means

The optional LSE (log-sum-exp) output tensor of the flash-attention forward op must be Float32 when provided. LSE stores per-head log-normalization constants in fp32 for numerical stability and for later combine/reduction steps.

Source

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

        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
        # ///////////////////////////////////////////////////////////////////////////////
        (
            sQ_layout_atom,
            sK_layout_atom,
            sV_layout_atom,

View on GitHub (pinned to 0132848349)

Solutions

Allocate LSE as float32: torch.empty((b,h,s), dtype=torch.float32, device=...)
Pass None if you don't need LSE output

Example fix

// before
LSE = torch.empty((b,h,s), dtype=torch.float16, device='cuda')
// after
LSE = torch.empty((b,h,s), dtype=torch.float32, device='cuda')

Defensive patterns

Strategy: type-guard

Validate before calling

if LSE is not None:
    assert LSE.dtype == torch.float32

Type guard

def valid_lse(LSE) -> bool:
    return LSE is None or LSE.dtype == torch.float32

Prevention

Allocate LSE explicitly with dtype=torch.float32
Pass None when LSE is not needed

When it happens

Trigger: Calling FlashAttentionForward with an mLSE tensor whose element type is Float16/BFloat16 instead of Float32 (passing None is allowed and skips the check).

Common situations: Allocating LSE with the same dtype as O for symmetry; reusing an fp16 workspace buffer for LSE; LSE needed downstream by a split-KV combine or backwards pass that assumes fp32.

Related errors

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