[fa] Make the FlashAttention backend extensible by subclasses

[klshuster]

Motivation

FlashAttentionBackend is hard to specialize without copying it. Three things a subclass may need are not exposed today, leaving only the options of duplicating the backend or patching module globals: the logsumexp the kernel already computes, per-layer control of split-KV, and the choice of FA4 build. This adds a seam for each. All three are no-ops for current behavior.

Modifications

return_lse on forward_extend / forward_decode. The kernel already computes the logsumexp for the cascade merge and the backend discards it, and the forwards accept neither return_lse nor **kwargs, so a caller that needs the softmax normalizer has no way to get it. With return_lse=True they return (output, lse), lse shaped [total_q, num_heads] (the varlen kernel emits it head-major). _check_return_lse raises on the combinations the paged path cannot serve — non-FA4, MLA, fa_skip_kv_cache, cross-attention, context parallelism, local and cascade attention — instead of returning a quietly wrong tensor. Defaults to False.

_resolve_num_splits(layer, forward_batch, is_swa_layer). Split-KV was read from self.num_splits at 13 kernel arguments, making it one value for the whole run. It is now resolved once per forward and threaded through, so a subclass can vary it per layer or per batch — for example forcing 1 on sliding-window layers, which degenerate when split. The base implementation returns self.num_splits. Since FA3's scheduler_metadata is precomputed from self.num_splits, the existing sched_meta gate also requires num_splits == self.num_splits so a subclass cannot reuse a schedule built for a different count; always true for the base class.

Kernel entry points as class attributes. flash_attn_with_kvcache and flash_attn_varlen_func are bound on the class and the 15 invocations go through self._flash_attn_*, so a subclass can substitute a different FA4 build.

Accuracy Tests

Two cases added to test/registered/attention/unittests/dense/test_fa4.py, self-contained so the shared dense kit is untouched:

• test_return_lse — calls the forwards directly, since return_lse is a backend-level contract and the RadixAttention dispatcher's custom-op schema cannot carry it, then checks the output and the LSE against a causal fp32 reference for MHA extend and GQA decode, plus the [total_q, num_heads] shape.
• test_return_lse_rejects_unsupported_paths — the guarded paths raise, and the check is a no-op when the LSE is not requested.

Both pass on a GB300 with no skips. Perturbing the returned LSE by a constant fails test_return_lse on 100% of elements, confirming the comparison is load-bearing.

Not covered: windowed layers, because the dense harness leaves attention_chunk_size unset, so a sliding_window_size case is skipped for fa4 before it asserts anything. The FA3 sched_meta gate was also not exercised, as the FA3 dense suite skips on that hardware; it is inert for the base class, whose hook returns self.num_splits.

Speed Tests and Profiling

Not applicable. No behavior changes for existing callers: return_lse defaults to False, _resolve_num_splits returns self.num_splits, the _flash_attn_* attributes are bound to the same functions the module already imported, and the new sched_meta condition is always true unless a subclass varies the split count.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ⏳ Run #34800666242
Latest PR Test (Extra): ❌ Run #34800666040
Latest PR Test (AMD ROCm 10): ⏳ Run #34800666284

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[klshuster]

@Fridge003, I rebased onto current main and cut this down a lot. Dropped _resolve_num_splits entirely since decode_num_splits / _get_fa_runtime_policy covers it, and removed the return_lse check and its test as you asked. The kernel indirection now just aliases the self.flash_attn_* attributes __init__ already assigns, so all 15 call sites are unchanged; what's left is return_lse plus one test.

[klshuster]

@Fridge003 do you need anything else from me to push this through? just rebased onto main again