sgl-project/sglang · error · NotImplementedError

FA4 CuTe FP8 backward is not supported yet (forward-only).

Error message

FA4 CuTe FP8 backward is not supported yet (forward-only).

What it means

The CuTe (FA4) flash attention interface supports FP8 inputs only in the forward direction; no FP8 backward kernel exists yet. Any input tensor requiring grad with FP8 Q or V triggers this NotImplementedError.

Source

Thrown at python/sglang/kernels/ops/attention/flash_attn/cute/interface.py:606

        softmax_scale = (
            1.0 / math.sqrt(head_dim)
            if qv is None or q is None
            else 1.0 / math.sqrt(head_dim + head_dim_v)
        )
    if softcap == 0.0:
        softcap = None
    qhead_per_kvhead = num_head // num_head_kv
    if pack_gqa is None:
        pack_gqa = qhead_per_kvhead > 1
    if pack_gqa:
        # pack_gqa reshapes SFQ's head/token layout, which the interleaved atom
        # can't express; fall back to the dense (non-interleaved) SFQ path.
        q_sf_interleaved = False

    is_fp8 = v.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
    requires_grad = any(t is not None and t.requires_grad for t in [q, k, v, qv])
    if is_fp8 and requires_grad:
        raise NotImplementedError(
            "FA4 CuTe FP8 backward is not supported yet (forward-only)."
        )
    # qk_blockscaled (fp8 Q/K, bf16 V): output follows V's dtype. v_blockscaled
    # (fp8 V dequanted in-kernel): output is bf16.
    if qk_blockscaled:
        out_torch_dtype = torch.bfloat16 if v_blockscaled else v.dtype
    else:
        out_torch_dtype = torch.bfloat16 if is_fp8 else q_dtype
    device = v.device
    q_batch_seqlen_shape = (
        (batch_size, seqlen_q) if cu_seqlens_q is None else (total_q,)
    )

    if qv is None:
        lse_shape = (
            (batch_size, num_head, seqlen_q)
            if cu_seqlens_q is None
            else (num_head, total_q)

View on GitHub (pinned to 0132848349)

Solutions

1. Detach the FP8 tensors: pass q.detach() etc. or run under torch.no_grad()
2. Switch Q/V to bf16 if backward is genuinely needed
3. Wait for upstream FP8 backward support in FA4 CuTe

Example fix

// before
loss = fa(q_fp8.requires_grad_(True), k, v_fp8).sum(); loss.backward()
// after
out = fa(q_fp8.detach(), k, v_fp8.detach())  # inference only

is_fp8 = any(t is not None and t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for t in (q, k, v))
needs_grad = any(t is not None and t.requires_grad for t in (q, k, v, qv))
assert not (is_fp8 and needs_grad)

def fp8_backward_safe(q, k, v) -> bool:
    fp8 = any(t is not None and t.dtype in (torch.float8_e4m3fn, torch.float8_e5m2) for t in (q, k, v))
    grad = any(t is not None and t.requires_grad for t in (q, k, v))
    return not (fp8 and grad)

try:
    out = fa(q, k, v)
except NotImplementedError:
    out = fa(q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16))  # bf16 fallback supports backward

When it happens

Trigger: Calling flash attention with v (or q) in torch.float8_e4m3fn/float8_e5m2 while q/k/v/qv have requires_grad=True (e.g. under autograd.backward or training loop).

Common situations: Fine-tuning or training with FP8 quantized activations without detaching; using FlexAttention-style APIs that assume differentiable paths in FP8 inference code.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Backward pass is not implemented yet and we do not have plan
• out must not require gradients
• Type mismatch: {self.q_dtype} != {self.k_dtype}
• Type mismatch: {self.q_dtype} != {self.v_dtype}
• Type mismatch: {self.sfq_dtype} != {self.sfk_dtype}

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