[Feature] Add optional FlashMLA backend for DSA Sparse MLA

docs/en/get_started/installation.md

@@ -41,6 +41,21 @@ For Hopper architecture GPUs, you can additionally install fa3 and enable it thr
 
 In addition, XTuner recommends installing flash-attn, and RL recommends installing flash-attn-3, which can significantly improve training speed. You can refer to the [official documentation](https://github.com/Dao-AILab/flash-attention) for installation.
 
+To use the FlashMLA backend for DSA Sparse MLA, install the numerically aligned
+FlashMLA branch:
+
+```{code-block} shell
+:caption: Install FlashMLA
+
+git clone -b precise-flashmla https://github.com/DeepLink-org/flashmla-optimization.git
+cd flashmla-optimization
+git submodule update --init --recursive
+FLASH_MLA_DISABLE_SM100=1 MAX_JOBS=8 NVCC_THREADS=2 pip install -v --no-build-isolation --no-deps .
+```
+
+`FLASH_MLA_DISABLE_SM100=1` compiles only the SM90 kernels for H200 / CUDA 12.8.
+Set `sparse_mla_backend="flashmla"` to enable this backend in XTuner.
+
 
 If you want to experience RL-related features in advance, you need to execute the following command to install RL-related dependencies. In addition, you need to install the inference engine of your choice. Taking LMDeploy as an example, you can refer to the [official documentation](https://github.com/InternLM/lmdeploy/) for installation.
 
@@ -279,4 +294,4 @@ The above log shows that only 10G of memory is needed to run. If you want to red
    ```bash
    pip uninstall opencv-python
    pip install opencv-python-headless
-   ```
+   ```

docs/zh_cn/get_started/installation.md

@@ -41,6 +41,20 @@ pip install git+https://github.com/InternLM/AdaptiveGEMM.git@main
 
 此外，XTuner 推荐安装 flash-attn，RL 推荐安装 flash-attn-3，能够显著提升训练速度。可以参考[官方文档](https://github.com/Dao-AILab/flash-attention)进行安装。
 
+如果需要使用 DSA Sparse MLA 的 FlashMLA 后端，可以安装精度对齐后的 FlashMLA 分支：
+
+```{code-block} shell
+:caption: 安装 FlashMLA
+
+git clone -b precise-flashmla https://github.com/DeepLink-org/flashmla-optimization.git
+cd flashmla-optimization
+git submodule update --init --recursive
+FLASH_MLA_DISABLE_SM100=1 MAX_JOBS=8 NVCC_THREADS=2 pip install -v --no-build-isolation --no-deps .
+```
+
+其中 `FLASH_MLA_DISABLE_SM100=1` 用于在 H200 / CUDA 12.8 环境下只编译 SM90 kernel。
+启用 XTuner 的 FlashMLA 后端时，设置 `sparse_mla_backend="flashmla"`。
+
 
 如果想抢先体验 RL 相关功能，则需要执行下述命令来安装 RL 部分依赖。除此之外，需要安装你选择的推理引擎。以LMDeploy为例，可参考[官网文档](https://github.com/InternLM/lmdeploy/)进行安装。
 

tests/module/attention/test_dsa_mla.py

@@ -10,6 +10,7 @@
 TestAcceleratedSparseMLA
     test_tilelang_forward_backward_matches_torch: TileLang 前反向数值与 PyTorch 后端一致。
     test_compiled_cudnn_backward_matches_tilelang: 编译后的 cuDNN DSA 前反向与 TileLang 一致。
+    test_compiled_flashmla_backward_matches_tilelang: FlashMLA 前向和 TileLang 反向与 TileLang 前反向一致。
 TestDSASequenceParallel
     test_packed_attention_matches_full_sequence: SP2 的输出、top-k 和输入梯度与完整序列一致。
     test_tilelang_indexer_matches_torch: SP2 query shard 的 TileLang indexer 与 PyTorch 一致。
@@ -74,6 +75,20 @@ def _cudnn_dsa_sparse_mla_available() -> bool:
     return result.returncode == 0
 
 
+@cache
+def _flashmla_sparse_mla_available() -> bool:
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9:
+        return False
+    result = subprocess.run(
+        [sys.executable, "-c", "from flash_mla import flash_mla_sparse_fwd"],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+        check=False,
+    )
+    return result.returncode == 0
+
+
 def _sparse_indices(seq_len: int, topk: int) -> torch.Tensor:
     indices = torch.full((seq_len, 1, topk), -1, device="cuda", dtype=torch.int64)
     for token_idx in range(seq_len):
@@ -98,6 +113,14 @@ def _cudnn_dsa_sparse_mla_inputs():
     return q, kv, _sparse_indices(seq_len, topk=64)
 
 
+def _flashmla_sparse_mla_inputs():
+    torch.manual_seed(0)
+    seq_len = 64
+    q = torch.randn(seq_len, 64, 576, device="cuda", dtype=torch.bfloat16)
+    kv = torch.randn(seq_len, 1, 576, device="cuda", dtype=torch.bfloat16)
+    return q, kv, _sparse_indices(seq_len, topk=128)
+
+
 def _tiny_dsa_attention(
     indexer_types: list[str] | None = None,
     layer_idx: int = 0,
@@ -334,6 +357,44 @@ def count_indexer_call(*_args):
         assert source_ids.dtype == torch.int32
         assert indexer_calls == 1
 
+    @pytest.mark.skipif(
+        not (_flashmla_sparse_mla_available() and _tilelang_sparse_mla_available()),
+        reason="requires CUDA, FlashMLA, and TileLang runtimes",
+    )
+    def test_compiled_flashmla_backward_matches_tilelang(self):
+        q, kv, indices = _flashmla_sparse_mla_inputs()
+        scaling = 1 / math.sqrt(q.shape[-1])
+
+        def compiled_sparse_mla(q: torch.Tensor, kv: torch.Tensor, backend: str) -> tuple[torch.Tensor, torch.Tensor]:
+            output = sparse_mla(
+                q,
+                kv,
+                indices,
+                scaling=scaling,
+                value_dim=512,
+                backend=backend,
+            )
+            return output.raw_output, output.softmax_lse
+
+        compiled_sparse_mla = torch.compile(compiled_sparse_mla, fullgraph=False)
+        q_tilelang = q.detach().clone().requires_grad_()
+        kv_tilelang = kv.detach().clone().requires_grad_()
+        q_flashmla = q.detach().clone().requires_grad_()
+        kv_flashmla = kv.detach().clone().requires_grad_()
+
+        expected, expected_lse = compiled_sparse_mla(q_tilelang, kv_tilelang, "tilelang")
+        actual, actual_lse = compiled_sparse_mla(q_flashmla, kv_flashmla, "flashmla")
+        grad_output = torch.randn_like(expected)
+        expected.backward(grad_output)
+        actual.backward(grad_output)
+
+        assert torch.equal(actual, expected)
+        assert torch.equal(actual_lse, expected_lse)
+        torch.testing.assert_close(actual, expected, atol=BF16_ATOL, rtol=BF16_RTOL)
+        torch.testing.assert_close(actual_lse, expected_lse, atol=BF16_ATOL, rtol=BF16_RTOL)
+        torch.testing.assert_close(q_flashmla.grad, q_tilelang.grad, atol=BF16_ATOL, rtol=BF16_RTOL)
+        torch.testing.assert_close(kv_flashmla.grad, kv_tilelang.grad, atol=DKV_ATOL, rtol=DKV_RTOL)
+
 
 # The multiprocess cases must run before TileLang JIT is initialized in the
 # pytest parent. With TileLang 0.1.11, spawning them afterwards crashes rank 0

xtuner/v1/model/moe/glm52/dsa_mla.py

@@ -21,6 +21,7 @@
     SparseMLABackend,
     SparseMLAProtocol,
     ensure_cudnn_dsa_runtime_available,
+    ensure_flashmla_runtime_available,
     ensure_tilelang_runtime_available,
     get_dsa_topk_indices,
     get_sparse_mla,
@@ -219,16 +220,20 @@ def build(
     ) -> "DSAMultiLatentAttention":
         if not self.freeze_dsa_indexer:
             raise ValueError("freeze_dsa_indexer=False is not supported until the indexer has a differentiable output")
-        indexer_backend = self.indexer_backend or self.sparse_mla_backend
+        indexer_backend = self.indexer_backend or (
+            "tilelang" if self.sparse_mla_backend == "flashmla" else self.sparse_mla_backend
+        )
         _validate_indexer_backend_config(
             indexer_backend,
             index_head_dim=self.index_head_dim,
             index_n_heads=self.index_n_heads,
         )
-        if self.sparse_mla_backend in ("tilelang", "cudnn_dsa"):
+        if self.sparse_mla_backend in ("tilelang", "cudnn_dsa", "flashmla"):
             ensure_tilelang_runtime_available()
         if self.sparse_mla_backend == "cudnn_dsa":
             ensure_cudnn_dsa_runtime_available()
+        if self.sparse_mla_backend == "flashmla":
+            ensure_flashmla_runtime_available()
 
         return DSAMultiLatentAttention(
             **self.model_dump(),
@@ -281,7 +286,9 @@ def __init__(
         self.indexer_rope_interleave = indexer_rope_interleave
         self.indexer_types = indexer_types
         self.sparse_mla_backend = sparse_mla_backend
-        self.indexer_backend = indexer_backend or sparse_mla_backend
+        self.indexer_backend = indexer_backend or (
+            "tilelang" if sparse_mla_backend == "flashmla" else sparse_mla_backend
+        )
         self.freeze_dsa_indexer = freeze_dsa_indexer
         self.sparse_mla_func: SparseMLAProtocol = get_sparse_mla(sparse_mla_backend)
 

xtuner/v1/ops/sparse_mla/__init__.py

@@ -18,6 +18,10 @@ def get_sparse_mla(backend: SparseMLABackend) -> SparseMLAProtocol:
         from .cudnn_dsa import cudnn_dsa_sparse_mla
 
         return cudnn_dsa_sparse_mla
+    if backend == "flashmla":
+        from .flashmla import flashmla_sparse_mla
+
+        return flashmla_sparse_mla
     raise ValueError(f"Unsupported SparseMLA backend: {backend}")
 
 
@@ -78,6 +82,12 @@ def ensure_cudnn_dsa_runtime_available() -> None:
     return _impl()
 
 
+def ensure_flashmla_runtime_available() -> None:
+    from .flashmla import ensure_flashmla_runtime_available as _impl
+
+    return _impl()
+
+
 def sparse_mla_fwd_interface(*args, **kwargs):
     from .tilelang_sparse_mla_fwd import sparse_mla_fwd_interface as _impl
 
@@ -104,6 +114,7 @@ def indexer_fwd_interface(*args, **kwargs):
     "SparseMLAProtocol",
     "dsa_topk_indices",
     "ensure_cudnn_dsa_runtime_available",
+    "ensure_flashmla_runtime_available",
     "ensure_tilelang_runtime_available",
     "get_dsa_topk_indices",
     "get_sparse_mla",

xtuner/v1/ops/sparse_mla/flashmla.py

@@ -0,0 +1,110 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+import torch
+from torch import Tensor
+
+from .protocol import SparseMLAOutputs
+from .tilelang import _validate_tilelang_sparse_mla_inputs
+
+
+def flashmla_sparse_mla(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    scaling: float | None,
+    value_dim: int | None = None,
+) -> SparseMLAOutputs:
+    _validate_flashmla_sparse_mla_inputs(q, kv, indices, value_dim)
+    indices = indices.to(torch.int32).contiguous()
+    topk_length = (indices[:, 0, :] != -1).sum(dim=-1, dtype=torch.int32).contiguous()
+    raw_output, softmax_lse = _flashmla_tilelang_sparse_mla_forward(
+        q.contiguous(), kv.contiguous(), indices, topk_length, scaling
+    )
+    return SparseMLAOutputs(raw_output=raw_output, softmax_lse=softmax_lse)
+
+
+def _validate_flashmla_sparse_mla_inputs(
+    q: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    value_dim: int | None,
+) -> None:
+    _validate_tilelang_sparse_mla_inputs(q, kv, indices, value_dim)
+    if kv.shape[1] != 1 or indices.shape[1] != 1:
+        raise RuntimeError("FlashMLA SparseMLA currently supports kv_group=1 only.")
+    if indices.shape[-1] % 128 != 0:
+        raise RuntimeError("FlashMLA SparseMLA requires topk to be divisible by 128.")
+
+
+@torch.library.custom_op("sparse_mla::flashmla_tilelang_sparse_mla_forward", mutates_args=(), device_types="cuda")
+def _flashmla_tilelang_sparse_mla_forward(
+    q: Tensor,
+    kv: Tensor,
+    indices: Tensor,
+    topk_length: Tensor,
+    scaling: float | None,
+) -> tuple[Tensor, Tensor]:
+    from flash_mla import flash_mla_sparse_fwd
+
+    raw_output, _max_logits, softmax_lse = flash_mla_sparse_fwd(
+        q,
+        kv,
+        indices,
+        sm_scale=float(scaling) if scaling is not None else q.shape[-1] ** -0.5,
+        topk_length=topk_length,
+    )
+    return raw_output, softmax_lse
+
+
+@_flashmla_tilelang_sparse_mla_forward.register_fake
+def _(
+    q: Tensor,
+    kv: Tensor,
+    indices: Tensor,
+    topk_length: Tensor,
+    scaling: float | None,
+) -> tuple[Tensor, Tensor]:
+    return q.new_empty((*q.shape[:-1], 512)), q.new_empty(q.shape[:-1], dtype=torch.float32)
+
+
+def _setup_flashmla_tilelang_sparse_mla_context(ctx, inputs, output) -> None:
+    q, kv, indices, topk_length, scaling = inputs
+    raw_output, softmax_lse = output
+    ctx.scaling = scaling
+    ctx.save_for_backward(q, kv, indices, raw_output, softmax_lse)
+
+
+def _flashmla_tilelang_sparse_mla_backward(ctx, grad_output: Tensor, grad_lse: Tensor):
+    del grad_lse
+    q, kv, indices, raw_output, softmax_lse = ctx.saved_tensors
+    from .tilelang import _tilelang_sparse_mla_backward_op
+
+    # FlashMLA returns natural-log LSE, while TileLang backward consumes log2 LSE.
+    dq, dkv = _tilelang_sparse_mla_backward_op(
+        q,
+        kv,
+        raw_output,
+        grad_output.contiguous(),
+        indices,
+        (softmax_lse / 0.6931471805599453).contiguous(),
+        ctx.scaling,
+    )
+    return dq, dkv, None, None, None
+
+
+_flashmla_tilelang_sparse_mla_forward.register_autograd(
+    _flashmla_tilelang_sparse_mla_backward,
+    setup_context=_setup_flashmla_tilelang_sparse_mla_context,
+)
+
+
+def ensure_flashmla_runtime_available() -> None:
+    try:
+        from flash_mla import flash_mla_sparse_fwd  # noqa: F401
+    except Exception as exc:
+        raise RuntimeError("FlashMLA SparseMLA requires the FlashMLA forward runtime.") from exc
+
+    if torch.cuda.is_available():
+        major, _ = torch.cuda.get_device_capability()
+        if major < 9:
+            raise RuntimeError(f"FlashMLA SparseMLA requires SM90+, found SM{major}0.")

xtuner/v1/ops/sparse_mla/protocol.py

@@ -6,7 +6,7 @@
 from xtuner.v1.data_proto import SequenceContext
 
 
-SparseMLABackend = Literal["torch", "tilelang", "cudnn_dsa"]
+SparseMLABackend = Literal["torch", "tilelang", "cudnn_dsa", "flashmla"]
 # ``deep_gemm_fp8`` names the runtime dependency and its FP8 MQA score path.
 DSAIndexerBackend = Literal["torch", "tilelang", "cudnn_dsa", "deep_gemm_fp8"]