[fa] Make the FlashAttention backend extensible by subclasses (#33426)

Signed-off-by: Kurt Shuster <kurt@thinkingmachines.ai>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

test/registered/attention/unittests/dense/test_fa4.py

@@ -3,9 +3,13 @@ import unittest
 import torch
 
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.attention_unittest.attention_methods.dense_attention import (
+    DENSE_ATOL,
+    DENSE_RTOL,
     DenseAttentionCase,
+    build_dense_attention_fixture,
     make_dense_cases,
     run_dense_attention_case,
 )
@@ -423,6 +427,102 @@ class TestFA4DenseAttentionBackendCorrectness(CustomTestCase):
                     hidden_size=self.HIDDEN_SIZE,
                 )
 
+    RETURN_LSE_CASES = (
+        DenseAttentionCase(
+            name="return_lse_mha_extend",
+            backend="fa4",
+            forward_mode=ForwardMode.EXTEND,
+            num_heads=4,
+            num_kv_heads=4,
+            page_size=1,
+            prefix_lens=(2, 4),
+            extend_lens=(3, 1),
+        ),
+        DenseAttentionCase(
+            name="return_lse_gqa_decode",
+            backend="fa4",
+            forward_mode=ForwardMode.DECODE,
+            num_heads=8,
+            num_kv_heads=2,
+            page_size=1,
+            prefix_lens=(5, 9),
+        ),
+    )
+
+    def _reference_out_and_lse(self, fixture):
+        """Causal fp32 reference: output and per-query LSE."""
+        case, module = fixture.case, fixture.reference_module
+        dim = module.head_dim
+        rep = case.num_heads // case.num_kv_heads
+        q, k, v = module.project_qkv(fixture.input_hidden)
+        q = q.view(-1, case.num_heads, dim).float()
+        k = k.view(-1, case.num_kv_heads, dim).float()
+        v = v.view(-1, case.num_kv_heads, dim).float()
+
+        outs, lses, seen = [], [], 0
+        for req, prefix in enumerate(fixture.prefix_hidden):
+            _, prefix_k, prefix_v = module.project_qkv(prefix)
+            n = case.input_lens[req]
+            keys = torch.cat(
+                [prefix_k.view(-1, case.num_kv_heads, dim).float(), k[seen : seen + n]]
+            ).repeat_interleave(rep, dim=1)
+            values = torch.cat(
+                [prefix_v.view(-1, case.num_kv_heads, dim).float(), v[seen : seen + n]]
+            ).repeat_interleave(rep, dim=1)
+            for offset in range(n):
+                end = case.prefix_lens[req] + offset + 1
+                scores = (
+                    torch.einsum("hd,khd->hk", q[seen + offset], keys[:end])
+                    * module.scaling
+                )
+                probs = torch.softmax(scores, dim=-1)
+                outs.append(torch.einsum("hk,khd->hd", probs, values[:end]).reshape(-1))
+                lses.append(torch.logsumexp(scores, dim=-1))
+            seen += n
+        return torch.stack(outs), torch.stack(lses)
+
+    def test_return_lse(self):
+        """Calls the backend directly: return_lse is a backend-level contract,
+        and the RadixAttention dispatcher's custom-op schema cannot carry it.
+        """
+        for case in self.RETURN_LSE_CASES:
+            with self.subTest(case=case.name):
+                fixture = build_dense_attention_fixture(
+                    self, case, head_dim=self.HEAD_DIM, hidden_size=self.HIDDEN_SIZE
+                )
+                module = fixture.actual_module
+                forward = (
+                    fixture.backend.forward_decode
+                    if case.forward_mode.is_decode()
+                    else fixture.backend.forward_extend
+                )
+                with (
+                    torch.no_grad(),
+                    forward_context(ForwardContext(attn_backend=fixture.backend)),
+                ):
+                    fixture.backend.init_forward_metadata(fixture.forward_batch)
+                    q, k, v = module.project_qkv(fixture.input_hidden)
+                    result = forward(
+                        q, k, v, module.attn, fixture.forward_batch, return_lse=True
+                    )
+
+                self.assertIsInstance(result, tuple)
+                out, lse = result
+                self.assertEqual(
+                    tuple(lse.shape), (case.num_input_tokens, case.num_heads)
+                )
+
+                expected_out, expected_lse = self._reference_out_and_lse(fixture)
+                torch.testing.assert_close(
+                    lse.float(), expected_lse, atol=DENSE_ATOL, rtol=DENSE_RTOL
+                )
+                torch.testing.assert_close(
+                    out.float().reshape(expected_out.shape),
+                    expected_out,
+                    atol=DENSE_ATOL,
+                    rtol=DENSE_RTOL,
+                )
+
 
 if __name__ == "__main__":
     unittest.main()