From c4f5cd60dae386d106c9b8a12dbab24e2e9dda0b Mon Sep 17 00:00:00 2001
From: Matthew Bonanni <mbonanni@redhat.com>
Date: Mon, 13 Jul 2026 20:29:56 -0400
Subject: [PATCH] [1/N] Add dense MHA path for sparse MLA short sequences
 (#47327)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
---
 benchmarks/attention_benchmarks/benchmark.py  | 178 +++++-
 benchmarks/attention_benchmarks/common.py     | 105 ++--
 .../configs/mla_sparse_mha_vs_mqa.yaml        | 474 +++++++++++++++
 benchmarks/attention_benchmarks/mla_runner.py | 188 +++++-
 docs/design/attention_backends.md             |   2 +-
 .../test_dspark_noncausal_sparse_mla.py       |  17 +
 tests/v1/attention/test_mla_backends.py       |  13 +-
 .../v1/attention/test_sparse_mla_backends.py  | 269 +++++++++
 vllm/config/attention.py                      |   5 +
 .../layers/attention/mla_attention.py         | 568 ++++++++++--------
 .../layers/attention/sparse_mla_attention.py  | 316 ++++++++++
 vllm/v1/attention/backend.py                  |  85 +--
 .../backends/mla/flashattn_mla_sparse.py      |  99 ++-
 .../backends/mla/flashinfer_mla_sparse.py     | 131 ++--
 .../mla/flashinfer_mla_sparse_sm120.py        |   6 +-
 .../attention/backends/mla/flashmla_sparse.py | 128 ++--
 vllm/v1/attention/backends/mla/indexer.py     |   9 +-
 .../backends/mla/prefill/flash_attn.py        |  19 +-
 .../backends/mla/rocm_aiter_mla_sparse.py     |   6 +-
 vllm/v1/attention/backends/mla/sparse_swa.py  |  12 +-
 .../attention/backends/mla/xpu_mla_sparse.py  |   6 +-
 21 files changed, 1990 insertions(+), 646 deletions(-)
 create mode 100644 benchmarks/attention_benchmarks/configs/mla_sparse_mha_vs_mqa.yaml
 create mode 100644 vllm/model_executor/layers/attention/sparse_mla_attention.py

diff --git a/benchmarks/attention_benchmarks/benchmark.py b/benchmarks/attention_benchmarks/benchmark.py
index 9860d4b2d1..2fdd5ce88d 100644
--- a/benchmarks/attention_benchmarks/benchmark.py
+++ b/benchmarks/attention_benchmarks/benchmark.py
@@ -75,7 +75,11 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:
     from mla_runner import run_mla_benchmark as run_mla
 
     return run_mla(
-        config.backend, config, prefill_backend=config.prefill_backend, **kwargs
+        config.backend,
+        config,
+        prefill_backend=config.prefill_backend,
+        sparse_mla_force_mqa=config.sparse_mla_force_mqa,
+        **kwargs,
     )
 
 
@@ -592,6 +596,30 @@ def main():
         default="profile",
         help="Output file name for ncu profile (default: 'profile').",
     )
+    parser.add_argument(
+        "--torch-profile",
+        action="store_true",
+        default=False,
+        help="Collect a PyTorch profiler Chrome trace for each benchmark run.",
+    )
+    parser.add_argument(
+        "--torch-profile-dir",
+        default=None,
+        help="Directory for PyTorch profiler traces.",
+    )
+    parser.add_argument(
+        "--torch-profile-iters",
+        type=int,
+        default=3,
+        help="Number of forward passes to record per PyTorch profiler trace.",
+    )
+    parser.add_argument(
+        "--sparse-mla-mha-variants",
+        nargs="+",
+        default=None,
+        choices=["dense_mha", "mqa"],
+        help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.",
+    )
 
     # Parameter sweep (use YAML config for advanced sweeps)
     parser.add_argument(
@@ -641,6 +669,7 @@ def main():
 
         # Prefill backends (e.g., ["fa3", "fa4"])
         args.prefill_backends = yaml_config.get("prefill_backends", None)
+        args.prefill_backend = yaml_config.get("prefill_backend", None)
 
         # FP8 output benchmark knobs; CLI wins.
         if args.fp8_output_scale is None:
@@ -683,6 +712,9 @@ def main():
             args.num_q_heads = model.get("num_q_heads", args.num_q_heads)
             args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)
             args.block_size = model.get("block_size", args.block_size)
+            args.max_model_len = model.get(
+                "max_model_len", getattr(args, "max_model_len", None)
+            )
             # MLA-specific dimensions
             args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)
             args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)
@@ -701,6 +733,21 @@ def main():
             args.cuda_graphs = yaml_config["cuda_graphs"]
         if "ncu_profile" in yaml_config:
             args.ncu_profile = yaml_config["ncu_profile"]
+        if "torch_profile" in yaml_config:
+            args.torch_profile = yaml_config["torch_profile"]
+        if "torch_profile_dir" in yaml_config:
+            args.torch_profile_dir = yaml_config["torch_profile_dir"]
+        if "torch_profile_iters" in yaml_config:
+            args.torch_profile_iters = yaml_config["torch_profile_iters"]
+        args.sparse_mla_topk_pattern = yaml_config.get(
+            "sparse_mla_topk_pattern", "random"
+        )
+        args.sparse_mla_dense_mha_max_seq_len = yaml_config.get(
+            "sparse_mla_dense_mha_max_seq_len", None
+        )
+        args.sparse_mla_mha_variants = yaml_config.get(
+            "sparse_mla_mha_variants", args.sparse_mla_mha_variants
+        )
 
         # Parameter sweep configuration
         if "parameter_sweep" in yaml_config:
@@ -842,8 +889,6 @@ def main():
                         num_kv_heads=args.num_kv_heads,
                         block_size=args.block_size,
                         device=args.device,
-                        repeats=args.repeats,
-                        warmup_iters=args.warmup_iters,
                         profile_memory=args.profile_memory,
                         kv_cache_dtype=args.kv_cache_dtype,
                         use_cuda_graphs=args.cuda_graphs,
@@ -1063,6 +1108,133 @@ def main():
                     f"\n  [yellow]Prefill always faster for batch_size={bs}[/]"
                 )
 
+    # Handle MHA vs MQA comparison mode for sparse MLA
+    elif hasattr(args, "mode") and args.mode == "mha_vs_mqa":
+        console.print("[yellow]Mode: MHA vs MQA comparison for sparse MLA[/]")
+
+        sparse_mla_topk_pattern = getattr(args, "sparse_mla_topk_pattern", "random")
+        dense_mha_max_seq_len = getattr(args, "sparse_mla_dense_mha_max_seq_len", None)
+        prefill_backend = getattr(args, "prefill_backend", None)
+        if prefill_backend:
+            console.print(f"Prefill backend: {prefill_backend}")
+        available_variants = [
+            ("dense_mha", False, "dense"),
+            ("mqa", True, "auto"),
+        ]
+        requested_variants = getattr(args, "sparse_mla_mha_variants", None)
+        if requested_variants is not None:
+            valid_variants = {label for label, _, _ in available_variants}
+            invalid_variants = sorted(set(requested_variants) - valid_variants)
+            if invalid_variants:
+                raise ValueError(
+                    "Invalid sparse_mla_mha_variants entries: "
+                    f"{invalid_variants}. Valid variants are: "
+                    f"{sorted(valid_variants)}"
+                )
+            requested_variant_set = set(requested_variants)
+            variants = [
+                variant
+                for variant in available_variants
+                if variant[0] in requested_variant_set
+            ]
+        else:
+            variants = available_variants
+        formatter = ResultsFormatter(console)
+        total = 0
+        for spec in args.batch_specs:
+            q_len = max(request.q_len for request in parse_batch_spec(spec))
+            for variant_label, _, _ in variants:
+                if (
+                    variant_label == "dense_mha"
+                    and dense_mha_max_seq_len is not None
+                    and q_len > dense_mha_max_seq_len
+                ):
+                    continue
+                total += len(backends)
+
+        with tqdm(total=total, desc="Benchmarking") as pbar:
+            for spec in args.batch_specs:
+                q_len = max(request.q_len for request in parse_batch_spec(spec))
+                for backend in backends:
+                    for variant_label, force_mqa, mha_mode in variants:
+                        if (
+                            variant_label == "dense_mha"
+                            and dense_mha_max_seq_len is not None
+                            and q_len > dense_mha_max_seq_len
+                        ):
+                            continue
+                        config = BenchmarkConfig(
+                            backend=f"{backend}_{variant_label}",
+                            batch_spec=spec,
+                            num_layers=args.num_layers,
+                            head_dim=args.head_dim,
+                            num_q_heads=args.num_q_heads,
+                            num_kv_heads=args.num_kv_heads,
+                            block_size=args.block_size,
+                            device=args.device,
+                            max_model_len=getattr(args, "max_model_len", None),
+                            kv_cache_dtype=args.kv_cache_dtype,
+                            profile_memory=args.profile_memory,
+                            use_cuda_graphs=args.cuda_graphs,
+                            ncu_profile=args.ncu_profile,
+                            torch_profile=args.torch_profile,
+                            torch_profile_dir=args.torch_profile_dir,
+                            torch_profile_iters=args.torch_profile_iters,
+                            warmup_ms=args.warmup_ms,
+                            kv_lora_rank=getattr(args, "kv_lora_rank", None),
+                            qk_nope_head_dim=getattr(args, "qk_nope_head_dim", None),
+                            qk_rope_head_dim=getattr(args, "qk_rope_head_dim", None),
+                            v_head_dim=getattr(args, "v_head_dim", None),
+                            sparse_mla_force_mqa=force_mqa,
+                            sparse_mla_mha_mode=mha_mode,
+                            sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len,
+                            sparse_mla_topk_pattern=sparse_mla_topk_pattern,
+                            prefill_backend=prefill_backend,
+                        )
+
+                        # run_mla_benchmark needs the real backend name
+                        from mla_runner import run_mla_benchmark as run_mla
+
+                        run_label = f"{backend}_{variant_label} {spec}"
+                        pbar.set_postfix_str(run_label)
+
+                        try:
+                            result = run_mla(
+                                backend,
+                                config,
+                                prefill_backend=prefill_backend,
+                                sparse_mla_force_mqa=force_mqa,
+                            )
+                        except Exception as e:
+                            result = BenchmarkResult(
+                                config=config,
+                                mean_time=float("inf"),
+                                median_time=float("inf"),
+                                std_time=0,
+                                min_time=float("inf"),
+                                max_time=float("inf"),
+                                error=str(e),
+                            )
+
+                        all_results.append(result)
+                        if args.output_csv:
+                            formatter.save_csv(all_results, args.output_csv)
+                        if args.output_json:
+                            formatter.save_json(all_results, args.output_json)
+
+                        if not result.success:
+                            console.print(
+                                f"[red]Error {backend}_{variant_label} "
+                                f"{spec}: {result.error}[/]"
+                            )
+
+                        pbar.update(1)
+
+        # Display results with variant labels as separate "backends"
+        console.print("\n[bold green]MHA vs MQA Results:[/]")
+        variant_backends = [f"{b}_{v}" for b in backends for v, _, _ in variants]
+        formatter.print_table(all_results, variant_backends)
+
     # Handle model parameter sweep mode
     elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:
         # Model parameter sweep
diff --git a/benchmarks/attention_benchmarks/common.py b/benchmarks/attention_benchmarks/common.py
index 106d785480..b8ca5739c5 100644
--- a/benchmarks/attention_benchmarks/common.py
+++ b/benchmarks/attention_benchmarks/common.py
@@ -4,8 +4,10 @@
 """Common utilities for attention benchmarking."""
 
 import csv
+import gc
 import json
 import math
+from collections.abc import Sequence
 from dataclasses import asdict, dataclass
 from pathlib import Path
 from typing import Any
@@ -44,10 +46,13 @@ def run_do_bench(
     kwargs: dict[str, Any] = {"return_mode": "all"}
     if use_cuda_graphs:
         result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)
+        gc.collect()
+        torch.accelerator.empty_cache()
     else:
         if warmup_ms is not None:
             kwargs["warmup"] = warmup_ms
         result = triton.testing.do_bench(benchmark_fn, **kwargs)
+    torch.accelerator.synchronize()
     return result
 
 
@@ -91,42 +96,6 @@ except ImportError:
     AttentionLayerBase = object  # Fallback
 
 
-class MockKVBProj:
-    """Mock KV projection layer for MLA prefill mode.
-
-    Mimics ColumnParallelLinear behavior for kv_b_proj in MLA backends.
-    Projects kv_c_normed to [qk_nope_head_dim + v_head_dim] per head.
-    """
-
-    def __init__(self, num_heads: int, qk_nope_head_dim: int, v_head_dim: int):
-        self.num_heads = num_heads
-        self.qk_nope_head_dim = qk_nope_head_dim
-        self.v_head_dim = v_head_dim
-        self.out_dim = qk_nope_head_dim + v_head_dim
-        self.weight = torch.empty(0, dtype=torch.bfloat16)
-
-    def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:
-        """
-        Project kv_c_normed to output space.
-
-        Args:
-            x: Input tensor [num_tokens, kv_lora_rank]
-
-        Returns:
-            Tuple containing output tensor
-                [num_tokens, num_heads, qk_nope_head_dim + v_head_dim]
-        """
-        num_tokens = x.shape[0]
-        result = torch.randn(
-            num_tokens,
-            self.num_heads,
-            self.out_dim,
-            device=x.device,
-            dtype=x.dtype,
-        )
-        return (result,)  # Return as tuple to match ColumnParallelLinear API
-
-
 class MockIndexer:
     """Mock Indexer for sparse MLA backends.
 
@@ -158,6 +127,60 @@ class MockIndexer:
         )
         self.topk_indices_buffer[:num_tokens] = indices
 
+    def fill_indices(
+        self,
+        num_tokens: int,
+        max_kv_len: int,
+        pattern: str = "random",
+        requests: Sequence[Any] | None = None,
+    ):
+        if pattern == "random":
+            self.fill_random_indices(num_tokens, max_kv_len)
+            return
+        if pattern == "prefix":
+            indices = torch.arange(
+                self.topk_tokens,
+                dtype=torch.int32,
+                device=self.topk_indices_buffer.device,
+            )
+            indices = (indices % max_kv_len).expand(num_tokens, -1)
+            self.topk_indices_buffer[:num_tokens] = indices
+            return
+        if pattern == "sliding_window":
+            if requests is None:
+                start = max(max_kv_len - self.topk_tokens, 0)
+                indices = torch.arange(
+                    start,
+                    start + self.topk_tokens,
+                    dtype=torch.int32,
+                    device=self.topk_indices_buffer.device,
+                )
+                indices = indices.clamp(max=max_kv_len - 1).expand(num_tokens, -1)
+                self.topk_indices_buffer[:num_tokens] = indices
+                return
+
+            rows = []
+            offsets = torch.arange(
+                self.topk_tokens,
+                dtype=torch.int32,
+                device=self.topk_indices_buffer.device,
+            ) - (self.topk_tokens - 1)
+            for request in requests:
+                q_len = request.q_len
+                kv_len = request.kv_len
+                context_len = kv_len - q_len
+                positions = torch.arange(
+                    context_len,
+                    kv_len,
+                    dtype=torch.int32,
+                    device=self.topk_indices_buffer.device,
+                )
+                row_indices = positions[:, None] + offsets[None, :]
+                rows.append(row_indices.clamp(min=0, max=kv_len - 1))
+            self.topk_indices_buffer[:num_tokens] = torch.cat(rows, dim=0)
+            return
+        raise ValueError(f"Unknown sparse MLA topk pattern: {pattern}")
+
 
 class MockLayer(AttentionLayerBase):
     """Mock attention layer with scale parameters and impl.
@@ -252,10 +275,14 @@ class BenchmarkConfig:
     num_kv_heads: int
     block_size: int
     device: str
+    max_model_len: int | None = None
     dtype: torch.dtype = torch.float16
     profile_memory: bool = False
-    use_cuda_graphs: bool = False
+    use_cuda_graphs: bool = True
     ncu_profile: bool = False
+    torch_profile: bool = False
+    torch_profile_dir: str | None = None
+    torch_profile_iters: int = 3
     warmup_ms: int | None = None
 
     # "auto" or "fp8"
@@ -271,6 +298,10 @@ class BenchmarkConfig:
     # Backend-specific tuning
     num_kv_splits: int | None = None  # CUTLASS MLA
     reorder_batch_threshold: int | None = None  # FlashAttn MLA, FlashMLA
+    sparse_mla_force_mqa: bool = False  # Force MQA path for sparse MLA
+    sparse_mla_mha_mode: str = "auto"  # "auto" or "dense"
+    sparse_mla_dense_mha_max_seq_len: int | None = None
+    sparse_mla_topk_pattern: str = "random"  # "random", "prefix", "sliding_window"
     num_splits: int | None = None  # FlashAttention split-K (0=auto, 1=disabled)
 
 
diff --git a/benchmarks/attention_benchmarks/configs/mla_sparse_mha_vs_mqa.yaml b/benchmarks/attention_benchmarks/configs/mla_sparse_mha_vs_mqa.yaml
new file mode 100644
index 0000000000..fd3911ea19
--- /dev/null
+++ b/benchmarks/attention_benchmarks/configs/mla_sparse_mha_vs_mqa.yaml
@@ -0,0 +1,474 @@
+# Sparse MLA benchmark: forward_mha vs forward_mqa
+#
+# Usage:
+#   python benchmark.py --config configs/mla_sparse_mha_vs_mqa.yaml
+#
+# Heatmap grid:
+#   - batch_size: 1, 2, 4, 8, 16, 32
+#   - seq_len: 32, 64, 128, 256, 512, 1024, 2048
+#   - q_len: powers of two through seq_len
+#
+# Specs with q_len < seq_len include context; the q_len == seq_len diagonal
+# covers pure prefill.
+# The model shape below is the DP case. For the TP8 run, manually change
+# model.num_q_heads from 128 to 16 before rerunning this benchmark.
+
+mode: mha_vs_mqa
+
+model:
+  name: "deepseek-v3"
+  num_layers: 60
+  num_q_heads: 128
+  num_kv_heads: 1
+  head_dim: 576
+  kv_lora_rank: 512
+  qk_nope_head_dim: 128
+  qk_rope_head_dim: 64
+  v_head_dim: 128
+  block_size: 128
+  max_model_len: 2048
+
+batch_specs:
+  # Batch size 1
+  # seq_len = 32
+  - "1q1s32"
+  - "1q2s32"
+  - "1q4s32"
+  - "1q8s32"
+  - "1q16s32"
+  - "1q32"
+  # seq_len = 64
+  - "1q1s64"
+  - "1q2s64"
+  - "1q4s64"
+  - "1q8s64"
+  - "1q16s64"
+  - "1q32s64"
+  - "1q64"
+  # seq_len = 128
+  - "1q1s128"
+  - "1q2s128"
+  - "1q4s128"
+  - "1q8s128"
+  - "1q16s128"
+  - "1q32s128"
+  - "1q64s128"
+  - "1q128"
+  # seq_len = 256
+  - "1q1s256"
+  - "1q2s256"
+  - "1q4s256"
+  - "1q8s256"
+  - "1q16s256"
+  - "1q32s256"
+  - "1q64s256"
+  - "1q128s256"
+  - "1q256"
+  # seq_len = 512
+  - "1q1s512"
+  - "1q2s512"
+  - "1q4s512"
+  - "1q8s512"
+  - "1q16s512"
+  - "1q32s512"
+  - "1q64s512"
+  - "1q128s512"
+  - "1q256s512"
+  - "1q512"
+  # seq_len = 1024
+  - "1q1s1024"
+  - "1q2s1024"
+  - "1q4s1024"
+  - "1q8s1024"
+  - "1q16s1024"
+  - "1q32s1024"
+  - "1q64s1024"
+  - "1q128s1024"
+  - "1q256s1024"
+  - "1q512s1024"
+  - "1q1024"
+  # seq_len = 2048
+  - "1q1s2048"
+  - "1q2s2048"
+  - "1q4s2048"
+  - "1q8s2048"
+  - "1q16s2048"
+  - "1q32s2048"
+  - "1q64s2048"
+  - "1q128s2048"
+  - "1q256s2048"
+  - "1q512s2048"
+  - "1q1024s2048"
+  - "1q2048"
+
+  # Batch size 2
+  # seq_len = 32
+  - "2q1s32"
+  - "2q2s32"
+  - "2q4s32"
+  - "2q8s32"
+  - "2q16s32"
+  - "2q32"
+  # seq_len = 64
+  - "2q1s64"
+  - "2q2s64"
+  - "2q4s64"
+  - "2q8s64"
+  - "2q16s64"
+  - "2q32s64"
+  - "2q64"
+  # seq_len = 128
+  - "2q1s128"
+  - "2q2s128"
+  - "2q4s128"
+  - "2q8s128"
+  - "2q16s128"
+  - "2q32s128"
+  - "2q64s128"
+  - "2q128"
+  # seq_len = 256
+  - "2q1s256"
+  - "2q2s256"
+  - "2q4s256"
+  - "2q8s256"
+  - "2q16s256"
+  - "2q32s256"
+  - "2q64s256"
+  - "2q128s256"
+  - "2q256"
+  # seq_len = 512
+  - "2q1s512"
+  - "2q2s512"
+  - "2q4s512"
+  - "2q8s512"
+  - "2q16s512"
+  - "2q32s512"
+  - "2q64s512"
+  - "2q128s512"
+  - "2q256s512"
+  - "2q512"
+  # seq_len = 1024
+  - "2q1s1024"
+  - "2q2s1024"
+  - "2q4s1024"
+  - "2q8s1024"
+  - "2q16s1024"
+  - "2q32s1024"
+  - "2q64s1024"
+  - "2q128s1024"
+  - "2q256s1024"
+  - "2q512s1024"
+  - "2q1024"
+  # seq_len = 2048
+  - "2q1s2048"
+  - "2q2s2048"
+  - "2q4s2048"
+  - "2q8s2048"
+  - "2q16s2048"
+  - "2q32s2048"
+  - "2q64s2048"
+  - "2q128s2048"
+  - "2q256s2048"
+  - "2q512s2048"
+  - "2q1024s2048"
+  - "2q2048"
+
+  # Batch size 4
+  # seq_len = 32
+  - "4q1s32"
+  - "4q2s32"
+  - "4q4s32"
+  - "4q8s32"
+  - "4q16s32"
+  - "4q32"
+  # seq_len = 64
+  - "4q1s64"
+  - "4q2s64"
+  - "4q4s64"
+  - "4q8s64"
+  - "4q16s64"
+  - "4q32s64"
+  - "4q64"
+  # seq_len = 128
+  - "4q1s128"
+  - "4q2s128"
+  - "4q4s128"
+  - "4q8s128"
+  - "4q16s128"
+  - "4q32s128"
+  - "4q64s128"
+  - "4q128"
+  # seq_len = 256
+  - "4q1s256"
+  - "4q2s256"
+  - "4q4s256"
+  - "4q8s256"
+  - "4q16s256"
+  - "4q32s256"
+  - "4q64s256"
+  - "4q128s256"
+  - "4q256"
+  # seq_len = 512
+  - "4q1s512"
+  - "4q2s512"
+  - "4q4s512"
+  - "4q8s512"
+  - "4q16s512"
+  - "4q32s512"
+  - "4q64s512"
+  - "4q128s512"
+  - "4q256s512"
+  - "4q512"
+  # seq_len = 1024
+  - "4q1s1024"
+  - "4q2s1024"
+  - "4q4s1024"
+  - "4q8s1024"
+  - "4q16s1024"
+  - "4q32s1024"
+  - "4q64s1024"
+  - "4q128s1024"
+  - "4q256s1024"
+  - "4q512s1024"
+  - "4q1024"
+  # seq_len = 2048
+  - "4q1s2048"
+  - "4q2s2048"
+  - "4q4s2048"
+  - "4q8s2048"
+  - "4q16s2048"
+  - "4q32s2048"
+  - "4q64s2048"
+  - "4q128s2048"
+  - "4q256s2048"
+  - "4q512s2048"
+  - "4q1024s2048"
+  - "4q2048"
+
+  # Batch size 8
+  # seq_len = 32
+  - "8q1s32"
+  - "8q2s32"
+  - "8q4s32"
+  - "8q8s32"
+  - "8q16s32"
+  - "8q32"
+  # seq_len = 64
+  - "8q1s64"
+  - "8q2s64"
+  - "8q4s64"
+  - "8q8s64"
+  - "8q16s64"
+  - "8q32s64"
+  - "8q64"
+  # seq_len = 128
+  - "8q1s128"
+  - "8q2s128"
+  - "8q4s128"
+  - "8q8s128"
+  - "8q16s128"
+  - "8q32s128"
+  - "8q64s128"
+  - "8q128"
+  # seq_len = 256
+  - "8q1s256"
+  - "8q2s256"
+  - "8q4s256"
+  - "8q8s256"
+  - "8q16s256"
+  - "8q32s256"
+  - "8q64s256"
+  - "8q128s256"
+  - "8q256"
+  # seq_len = 512
+  - "8q1s512"
+  - "8q2s512"
+  - "8q4s512"
+  - "8q8s512"
+  - "8q16s512"
+  - "8q32s512"
+  - "8q64s512"
+  - "8q128s512"
+  - "8q256s512"
+  - "8q512"
+  # seq_len = 1024
+  - "8q1s1024"
+  - "8q2s1024"
+  - "8q4s1024"
+  - "8q8s1024"
+  - "8q16s1024"
+  - "8q32s1024"
+  - "8q64s1024"
+  - "8q128s1024"
+  - "8q256s1024"
+  - "8q512s1024"
+  - "8q1024"
+  # seq_len = 2048
+  - "8q1s2048"
+  - "8q2s2048"
+  - "8q4s2048"
+  - "8q8s2048"
+  - "8q16s2048"
+  - "8q32s2048"
+  - "8q64s2048"
+  - "8q128s2048"
+  - "8q256s2048"
+  - "8q512s2048"
+  - "8q1024s2048"
+  - "8q2048"
+
+  # Batch size 16
+  # seq_len = 32
+  - "16q1s32"
+  - "16q2s32"
+  - "16q4s32"
+  - "16q8s32"
+  - "16q16s32"
+  - "16q32"
+  # seq_len = 64
+  - "16q1s64"
+  - "16q2s64"
+  - "16q4s64"
+  - "16q8s64"
+  - "16q16s64"
+  - "16q32s64"
+  - "16q64"
+  # seq_len = 128
+  - "16q1s128"
+  - "16q2s128"
+  - "16q4s128"
+  - "16q8s128"
+  - "16q16s128"
+  - "16q32s128"
+  - "16q64s128"
+  - "16q128"
+  # seq_len = 256
+  - "16q1s256"
+  - "16q2s256"
+  - "16q4s256"
+  - "16q8s256"
+  - "16q16s256"
+  - "16q32s256"
+  - "16q64s256"
+  - "16q128s256"
+  - "16q256"
+  # seq_len = 512
+  - "16q1s512"
+  - "16q2s512"
+  - "16q4s512"
+  - "16q8s512"
+  - "16q16s512"
+  - "16q32s512"
+  - "16q64s512"
+  - "16q128s512"
+  - "16q256s512"
+  - "16q512"
+  # seq_len = 1024
+  - "16q1s1024"
+  - "16q2s1024"
+  - "16q4s1024"
+  - "16q8s1024"
+  - "16q16s1024"
+  - "16q32s1024"
+  - "16q64s1024"
+  - "16q128s1024"
+  - "16q256s1024"
+  - "16q512s1024"
+  - "16q1024"
+  # seq_len = 2048
+  - "16q1s2048"
+  - "16q2s2048"
+  - "16q4s2048"
+  - "16q8s2048"
+  - "16q16s2048"
+  - "16q32s2048"
+  - "16q64s2048"
+  - "16q128s2048"
+  - "16q256s2048"
+  - "16q512s2048"
+  - "16q1024s2048"
+  - "16q2048"
+
+  # Batch size 32
+  # seq_len = 32
+  - "32q1s32"
+  - "32q2s32"
+  - "32q4s32"
+  - "32q8s32"
+  - "32q16s32"
+  - "32q32"
+  # seq_len = 64
+  - "32q1s64"
+  - "32q2s64"
+  - "32q4s64"
+  - "32q8s64"
+  - "32q16s64"
+  - "32q32s64"
+  - "32q64"
+  # seq_len = 128
+  - "32q1s128"
+  - "32q2s128"
+  - "32q4s128"
+  - "32q8s128"
+  - "32q16s128"
+  - "32q32s128"
+  - "32q64s128"
+  - "32q128"
+  # seq_len = 256
+  - "32q1s256"
+  - "32q2s256"
+  - "32q4s256"
+  - "32q8s256"
+  - "32q16s256"
+  - "32q32s256"
+  - "32q64s256"
+  - "32q128s256"
+  - "32q256"
+  # seq_len = 512
+  - "32q1s512"
+  - "32q2s512"
+  - "32q4s512"
+  - "32q8s512"
+  - "32q16s512"
+  - "32q32s512"
+  - "32q64s512"
+  - "32q128s512"
+  - "32q256s512"
+  - "32q512"
+  # seq_len = 1024
+  - "32q1s1024"
+  - "32q2s1024"
+  - "32q4s1024"
+  - "32q8s1024"
+  - "32q16s1024"
+  - "32q32s1024"
+  - "32q64s1024"
+  - "32q128s1024"
+  - "32q256s1024"
+  - "32q512s1024"
+  - "32q1024"
+  # seq_len = 2048
+  - "32q1s2048"
+  - "32q2s2048"
+  - "32q4s2048"
+  - "32q8s2048"
+  - "32q16s2048"
+  - "32q32s2048"
+  - "32q64s2048"
+  - "32q128s2048"
+  - "32q256s2048"
+  - "32q512s2048"
+  - "32q1024s2048"
+  - "32q2048"
+
+backends:
+  - FLASHMLA_SPARSE
+
+device: "cuda:0"
+profile_memory: false
+sparse_mla_dense_mha_max_seq_len: 2048
+sparse_mla_topk_pattern: "random"
+
+output:
+  csv: "benchmark_output/mla_sparse_mha_vs_mqa.csv"
+  json: "benchmark_output/mla_sparse_mha_vs_mqa.json"
diff --git a/benchmarks/attention_benchmarks/mla_runner.py b/benchmarks/attention_benchmarks/mla_runner.py
index c9b3fb29bb..f78aca3bce 100644
--- a/benchmarks/attention_benchmarks/mla_runner.py
+++ b/benchmarks/attention_benchmarks/mla_runner.py
@@ -9,6 +9,8 @@ needing full VllmConfig integration.
 """
 
 import statistics
+import tempfile
+from pathlib import Path
 
 import numpy as np
 import torch
@@ -17,7 +19,6 @@ from common import (
     BenchmarkResult,
     MockHfConfig,
     MockIndexer,
-    MockKVBProj,
     MockLayer,
     run_do_bench,
     run_ncu_profile,
@@ -33,8 +34,59 @@ from vllm.config import (
     VllmConfig,
     set_current_vllm_config,
 )
+from vllm.model_executor.layers.linear import ColumnParallelLinear
 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum
 
+
+def _safe_profile_name(value: str) -> str:
+    return "".join(c if c.isalnum() or c in "._-" else "_" for c in value)
+
+
+def _create_kv_b_proj(
+    mla_dims: dict,
+    device: torch.device,
+):
+    kv_b_proj = ColumnParallelLinear(
+        mla_dims["kv_lora_rank"],
+        mla_dims["num_q_heads"]
+        * (mla_dims["qk_nope_head_dim"] + mla_dims["v_head_dim"]),
+        bias=False,
+        params_dtype=torch.bfloat16,
+        quant_config=None,
+        prefix="benchmark.kv_b_proj",
+    ).to(device)
+    with torch.no_grad():
+        kv_b_proj.weight.copy_(torch.randn_like(kv_b_proj.weight))
+    return kv_b_proj
+
+
+def _ensure_single_rank_model_parallel() -> None:
+    import torch.distributed as dist
+
+    from vllm.distributed import (
+        ensure_model_parallel_initialized,
+        init_distributed_environment,
+        model_parallel_is_initialized,
+    )
+
+    if not dist.is_available():
+        return
+    if not dist.is_initialized():
+        with tempfile.NamedTemporaryFile(
+            prefix="vllm_bench_dist_", delete=False
+        ) as init_file:
+            distributed_init_method = f"file://{init_file.name}"
+        init_distributed_environment(
+            world_size=1,
+            rank=0,
+            distributed_init_method=distributed_init_method,
+            local_rank=0,
+            backend="nccl",
+        )
+    if not model_parallel_is_initialized():
+        ensure_model_parallel_initialized(1, 1)
+
+
 # ============================================================================
 # VllmConfig Creation
 # ============================================================================
@@ -66,10 +118,12 @@ def create_minimal_vllm_config(
     block_size: int = 128,
     max_num_seqs: int = 256,
     max_num_batched_tokens: int = 8192,
+    max_model_len: int = 32768,
     mla_dims: dict | None = None,
     index_topk: int | None = None,
     prefill_backend: str | None = None,
     kv_cache_dtype: str = "auto",
+    sparse_mla_force_mqa: bool = False,
 ) -> VllmConfig:
     """
     Create minimal VllmConfig for MLA benchmarks.
@@ -86,6 +140,8 @@ def create_minimal_vllm_config(
         prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",
                         "trtllm"). Configures the attention config to force
                         the specified prefill backend.
+        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
+                    forward_mqa (even prefill tokens).
 
     Returns:
         VllmConfig for benchmarking
@@ -131,7 +187,7 @@ def create_minimal_vllm_config(
             trust_remote_code=True,
             dtype="bfloat16",
             seed=0,
-            max_model_len=32768,
+            max_model_len=max_model_len,
             quantization=None,
             enforce_eager=False,
             max_logprobs=20,
@@ -163,7 +219,7 @@ def create_minimal_vllm_config(
     scheduler_config = SchedulerConfig(
         max_num_seqs=max_num_seqs,
         max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),
-        max_model_len=32768,
+        max_model_len=max_model_len,
         is_encoder_decoder=False,
         enable_chunked_prefill=True,
     )
@@ -192,6 +248,9 @@ def create_minimal_vllm_config(
                 "flash_attn_version"
             ]
 
+    if sparse_mla_force_mqa:
+        vllm_config.attention_config.sparse_mla_force_mqa = True
+
     return vllm_config
 
 
@@ -548,12 +607,7 @@ def _create_backend_impl(
     # Calculate scale
     scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])
 
-    # Create mock kv_b_proj layer for prefill mode
-    mock_kv_b_proj = MockKVBProj(
-        num_heads=mla_dims["num_q_heads"],
-        qk_nope_head_dim=mla_dims["qk_nope_head_dim"],
-        v_head_dim=mla_dims["v_head_dim"],
-    )
+    kv_b_proj = _create_kv_b_proj(mla_dims, device)
 
     # Create indexer for sparse backends
     indexer = None
@@ -584,7 +638,7 @@ def _create_backend_impl(
         "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],
         "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],
         "v_head_dim": mla_dims["v_head_dim"],
-        "kv_b_proj": mock_kv_b_proj,
+        "kv_b_proj": kv_b_proj,
     }
 
     # Add indexer for sparse backends
@@ -785,14 +839,35 @@ def _run_single_benchmark(
     # Fill indexer with random indices for sparse backends
     is_sparse = backend_cfg.get("is_sparse", False)
     if is_sparse and indexer is not None:
-        indexer.fill_random_indices(total_q, max_kv_len)
+        indexer.fill_indices(
+            total_q,
+            max_kv_len,
+            getattr(config, "sparse_mla_topk_pattern", "random"),
+        )
 
     # Determine which forward methods to use based on metadata.
-    # Sparse MLA backends always use forward_mqa
-    has_decode = is_sparse or getattr(metadata, "decode", None) is not None
-    has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None
+    # Non-sparse backends use .decode/.prefill sub-objects.
+    # Sparse backends use num_decode_tokens/num_prefills directly.
+    #
+    # sparse_mla_force_mqa overrides: even for prefill metadata, use MQA.
+    force_mqa = getattr(config, "sparse_mla_force_mqa", False)
+    force_dense_mha = getattr(config, "sparse_mla_mha_mode", "auto") == "dense"
+    if force_mqa:
+        has_decode = True
+        has_prefill = False
+    elif is_sparse:
+        has_decode = metadata.num_decode_tokens > 0
+        has_prefill = metadata.num_prefills > 0
+    else:
+        has_decode = metadata.decode is not None
+        has_prefill = metadata.prefill is not None
     if not has_decode and not has_prefill:
         raise RuntimeError("Metadata has neither decode nor prefill metadata")
+    if is_sparse and force_dense_mha and not has_prefill:
+        raise RuntimeError(
+            "Sparse MLA dense_mha benchmark did not produce prefill metadata. "
+            "Check reorder_batch_threshold/path forcing."
+        )
 
     num_decode = (
         metadata.num_decode_tokens
@@ -871,7 +946,6 @@ def _run_single_benchmark(
                 metadata,
                 prefill_inputs["k_scale"],
                 prefill_fp8_output if fused_output else prefill_inputs["output"],
-                prefill_output_scale if fused_output else None,
             )
             if fused_output:
                 out = prefill_fp8_output
@@ -898,6 +972,48 @@ def _run_single_benchmark(
             throughput_tokens_per_sec=0.0,
         )
 
+    if config.torch_profile:
+        profile_dir = Path(
+            config.torch_profile_dir or "benchmark_outputs/torch_profiles"
+        )
+        profile_dir.mkdir(parents=True, exist_ok=True)
+        trace_name = _safe_profile_name(f"{config.backend}_{config.batch_spec}")
+        trace_path = profile_dir / f"{trace_name}.json"
+        iters = max(config.torch_profile_iters, 1)
+
+        forward_fn()
+        torch.accelerator.synchronize()
+        with torch.profiler.profile(
+            activities=[
+                torch.profiler.ProfilerActivity.CPU,
+                torch.profiler.ProfilerActivity.CUDA,
+            ],
+            record_shapes=True,
+            profile_memory=True,
+            with_stack=False,
+        ) as prof:
+            for _ in range(iters):
+                forward_fn()
+                torch.accelerator.synchronize()
+                prof.step()
+        prof.export_chrome_trace(str(trace_path))
+        print(f"Saved PyTorch profiler trace to {trace_path}")
+        print(
+            prof.key_averages().table(
+                sort_by="cuda_time_total",
+                row_limit=25,
+            )
+        )
+        return BenchmarkResult(
+            config=config,
+            mean_time=0.0,
+            median_time=0.0,
+            std_time=0.0,
+            min_time=0.0,
+            max_time=0.0,
+            throughput_tokens_per_sec=0.0,
+        )
+
     all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)
 
     # Convert ms to seconds per layer
@@ -920,6 +1036,7 @@ def _run_mla_benchmark_batched(
     configs_with_params: list[tuple],  # [(config, threshold, num_splits), ...]
     index_topk: int = 2048,
     prefill_backend: str | None = None,
+    sparse_mla_force_mqa: bool = False,
     output_scale: float | None = None,
     fuse_quant_op: bool = False,
 ) -> list[BenchmarkResult]:
@@ -940,6 +1057,8 @@ def _run_mla_benchmark_batched(
         index_topk: Topk value for sparse MLA backends (default 2048)
         prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
             When set, forces the specified FlashAttention version for prefill.
+        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
+            forward_mqa (even prefill tokens).
 
     Returns:
         List of BenchmarkResult objects
@@ -980,21 +1099,41 @@ def _run_mla_benchmark_batched(
         sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))
         for cfg, *_ in configs_with_params
     )
+    max_model_len = max(
+        max_total_q,
+        max(
+            getattr(cfg, "max_model_len", None) or 32768
+            for cfg, *_ in configs_with_params
+        ),
+    )
 
     # Create and set vLLM config for MLA (reused across all benchmarks)
     vllm_config = create_minimal_vllm_config(
         model_name="deepseek-v3",  # Used only for model path
         block_size=block_size,
         max_num_batched_tokens=max_total_q,
+        max_model_len=max_model_len,
         mla_dims=mla_dims,  # Use custom dims from config or default
         index_topk=index_topk if is_sparse else None,
         prefill_backend=prefill_backend,
         kv_cache_dtype=kv_cache_dtype,
+        sparse_mla_force_mqa=sparse_mla_force_mqa,
     )
 
     results = []
 
+    # Initialize workspace manager (needed by metadata builders)
+    from vllm.v1.worker.workspace import (
+        init_workspace_manager,
+        is_workspace_manager_initialized,
+    )
+
+    if not is_workspace_manager_initialized():
+        init_workspace_manager(device)
+
     with set_current_vllm_config(vllm_config):
+        _ensure_single_rank_model_parallel()
+
         # Create backend impl, layer, builder, and indexer (reused across benchmarks)
         impl, layer, builder_instance, indexer = _create_backend_impl(
             backend_cfg,
@@ -1040,9 +1179,20 @@ def _run_mla_benchmark_batched(
         for config, threshold, num_splits in configs_with_params:
             # Set threshold for this benchmark (FlashAttn/FlashMLA only)
             original_threshold = None
-            if threshold is not None and builder_instance:
+            effective_threshold = threshold
+            force_dense_mha = (
+                is_sparse
+                and getattr(config, "sparse_mla_mha_mode", "auto") == "dense"
+                and not getattr(config, "sparse_mla_force_mqa", False)
+            )
+            if force_dense_mha:
+                # Sparse MLA normally treats q_len <= 1 as decode. Use an
+                # impossible threshold so dense_mha benchmarks actually run
+                # the prefill/MHA path, including q_len=1 short extends.
+                effective_threshold = -1
+            if effective_threshold is not None and builder_instance:
                 original_threshold = builder_instance.reorder_batch_threshold
-                builder_instance.reorder_batch_threshold = threshold
+                builder_instance.reorder_batch_threshold = effective_threshold
 
             # Set num_splits for CUTLASS
             original_num_splits = None
@@ -1090,6 +1240,7 @@ def run_mla_benchmark(
     num_kv_splits: int | None = None,
     index_topk: int = 2048,
     prefill_backend: str | None = None,
+    sparse_mla_force_mqa: bool = False,
     output_scale: float | None = None,
     fuse_quant_op: bool = False,
 ) -> BenchmarkResult | list[BenchmarkResult]:
@@ -1111,6 +1262,8 @@ def run_mla_benchmark(
         index_topk: Topk value for sparse MLA backends (default 2048)
         prefill_backend: Prefill backend name (e.g., "fa3", "fa4").
             When set, forces the specified FlashAttention version for prefill.
+        sparse_mla_force_mqa: If True, forces all sparse MLA tokens through
+            forward_mqa (even prefill tokens).
         output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).
         fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill
             kernel vs a standalone post-quant kernel. See _run_single_benchmark.
@@ -1142,6 +1295,7 @@ def run_mla_benchmark(
         configs_with_params,
         index_topk,
         prefill_backend=prefill_backend,
+        sparse_mla_force_mqa=sparse_mla_force_mqa,
         output_scale=output_scale,
         fuse_quant_op=fuse_quant_op,
     )
diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md
index 4bf3822f3a..c8d02fb640 100644
--- a/docs/design/attention_backends.md
+++ b/docs/design/attention_backends.md
@@ -205,7 +205,7 @@ hardware and configuration.
 
 | Backend | Description | Dtypes | Compute Cap. | Notes |
 | ------- | ----------- | ------ | ------------ | ----- |
-| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | FA4 on SM100+, FA3 on SM90, FA2 otherwise |
+| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
 | `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
 | `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
 | `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
diff --git a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py
index ebb29c9af0..71579c025b 100644
--- a/tests/v1/attention/test_dspark_noncausal_sparse_mla.py
+++ b/tests/v1/attention/test_dspark_noncausal_sparse_mla.py
@@ -304,6 +304,23 @@ def _run_sparse_backend_vs_sdpa(
         scale=kv_cache_scale,
     )
 
+    # The sparse builder clones the layer's dense-MHA prefill backend from
+    # static_forward_context; register a mock layer carrying one.
+    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+
+    prefill_backend = get_mla_prefill_backend(vllm_config)(
+        num_heads=num_heads,
+        scale=scale,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        vllm_config=vllm_config,
+    )
+    vllm_config.compilation_config.static_forward_context["placeholder"] = (
+        SimpleNamespace(prefill_backend=prefill_backend)
+    )
+
     builder = backend_cls.get_builder_cls()(
         kv_cache_spec, ["placeholder"], vllm_config, device
     )
diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py
index 18d1d309c1..3afa8eca92 100644
--- a/tests/v1/attention/test_mla_backends.py
+++ b/tests/v1/attention/test_mla_backends.py
@@ -1031,6 +1031,11 @@ def run_attention_backend(
 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])
 @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])
 @pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST)
+@pytest.mark.parametrize(
+    ("qk_nope_head_dim", "v_head_dim"),
+    [(128, 128), (192, 256)],
+    ids=["deepseek", "glm"],
+)
 def test_backend_correctness(
     default_vllm_config,
     dist_init,
@@ -1042,6 +1047,8 @@ def test_backend_correctness(
     q_scale: float,
     k_scale: float,
     prefill_backend: MLAPrefillBackendEnum,
+    qk_nope_head_dim: int,
+    v_head_dim: int,
 ):
     """
     Test that all backends produce similar outputs to a reference implementation
@@ -1090,9 +1097,9 @@ def test_backend_correctness(
             MLAPrefillSelectorConfig(
                 dtype=torch.bfloat16,
                 mla_dimensions=MLADimensions(
-                    qk_nope_head_dim=128,
+                    qk_nope_head_dim=qk_nope_head_dim,
                     qk_rope_head_dim=64,
-                    v_head_dim=128,
+                    v_head_dim=v_head_dim,
                 ),
             ),
         )
@@ -1168,8 +1175,6 @@ def test_backend_correctness(
     dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)
     kv_lora_rank = 512
     qk_rope_head_dim = 64
-    qk_nope_head_dim = 128
-    v_head_dim = 128
     total_head_size = kv_lora_rank + qk_rope_head_dim
     assert kv_lora_rank + qk_rope_head_dim == head_size, (
         f"MLA dimensions don't match: {total_head_size} != {head_size}"
diff --git a/tests/v1/attention/test_sparse_mla_backends.py b/tests/v1/attention/test_sparse_mla_backends.py
index 6e389604c8..d88f107ece 100644
--- a/tests/v1/attention/test_sparse_mla_backends.py
+++ b/tests/v1/attention/test_sparse_mla_backends.py
@@ -442,6 +442,23 @@ def test_sparse_backend_decode_correctness(
         scale=kv_cache_scale,
     )
 
+    # The sparse builder clones the layer's dense-MHA prefill backend from
+    # static_forward_context; register a mock layer carrying one.
+    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+
+    prefill_backend = get_mla_prefill_backend(vllm_config)(
+        num_heads=num_heads,
+        scale=scale,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        vllm_config=vllm_config,
+    )
+    vllm_config.compilation_config.static_forward_context["placeholder"] = (
+        SimpleNamespace(prefill_backend=prefill_backend)
+    )
+
     builder_cls = backend_cls.get_builder_cls()
     builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
     metadata = builder.build(
@@ -722,6 +739,258 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):
     assert out == expected
 
 
+PREFILL_BATCH_SPECS = {
+    "short_dense_mha": BatchSpec(seq_lens=[64, 128], query_lens=[64, 128]),
+    "short_context_dense_mha": BatchSpec(seq_lens=[128, 160], query_lens=[64, 32]),
+}
+
+
+@pytest.mark.skipif(
+    torch.cuda.get_device_capability()[0] < 10,
+    reason="Sparse MLA forward_mha requires FA4 (SM100+)",
+)
+@pytest.mark.parametrize("batch_name", list(PREFILL_BATCH_SPECS.keys()))
+@pytest.mark.parametrize("kv_cache_dtype", ["auto"])
+def test_sparse_backend_prefill_correctness(
+    default_vllm_config,
+    dist_init,
+    batch_name,
+    kv_cache_dtype,
+    workspace_init,
+):
+    """Test single-pass FA4 dense forward_mha for sparse MLA prefill."""
+    backend_cls = FlashMLASparseBackend
+    batch_spec = PREFILL_BATCH_SPECS[batch_name]
+
+    device = torch.device("cuda")
+    dtype = torch.bfloat16
+    block_size = 64
+
+    num_heads = 128
+    kv_lora_rank = 512
+    qk_nope_head_dim = 128
+    qk_rope_head_dim = 64
+    v_head_dim = 128
+    head_size = kv_lora_rank + qk_rope_head_dim
+    topk_tokens = 512
+
+    max_seqlen = max(batch_spec.seq_lens)
+    total_cache_tokens = sum(batch_spec.seq_lens)
+
+    vllm_config = create_vllm_config(
+        model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",
+        tensor_parallel_size=1,
+        max_model_len=max_seqlen,
+        num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),
+        block_size=block_size,
+        hf_config_override={
+            "index_topk": topk_tokens,
+            "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],
+        },
+    )
+    model_config = vllm_config.model_config
+    model_config.hf_text_config = SimpleNamespace(
+        q_lora_rank=None,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        model_type="deepseek_v2",
+    )
+    model_config.dtype = dtype
+    model_config.get_num_attention_heads = MethodType(
+        lambda self, parallel_config: num_heads, model_config
+    )
+    model_config.get_num_kv_heads = MethodType(
+        lambda self, parallel_config: 1, model_config
+    )
+    model_config.get_head_size = MethodType(lambda self: head_size, model_config)
+    model_config.get_sliding_window = MethodType(lambda self: None, model_config)
+
+    kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
+    scale = 1.0 / math.sqrt(qk_nope_head_dim + qk_rope_head_dim)
+
+    torch.manual_seed(42)
+
+    W_UK = torch.rand(
+        kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device
+    )
+    W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)
+
+    seq_lens = batch_spec.seq_lens
+    query_lens = batch_spec.query_lens
+
+    # Compute dense reference outputs.
+    total_query_tokens = sum(query_lens)
+    sparse_indices = torch.zeros(
+        total_query_tokens, topk_tokens, dtype=torch.int32, device=device
+    )
+
+    all_q, all_kv_c_new, all_k_pe_new = [], [], []
+    kv_c_contexts, k_pe_contexts = [], []
+    reference_outputs = []
+
+    for i in range(batch_spec.batch_size):
+        s_len = seq_lens[i]
+        q_len = query_lens[i]
+        ctx_len = s_len - q_len
+
+        q_mha = torch.rand(
+            q_len,
+            num_heads,
+            qk_nope_head_dim + qk_rope_head_dim,
+            dtype=dtype,
+            device=device,
+        )
+        kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)
+        k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)
+
+        # Decompress all KV for reference
+        kv_b_weight = torch.cat([W_UK, W_UV], dim=-1).view(
+            kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
+        )
+        kv_decompressed = (kv_c_full @ kv_b_weight).view(
+            s_len, num_heads, qk_nope_head_dim + v_head_dim
+        )
+        k_nope_all, v_all = kv_decompressed.split(
+            [qk_nope_head_dim, v_head_dim], dim=-1
+        )
+        k_pe_expanded = k_pe_full.expand(-1, num_heads, -1)
+        k_all = torch.cat([k_nope_all, k_pe_expanded], dim=-1)
+
+        for j in range(q_len):
+            attend_end = ctx_len + j + 1
+            q_tok = q_mha[j : j + 1]  # (1, H, D_qk)
+            k_attend = k_all[:attend_end]  # (N, H, D_qk)
+            v_attend = v_all[:attend_end]  # (N, H, D_v)
+
+            q_sdpa = q_tok.unsqueeze(0).transpose(1, 2).float()
+            k_sdpa = k_attend.unsqueeze(0).transpose(1, 2).float()
+            v_sdpa = v_attend.unsqueeze(0).transpose(1, 2).float()
+
+            out = torch.nn.functional.scaled_dot_product_attention(
+                q_sdpa, k_sdpa, v_sdpa, scale=scale
+            )
+            out = out.transpose(1, 2).squeeze(0)  # (1, H, D_v)
+            reference_outputs.append(out.to(dtype).flatten(start_dim=-2))
+
+        all_q.append(q_mha)
+        all_kv_c_new.append(kv_c_full[ctx_len:])
+        all_k_pe_new.append(k_pe_full[ctx_len:])
+        kv_c_contexts.append(kv_c_full)
+        k_pe_contexts.append(k_pe_full)
+
+    query_cat = torch.cat(all_q, dim=0)
+    kv_c_cat = torch.cat(all_kv_c_new, dim=0)
+    k_pe_cat = torch.cat(all_k_pe_new, dim=0)
+    ref_output = torch.cat(reference_outputs, dim=0)
+
+    vllm_config.cache_config.cache_dtype = kv_cache_dtype
+    vllm_config.model_config.hf_config.index_topk = topk_tokens
+
+    common_attn_metadata = create_common_attn_metadata(
+        batch_spec,
+        vllm_config.cache_config.block_size,
+        device,
+        arange_block_indices=True,
+    )
+
+    kv_cache = create_and_prepopulate_kv_cache(
+        kv_c_contexts=kv_c_contexts,
+        k_pe_contexts=k_pe_contexts,
+        block_size=block_size,
+        head_size=head_size,
+        dtype=dtype,
+        device=device,
+        num_blocks=vllm_config.cache_config.num_gpu_blocks,
+        common_attn_metadata=common_attn_metadata,
+        randomize_blocks=False,
+        kv_cache_dtype=kv_cache_dtype,
+    )
+
+    # The sparse builder clones the layer's dense-MHA prefill backend from
+    # static_forward_context; register a mock layer carrying one.
+    from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend
+
+    prefill_backend = get_mla_prefill_backend(vllm_config)(
+        num_heads=num_heads,
+        scale=scale,
+        kv_lora_rank=kv_lora_rank,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+        v_head_dim=v_head_dim,
+        vllm_config=vllm_config,
+    )
+    vllm_config.compilation_config.static_forward_context["placeholder"] = (
+        SimpleNamespace(prefill_backend=prefill_backend)
+    )
+
+    builder_cls = backend_cls.get_builder_cls()
+    builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)
+    # Drive the queries through the dense-MHA prefill path directly (the routing
+    # threshold would otherwise classify these short queries as MQA decodes).
+    builder.reorder_batch_threshold = 1
+    metadata = builder.build(
+        common_prefix_len=0, common_attn_metadata=common_attn_metadata
+    )
+
+    mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)
+
+    kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1).view(
+        kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)
+    )
+
+    mock_kv_b_proj = ColumnParallelLinear(
+        input_size=kv_lora_rank,
+        output_size=num_heads * (qk_nope_head_dim + v_head_dim),
+        bias=False,
+    ).to(device=device, dtype=dtype)
+    mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())
+
+    impl_cls = backend_cls.get_impl_cls()
+    with set_current_vllm_config(vllm_config):
+        impl = impl_cls(
+            num_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+            num_kv_heads=1,
+            alibi_slopes=None,
+            sliding_window=None,
+            kv_cache_dtype=kv_cache_dtype,
+            logits_soft_cap=None,
+            attn_type="decoder",
+            kv_sharing_target_layer_name=None,
+            q_lora_rank=None,
+            kv_lora_rank=kv_lora_rank,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            kv_b_proj=mock_kv_b_proj,
+            indexer=mock_indexer,
+        )
+        impl.process_weights_after_loading(dtype)
+
+    out_buffer = torch.empty(
+        total_query_tokens, num_heads * v_head_dim, dtype=dtype, device=device
+    )
+
+    with torch.inference_mode():
+        impl.forward_mha(
+            q=query_cat,
+            kv_c_normed=kv_c_cat,
+            k_pe=k_pe_cat,
+            kv_c_and_k_pe_cache=kv_cache,
+            attn_metadata=metadata,
+            k_scale=torch.tensor(1.0, device=device),
+            output=out_buffer,
+        )
+
+    assert out_buffer.shape == ref_output.shape
+    assert torch.isfinite(out_buffer).all(), "Non-finite values in output"
+    torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01)
+
+
 @pytest.mark.parametrize(
     "seq_lens,query_lens,workspace_size,max_logits_bytes,expected",
     [
diff --git a/vllm/config/attention.py b/vllm/config/attention.py
index 48db183d5a..e2ad4729a1 100644
--- a/vllm/config/attention.py
+++ b/vllm/config/attention.py
@@ -59,6 +59,11 @@ class AttentionConfig:
     use_non_causal: bool = False
     """Whether to use non-causal (bidirectional) attention."""
 
+    sparse_mla_force_mqa: bool = False
+    """Force sparse MLA to use forward_mqa for all requests, including prefill.
+    When False (default), pure prefill batches use forward_mha when implemented.
+    Set to True to always use the MQA path."""
+
     flex_attn_block_m: int | None = None
     """Triton kernel BLOCK_M tile size for flex attention.
     Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1,
diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index b8f197b316..4f729f6a06 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -259,7 +259,6 @@ from vllm.v1.attention.backend import (
     AttentionType,
     CommonAttentionMetadata,
     MLAAttentionImpl,
-    SparseMLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.prefill import (
     MLAPrefillBackend,
@@ -490,16 +489,30 @@ class MLAAttention(nn.Module, AttentionLayerBase):
             raise ValueError(f"Duplicate layer name: {prefix}")
         compilation_config.static_forward_context[prefix] = self
 
-        prefill_backend_cls = get_mla_prefill_backend(vllm_config)
-        self.prefill_backend = prefill_backend_cls(
-            num_heads=self.num_heads,
-            scale=self.scale,
-            kv_lora_rank=self.kv_lora_rank,
-            qk_nope_head_dim=self.qk_nope_head_dim,
-            qk_rope_head_dim=self.qk_rope_head_dim,
-            v_head_dim=self.v_head_dim,
-            vllm_config=vllm_config,
-        )
+        self.prefill_backend: MLAPrefillBackend | None
+        try:
+            prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+        except ValueError:
+            if (
+                not self.impl.is_sparse
+                or vllm_config.attention_config.mla_prefill_backend is not None
+            ):
+                raise
+            logger.warning_once(
+                "No MLA prefill backend supports this model; sparse MLA will use the "
+                "top-k MQA path only (no dense-MHA prefill)."
+            )
+            self.prefill_backend = None
+        else:
+            self.prefill_backend = prefill_backend_cls(
+                num_heads=self.num_heads,
+                scale=self.scale,
+                kv_lora_rank=self.kv_lora_rank,
+                qk_nope_head_dim=self.qk_nope_head_dim,
+                qk_rope_head_dim=self.qk_rope_head_dim,
+                v_head_dim=self.v_head_dim,
+                vllm_config=vllm_config,
+            )
 
         self.kv_cache = torch.tensor([])
 
@@ -693,23 +706,28 @@ class MLAAttention(nn.Module, AttentionLayerBase):
         if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla":
             kv_cache = kv_cache.view(current_platform.fp8_dtype())
 
-        # Sparse MLA impls only support forward_mqa (decode-style attention)
-        is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl)
+        assert (
+            attn_metadata.num_decodes is not None
+            and attn_metadata.num_prefills is not None
+            and attn_metadata.num_decode_tokens is not None
+        )
+        num_mqa_tokens = attn_metadata.num_decode_tokens
+        num_mha_tokens = q.size(0) - num_mqa_tokens
 
-        if is_sparse_impl:
-            num_mqa_tokens = q.size(0)
-            num_mha_tokens = 0
-        else:
-            assert (
-                attn_metadata.num_decodes is not None
-                and attn_metadata.num_prefills is not None
-                and attn_metadata.num_decode_tokens is not None
+        if self.impl.is_sparse and num_mha_tokens > 0:
+            prefill_max_seq_len = attn_metadata.prefill_max_seq_len  # type: ignore[attr-defined]
+            use_mha = (
+                self.prefill_backend is not None
+                and prefill_max_seq_len <= attn_metadata.topk_tokens  # type: ignore[attr-defined]
+                and not self._vllm_config.attention_config.sparse_mla_force_mqa
             )
-            num_mqa_tokens = attn_metadata.num_decode_tokens
-            num_mha_tokens = q.size(0) - num_mqa_tokens
+            if not use_mha:
+                num_mqa_tokens = q.size(0)
+                num_mha_tokens = 0
 
         mha_use_quant_output = (
             quant_key is not None
+            and self.prefill_backend is not None
             and self.prefill_backend.supports_quant_output(quant_key)
             and attn_metadata is not None
             and attn_metadata.prefill is not None
@@ -807,7 +825,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
                 mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)
 
             # call decode attn
-            if not is_sparse_impl:
+            if not self.impl.is_sparse:
                 assert attn_metadata.decode is not None
             attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)  # type: ignore[attr-defined]
 
@@ -1398,6 +1416,182 @@ def backend_supports_prefill_query_quantization() -> bool:
     )
 
 
+def build_mla_chunked_context_metadata(
+    *,
+    context_lens_cpu: torch.Tensor,
+    prefill_query_start_loc_cpu: torch.Tensor,
+    num_prefills: int,
+    chunked_prefill_workspace: torch.Tensor,
+    chunked_prefill_workspace_size: int,
+    block_size: int,
+    align_chunk_to_block: bool,
+    device: torch.device,
+    dcp_world_size: int,
+    dcp_local_block_size: int,
+    dcp_virtual_block_size: int,
+) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":
+    """Build chunked-context metadata for an MLA prefill.
+
+    Shared by dense and sparse builders. Splits each prefill's context
+    into workspace-sized chunks and, under DCP, plans the per-rank interleaved
+    local chunks the all-gather reduction consumes.
+
+    Args:
+        context_lens_cpu: Per-prefill context length (seq_len - query_len).
+        prefill_query_start_loc_cpu: Prefill query cumulative offsets (0-based).
+        num_prefills: Number of prefill requests.
+        chunked_prefill_workspace: Scratch buffer the context gather writes to.
+        chunked_prefill_workspace_size: Row capacity of the workspace.
+        block_size: KV cache page size for chunk-start alignment.
+        align_chunk_to_block: Round the chunk size down to ``block_size``.
+        device: Target device for the returned tensors.
+        dcp_world_size: Decode-context-parallel world size (1 if disabled).
+        dcp_local_block_size: Per-rank interleave block size for DCP.
+        dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``.
+
+    Returns:
+        The chunked-context metadata, or None when no prefill has any context.
+    """
+    # NOTE: it is recommended you read the `Chunked Prefill` section in the
+    # comment at the top of the file before trying to understand this code.
+    max_context_len = context_lens_cpu.max().item()
+    if max_context_len <= 0:
+        return None
+    num_prefills_with_context = int((context_lens_cpu > 0).sum().item())
+
+    # Currently we allocate an equal amount of workspace for each prefill with
+    # context; we could probably use a more advanced algorithm here and allocate
+    # more workspace to prefills with longer context lengths.
+    max_context_chunk = chunked_prefill_workspace_size // num_prefills_with_context
+    if align_chunk_to_block:
+        # The `gather_and_maybe_dequant_cache` kernel cannot handle chunk
+        # starts that are not aligned to block_size, so round down.
+        max_context_chunk = round_down(max_context_chunk, block_size)
+    assert max_context_chunk > 0
+
+    num_chunks = cdiv(max_context_len, max_context_chunk)
+    # e.g. max_context_chunk=256, num_chunks=3, num_prefills=4 ->
+    #   [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
+    # Note(simon): this is done on CPU because of downstream's use of `to_list`.
+    chunk_starts = torch.empty(
+        num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
+    ).copy_(
+        torch.arange(num_chunks, dtype=torch.int32)
+        .multiply_(max_context_chunk)
+        .unsqueeze(1)
+    )
+    chunk_ends = torch.min(
+        context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
+    )
+    chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
+
+    cu_seq_lens_cpu = torch.zeros(
+        num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
+    )
+    torch.cumsum(chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32)
+    chunk_total_token = cu_seq_lens_cpu[:, -1]
+
+    max_tokens_over_chunk = chunk_total_token.max().item()
+    token_to_seq_cpu = torch.zeros(
+        (num_chunks, max_tokens_over_chunk), dtype=torch.int32, pin_memory=True
+    )
+    req_indices = torch.arange(num_prefills, dtype=torch.int32)
+    for i in range(num_chunks):
+        token_to_seq = torch.repeat_interleave(req_indices, chunk_seq_lens[i])
+        token_to_seq_cpu[i, : token_to_seq.shape[0]] = token_to_seq
+
+    prefill_tokens_with_context = prefill_query_start_loc_cpu[
+        num_prefills_with_context
+    ].item()
+
+    metadata_cls = MLACommonPrefillMetadata.ChunkedContextMetadata
+    if dcp_world_size > 1:
+        local_context_lens_allranks = get_dcp_local_seq_lens(
+            context_lens_cpu, dcp_world_size, None, dcp_local_block_size
+        )
+        # Note(qcs): Per-rank local context lengths, padded to
+        # `dcp_local_block_size`.
+        padded_local_context_lens_cpu: torch.Tensor = (
+            cdiv(context_lens_cpu, dcp_virtual_block_size) * dcp_local_block_size
+        )
+        # Note(hc): The above max_context_chunk already enforces block_size
+        # alignment; DCP only requires block_size be divisible by dcp_world_size,
+        # because DCP uses cp_gather_cache, which does not require chunk starts
+        # aligned to block_size.
+        assert max_context_chunk % dcp_world_size == 0
+        padded_local_max_context_chunk = (
+            cdiv(max_context_chunk, dcp_virtual_block_size) * dcp_local_block_size
+        )
+        local_chunk_starts = torch.empty(
+            num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
+        ).copy_(
+            torch.arange(num_chunks, dtype=torch.int32)
+            .multiply_(padded_local_max_context_chunk)
+            .unsqueeze(1)
+        )
+        local_chunk_ends = torch.min(
+            padded_local_context_lens_cpu.unsqueeze(0),
+            local_chunk_starts + padded_local_max_context_chunk,
+        )
+        padded_local_chunk_seq_lens = (local_chunk_ends - local_chunk_starts).clamp(
+            min=0
+        )
+
+        padded_local_cu_seq_lens_cpu = torch.zeros(
+            num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
+        )
+        torch.cumsum(
+            padded_local_chunk_seq_lens,
+            dim=1,
+            out=padded_local_cu_seq_lens_cpu[:, 1:],
+            dtype=torch.int32,
+        )
+        max_padded_local_tokens = padded_local_cu_seq_lens_cpu[:, -1].max().item()
+        padded_local_token_to_seq_cpu = torch.zeros(
+            (num_chunks, max_padded_local_tokens), dtype=torch.int32
+        )
+        for i in range(num_chunks):
+            tts = torch.repeat_interleave(req_indices, padded_local_chunk_seq_lens[i])
+            padded_local_token_to_seq_cpu[i, : tts.shape[0]] = tts
+
+        chunked_context_metadata = metadata_cls(
+            cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
+            starts=local_chunk_starts.to(device, non_blocking=True),
+            seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),
+            max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
+            seq_lens=chunk_seq_lens,
+            token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
+            chunk_total_token=chunk_total_token.tolist(),
+            workspace=chunked_prefill_workspace,
+            prefill_tokens_with_context=prefill_tokens_with_context,
+            padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
+            local_context_lens_allranks=local_context_lens_allranks.tolist(),
+            padded_local_cu_seq_lens=padded_local_cu_seq_lens_cpu.to(
+                device, non_blocking=True
+            ),
+            padded_local_token_to_seq=padded_local_token_to_seq_cpu.to(
+                device, non_blocking=True
+            ),
+            cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
+            chunk_size=padded_local_max_context_chunk,
+        )
+    else:
+        chunked_context_metadata = metadata_cls(
+            cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
+            starts=chunk_starts.to(device, non_blocking=True),
+            seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
+            max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
+            seq_lens=chunk_seq_lens,
+            token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
+            chunk_total_token=chunk_total_token,
+            workspace=chunked_prefill_workspace,
+            prefill_tokens_with_context=prefill_tokens_with_context,
+        )
+
+    assert max(chunked_context_metadata.max_seq_lens) <= chunked_prefill_workspace_size
+    return chunked_context_metadata
+
+
 class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
     """
     NOTE: Please read the comment at the top of the file before trying to
@@ -1662,8 +1856,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
             context_lens_cpu = (
                 seq_lens_cpu[reqs_start:num_reqs] - prefill_query_lens_cpu
             )
-            max_context_len_cpu = context_lens_cpu.max().item()
-            num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()
             prefill_query_start_loc = (
                 query_start_loc[reqs_start:] - query_start_loc[reqs_start]
             )
@@ -1671,188 +1863,19 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
                 query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start]
             )
 
-            chunked_context_metadata = None
-            if max_context_len_cpu > 0:
-                # NOTE: it is recommend you read the `Chunked Prefill` section
-                # in the comment at the top of the file before trying to
-                # understand the following code
-
-                # currently we allocate an equal amount of workspace for each
-                # prefill in the batch, we could probably use a more advanced
-                # algorithm here and allocate more workspace to prefills with
-                # longer context lengths
-                max_context_chunk = (
-                    self.chunked_prefill_workspace_size // num_prefills_with_context_cpu
-                )
-
-                # align max_context_chunk to page_size by rounding down,
-                # currently the `gather_and_maybe_dequant_cache` kernel
-                # cannot handle `context_chunk_starts` that are not aligned
-                # to page_size
-                max_context_chunk = round_down(max_context_chunk, self.page_size)
-
-                assert max_context_chunk > 0
-                num_chunks = cdiv(max_context_len_cpu, max_context_chunk)
-
-                # if `max_context_chunk = 256`, `num_chunks = 3`, and
-                #   `num_prefills_with_context = 4`, create a tensor that looks
-                # like
-                #  [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]
-                # Note(simon): this is done in CPU because of downstream's
-                # of `to_list`.
-                chunk_starts = torch.empty(
-                    num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
-                ).copy_(
-                    torch.arange(num_chunks, dtype=torch.int32)
-                    .multiply_(max_context_chunk)
-                    .unsqueeze(1)
-                )
-                chunk_ends = torch.min(
-                    context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
-                )
-                chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)
-
-                cu_seq_lens_cpu = torch.zeros(
-                    num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
-                )
-                torch.cumsum(
-                    chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32
-                )
-                chunk_total_token = cu_seq_lens_cpu[:, -1]
-
-                max_token_num_over_chunk = chunk_total_token.max().item()
-                token_to_seq_tensor_cpu = torch.zeros(
-                    [num_chunks, max_token_num_over_chunk],
-                    dtype=torch.int32,
-                    pin_memory=True,
-                )
-                range_idx = torch.arange(num_prefills, dtype=torch.int32)
-                for i in range(num_chunks):
-                    chunk_token_to_seq_tensor = torch.repeat_interleave(
-                        range_idx, chunk_seq_lens[i]
-                    )
-                    chunk_len = chunk_token_to_seq_tensor.shape[0]
-                    token_to_seq_tensor_cpu[i, :chunk_len] = chunk_token_to_seq_tensor
-
-                if self.dcp_world_size > 1:
-                    local_context_lens_allranks = get_dcp_local_seq_lens(
-                        context_lens_cpu,
-                        self.dcp_world_size,
-                        None,
-                        self.dcp_local_block_size,
-                    )
-                    # Note(qcs): The max local context lengths
-                    # padded to `dcp_local_block_size`.
-                    padded_local_context_lens_cpu: torch.Tensor = (
-                        cdiv(
-                            context_lens_cpu,
-                            self.dcp_virtual_block_size,
-                        )
-                        * self.dcp_local_block_size
-                    )
-                    # Note(hc): The above max_context_chunk already enforces
-                    # block_size alignment, DCP just need the block_size can
-                    # be divisible by dcp_world_size, because DCP use
-                    # cp_gather_cache which not require `cp_chunk_starts`
-                    # aligned to page_size.
-                    assert max_context_chunk % self.dcp_world_size == 0
-                    padded_local_max_context_chunk_across_ranks = (
-                        cdiv(
-                            max_context_chunk,
-                            self.dcp_virtual_block_size,
-                        )
-                        * self.dcp_local_block_size
-                    )
-                    local_chunk_starts = torch.empty(
-                        num_chunks, num_prefills, dtype=torch.int32, pin_memory=True
-                    ).copy_(
-                        torch.arange(num_chunks, dtype=torch.int32)
-                        .multiply_(padded_local_max_context_chunk_across_ranks)
-                        .unsqueeze(1)
-                    )
-                    local_chunk_ends = torch.min(
-                        padded_local_context_lens_cpu.unsqueeze(0),
-                        local_chunk_starts
-                        + padded_local_max_context_chunk_across_ranks,
-                    )
-                    padded_local_chunk_seq_lens = (
-                        local_chunk_ends - local_chunk_starts
-                    ).clamp(min=0)
-
-                    padded_local_cu_chunk_seq_lens_cpu = torch.zeros(
-                        num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True
-                    )
-                    torch.cumsum(
-                        padded_local_chunk_seq_lens,
-                        dim=1,
-                        out=padded_local_cu_chunk_seq_lens_cpu[:, 1:],
-                        dtype=torch.int32,
-                    )
-                    max_padded_local_tokens_over_chunk = (
-                        padded_local_cu_chunk_seq_lens_cpu[:, -1].max().item()
-                    )
-                    padded_local_token_to_seq_tensor_cpu = torch.zeros(
-                        [num_chunks, max_padded_local_tokens_over_chunk],
-                        dtype=torch.int32,
-                    )
-                    for i in range(num_chunks):
-                        chunk_token_to_seq_tensor = torch.repeat_interleave(
-                            range_idx, padded_local_chunk_seq_lens[i]
-                        )
-                        chunk_len = chunk_token_to_seq_tensor.shape[0]
-                        padded_local_token_to_seq_tensor_cpu[i, :chunk_len] = (
-                            chunk_token_to_seq_tensor
-                        )
-
-                prefill_tokens_with_context = None
-                if num_prefills_with_context_cpu > 0:
-                    prefill_tokens_with_context = prefill_query_start_loc_cpu[
-                        num_prefills_with_context_cpu
-                    ].item()
-                _ChunkedMetadata = MLACommonPrefillMetadata.ChunkedContextMetadata
-                if self.dcp_world_size > 1:
-                    chunked_context_metadata = _ChunkedMetadata(
-                        cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
-                        starts=local_chunk_starts.to(device, non_blocking=True),
-                        seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),
-                        max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
-                        seq_lens=chunk_seq_lens,
-                        token_to_seq=token_to_seq_tensor_cpu.to(
-                            device, non_blocking=True
-                        ),
-                        chunk_total_token=chunk_total_token.tolist(),
-                        workspace=self.chunked_prefill_workspace,
-                        padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),
-                        local_context_lens_allranks=local_context_lens_allranks.tolist(),
-                        padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to(
-                            device, non_blocking=True
-                        ),
-                        padded_local_token_to_seq=padded_local_token_to_seq_tensor_cpu.to(
-                            device, non_blocking=True
-                        ),
-                        cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),
-                        chunk_size=padded_local_max_context_chunk_across_ranks,
-                        prefill_tokens_with_context=prefill_tokens_with_context,
-                    )
-                else:
-                    chunked_context_metadata = _ChunkedMetadata(
-                        cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),
-                        starts=chunk_starts.to(device, non_blocking=True),
-                        seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
-                        max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
-                        seq_lens=chunk_seq_lens,
-                        token_to_seq=token_to_seq_tensor_cpu.to(
-                            device, non_blocking=True
-                        ),
-                        chunk_total_token=chunk_total_token,
-                        workspace=self.chunked_prefill_workspace,
-                        prefill_tokens_with_context=prefill_tokens_with_context,
-                    )
-
-                assert (
-                    max(chunked_context_metadata.max_seq_lens)
-                    <= self.chunked_prefill_workspace_size
-                )
+            chunked_context_metadata = build_mla_chunked_context_metadata(
+                context_lens_cpu=context_lens_cpu,
+                prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
+                num_prefills=num_prefills,
+                chunked_prefill_workspace=self.chunked_prefill_workspace,
+                chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,
+                block_size=self.page_size,
+                align_chunk_to_block=True,
+                device=device,
+                dcp_world_size=self.dcp_world_size,
+                dcp_local_block_size=self.dcp_local_block_size,
+                dcp_virtual_block_size=self.dcp_virtual_block_size,
+            )
 
             prefill_metadata = MLACommonPrefillMetadata(
                 block_table=block_table_tensor[reqs_start:, ...],
@@ -1985,19 +2008,14 @@ def reorg_kvcache(
     return reorganized_kv_c_normed, reorganized_k_pe
 
 
-class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
+class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
     """
-    NOTE: Please read the comment at the top of the file before trying to
-    understand this class
+    Shared MLA base providing dense-MHA prefill (via the selected
+    MLAPrefillBackend) for both dense and sparse impls; subclasses add decode
+    (``forward_mqa``).
     """
 
-    def fused_output_quant_supported(self, quant_key):
-        return quant_key in (
-            kFp8StaticTensorSym,
-            kNvfp4Dynamic,
-            kFp8Dynamic128Sym,
-            kFp8Dynamic64Sym,
-        )
+    _use_flashinfer_concat_mla_k: bool
 
     def __init__(
         self,
@@ -2005,60 +2023,25 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
         head_size: int,
         scale: float,
         num_kv_heads: int,
-        alibi_slopes: list[float] | None,
-        sliding_window: int | None,
         kv_cache_dtype: str,
-        logits_soft_cap: float | None,
-        attn_type: str,
-        kv_sharing_target_layer_name: str | None,
-        # MLA Specific Arguments
-        q_lora_rank: int | None,
         kv_lora_rank: int,
         qk_nope_head_dim: int,
         qk_rope_head_dim: int,
         qk_head_dim: int,
         v_head_dim: int,
         kv_b_proj: ColumnParallelLinear,
-        # DSV3.2 MLA Specific Arguments
-        indexer: object | None = None,
-        q_pad_num_heads: int | None = None,
     ) -> None:
-        if kv_sharing_target_layer_name is not None:
-            raise NotImplementedError("KV sharing is not supported for MLA")
-
         self.num_heads = num_heads
         self.head_size = head_size
         self.scale = float(scale)
         self.num_kv_heads = num_kv_heads
         self.kv_cache_dtype = kv_cache_dtype
-
-        self.q_lora_rank = q_lora_rank
         self.kv_lora_rank = kv_lora_rank
         self.qk_nope_head_dim = qk_nope_head_dim
         self.qk_rope_head_dim = qk_rope_head_dim
         self.qk_head_dim = qk_head_dim
         self.v_head_dim = v_head_dim
         self.kv_b_proj = kv_b_proj
-        self.indexer = indexer
-        self.q_pad_num_heads = q_pad_num_heads
-        self.supports_quant_query_input = True
-        self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
-
-        # Use flashinfer's optimized concat_mla_k kernel when available.
-        # The kernel is optimized for DeepSeek V3 dimensions:
-        # num_heads=128, nope_dim=128, rope_dim=64
-        self._use_flashinfer_concat_mla_k = (
-            has_flashinfer()
-            and (self.num_heads == 128)
-            and (self.qk_nope_head_dim == 128)
-            and (self.qk_rope_head_dim == 64)
-        )
-
-        self.dcp_world_size: int = -1
-
-        self.cp_kv_cache_interleave_size: int = (
-            get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
-        )
 
     def _concat_k_nope_k_pe(
         self, k_nope: torch.Tensor, k_pe: torch.Tensor
@@ -2341,7 +2324,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
 
         return output, output_lse
 
-    def forward_mha(
+    def forward_mha(  # type: ignore[override]
         self,
         q: torch.Tensor,
         kv_c_normed: torch.Tensor,
@@ -2428,6 +2411,83 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
             output_prefill = output_prefill.flatten(start_dim=-2)
             output.copy_(output_prefill)
 
+
+class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
+    """
+    NOTE: Please read the comment at the top of the file before trying to
+    understand this class
+    """
+
+    def fused_output_quant_supported(self, quant_key):
+        return quant_key in (
+            kFp8StaticTensorSym,
+            kNvfp4Dynamic,
+            kFp8Dynamic128Sym,
+            kFp8Dynamic64Sym,
+        )
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA Specific Arguments
+        q_lora_rank: int | None,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        kv_b_proj: ColumnParallelLinear,
+        # DSV3.2 MLA Specific Arguments
+        indexer: object | None = None,
+        q_pad_num_heads: int | None = None,
+    ) -> None:
+        if kv_sharing_target_layer_name is not None:
+            raise NotImplementedError("KV sharing is not supported for MLA")
+
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            kv_cache_dtype,
+            kv_lora_rank,
+            qk_nope_head_dim,
+            qk_rope_head_dim,
+            qk_head_dim,
+            v_head_dim,
+            kv_b_proj,
+        )
+        self.q_lora_rank = q_lora_rank
+        self.indexer = indexer
+        self.q_pad_num_heads = q_pad_num_heads
+        self.supports_quant_query_input = True
+        self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
+
+        # Use flashinfer's optimized concat_mla_k kernel when available.
+        # The kernel is optimized for DeepSeek V3 dimensions:
+        # num_heads=128, nope_dim=128, rope_dim=64
+        self._use_flashinfer_concat_mla_k = (
+            has_flashinfer()
+            and (self.num_heads == 128)
+            and (self.qk_nope_head_dim == 128)
+            and (self.qk_rope_head_dim == 64)
+        )
+
+        self.dcp_world_size: int = -1
+
+        self.cp_kv_cache_interleave_size: int = (
+            get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size
+        )
+
     @abstractmethod
     def forward_mqa(
         self,
diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py
new file mode 100644
index 0000000000..4948dca676
--- /dev/null
+++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py
@@ -0,0 +1,316 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared forward_mha implementation and metadata builder for sparse MLA backends."""
+
+from shutil import which
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+import numpy as np
+import torch
+
+from vllm.distributed import get_dcp_group
+from vllm.logger import init_logger
+from vllm.model_executor.layers.attention.mla_attention import (
+    MLACommonBaseImpl,
+    MLACommonPrefillMetadata,
+    build_mla_chunked_context_metadata,
+    get_mla_dims,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer
+from vllm.utils.torch_utils import np_to_pinned_tensor
+from vllm.v1.attention.backend import (
+    AttentionMetadata,
+    AttentionMetadataBuilder,
+)
+from vllm.v1.attention.backends.utils import split_decodes_and_prefills
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.model_executor.layers.linear import ColumnParallelLinear
+    from vllm.v1.attention.backend import CommonAttentionMetadata
+    from vllm.v1.kv_cache_interface import AttentionSpec
+
+logger = init_logger(__name__)
+
+T = TypeVar("T", bound=AttentionMetadata)
+
+
+class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
+    metadata_cls: type[T]
+
+    def __init__(
+        self,
+        kv_cache_spec: "AttentionSpec",
+        layer_names: list[str],
+        vllm_config: "VllmConfig",
+        device: torch.device,
+    ) -> None:
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+        self.vllm_config = vllm_config
+        self.device = device
+        self.model_config = vllm_config.model_config
+        self.mla_dims = get_mla_dims(self.model_config)
+        self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk
+        self.req_id_per_token_buffer = torch.empty(
+            (vllm_config.scheduler_config.max_num_batched_tokens,),
+            dtype=torch.int32,
+            device=device,
+        )
+        parallel_config = vllm_config.parallel_config
+        try:
+            self.dcp_world_size = get_dcp_group().world_size
+            self.dcp_rank = get_dcp_group().rank_in_group
+        except AssertionError:
+            # DCP might not be initialized in testing
+            self.dcp_world_size = 1
+            self.dcp_rank = 0
+        self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
+        self.dcp_local_block_size = self.cp_kv_cache_interleave_size
+        self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
+
+        self.chunked_prefill_workspace_size = (
+            self.determine_chunked_prefill_workspace_size(vllm_config)
+        )
+        workspace_head_size = (
+            self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim
+        )
+        workspace_rows = self.chunked_prefill_workspace_size
+        if self.dcp_world_size > 1:
+            # DCP gathers each rank's local KV shard into the workspace, so it
+            # needs an extra 1/DCP rows beyond the TP allocation.
+            assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0
+            workspace_rows += self.chunked_prefill_workspace_size // self.dcp_world_size
+        self.chunked_prefill_workspace = torch.empty(
+            (workspace_rows, workspace_head_size),
+            dtype=self.model_config.dtype,
+            device=device,
+        )
+        layer_prefill_backend = vllm_config.compilation_config.static_forward_context[
+            layer_names[0]
+        ].prefill_backend
+        self._prefill_backend = (
+            layer_prefill_backend.clone() if layer_prefill_backend is not None else None
+        )
+
+    @staticmethod
+    def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int:
+        scheduler_config = vllm_config.scheduler_config
+        cache_config = vllm_config.cache_config
+        model_config = vllm_config.model_config
+        topk_tokens = model_config.hf_config.index_topk
+
+        workspace_size = min(
+            max(
+                8 * model_config.max_model_len,
+                4 * scheduler_config.max_num_seqs * cache_config.block_size,
+            ),
+            64 * 1024,
+            scheduler_config.max_num_seqs * topk_tokens,
+        )
+        return max(
+            workspace_size,
+            scheduler_config.max_num_seqs * cache_config.block_size,
+        )
+
+    def _build_req_id_per_token(
+        self,
+        common_attn_metadata: "CommonAttentionMetadata",
+    ) -> torch.Tensor:
+        num_tokens = common_attn_metadata.num_actual_tokens
+        starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
+        seg_lengths = np.diff(starts)
+        req_id_per_token = np.repeat(
+            np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
+        )
+        self.req_id_per_token_buffer.fill_(0)
+        self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
+            np_to_pinned_tensor(req_id_per_token), non_blocking=True
+        )
+        return self.req_id_per_token_buffer[:num_tokens]
+
+    def _build_chunked_context_fields(
+        self,
+        common_attn_metadata: "CommonAttentionMetadata",
+        num_decodes: int,
+        num_prefills: int,
+        prefill_query_lens_cpu: torch.Tensor | None,
+    ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":
+        if num_prefills == 0 or prefill_query_lens_cpu is None:
+            return None
+
+        seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+        assert seq_lens_cpu is not None
+        context_lens_cpu = (
+            seq_lens_cpu[num_decodes : num_decodes + num_prefills]
+            - prefill_query_lens_cpu
+        )
+        qsl_cpu = common_attn_metadata.query_start_loc_cpu
+        prefill_query_start_loc_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]
+
+        return build_mla_chunked_context_metadata(
+            context_lens_cpu=context_lens_cpu,
+            prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,
+            num_prefills=num_prefills,
+            chunked_prefill_workspace=self.chunked_prefill_workspace,
+            chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,
+            block_size=self.kv_cache_spec.block_size,
+            align_chunk_to_block=current_platform.is_cuda(),
+            device=self.device,
+            dcp_world_size=self.dcp_world_size,
+            dcp_local_block_size=self.dcp_local_block_size,
+            dcp_virtual_block_size=self.dcp_virtual_block_size,
+        )
+
+    def build(
+        self,
+        common_prefix_len: int,
+        common_attn_metadata: "CommonAttentionMetadata",
+        fast_build: bool = False,
+    ) -> T:
+        req_id_per_token = self._build_req_id_per_token(common_attn_metadata)
+
+        num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(
+            common_attn_metadata,
+            decode_threshold=self.reorder_batch_threshold or 1,
+        )
+        (
+            prefill_query_start_loc,
+            prefill_max_query_len,
+            prefill_query_lens_cpu,
+        ) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills)
+
+        prefill_max_seq_len = 0
+        prefill: MLACommonPrefillMetadata | None = None
+        if num_prefills > 0 and self._prefill_backend is not None:
+            seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+            assert seq_lens_cpu is not None
+            prefill_max_seq_len = int(
+                seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item()
+            )
+            prefill = MLACommonPrefillMetadata(
+                block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],
+                query_start_loc=prefill_query_start_loc,
+                max_query_len=prefill_max_query_len,
+                chunked_context=self._build_chunked_context_fields(
+                    common_attn_metadata,
+                    num_decodes,
+                    num_prefills,
+                    prefill_query_lens_cpu,
+                ),
+                q_data_type=self.model_config.dtype,
+                output_dtype=self.model_config.dtype,
+                prefill_backend=self._prefill_backend,
+            )
+            self._prefill_backend.prepare_metadata(prefill)
+
+        return self.metadata_cls(  # type: ignore[call-arg]
+            num_reqs=common_attn_metadata.num_reqs,
+            max_query_len=common_attn_metadata.max_query_len,
+            max_seq_len=common_attn_metadata.max_seq_len,
+            num_actual_tokens=common_attn_metadata.num_actual_tokens,
+            query_start_loc=common_attn_metadata.query_start_loc,
+            slot_mapping=common_attn_metadata.slot_mapping,
+            block_table=common_attn_metadata.block_table_tensor,
+            req_id_per_token=req_id_per_token,
+            seq_lens=common_attn_metadata.seq_lens,
+            block_size=self.kv_cache_spec.block_size,
+            topk_tokens=self.topk_tokens,
+            num_decodes=num_decodes,
+            num_prefills=num_prefills,
+            num_decode_tokens=num_decode_tokens,
+            prefill_max_seq_len=prefill_max_seq_len,
+            prefill=prefill,
+            cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,
+        )
+
+    @staticmethod
+    def _build_prefill_fields(
+        common_attn_metadata: "CommonAttentionMetadata",
+        num_decodes: int,
+        num_prefills: int,
+    ) -> tuple[
+        torch.Tensor | None,  # prefill_query_start_loc
+        int,  # prefill_max_query_len
+        torch.Tensor | None,  # prefill_query_lens_cpu
+    ]:
+        if num_prefills == 0:
+            return None, 0, None
+
+        offset = common_attn_metadata.query_start_loc[num_decodes]
+        prefill_query_start_loc = (
+            common_attn_metadata.query_start_loc[num_decodes:] - offset
+        )
+
+        qsl_cpu = common_attn_metadata.query_start_loc_cpu
+        prefill_qsl_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]
+        prefill_query_lens = prefill_qsl_cpu[1:] - prefill_qsl_cpu[:-1]
+        prefill_max_query_len = int(prefill_query_lens.max().item())
+
+        return (
+            prefill_query_start_loc,
+            prefill_max_query_len,
+            prefill_query_lens,
+        )
+
+
+class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):
+    """Sparse MLA base: shared dense-MHA prefill (from MLACommonBaseImpl) plus the
+    sparse top-k MQA decode path. Subclasses implement forward_mqa."""
+
+    is_sparse = True
+
+    def __init__(
+        self,
+        num_heads: int,
+        head_size: int,
+        scale: float,
+        num_kv_heads: int,
+        alibi_slopes: list[float] | None,
+        sliding_window: int | None,
+        kv_cache_dtype: str,
+        logits_soft_cap: float | None,
+        attn_type: str,
+        kv_sharing_target_layer_name: str | None,
+        # MLA-specific
+        q_lora_rank: int | None,
+        kv_lora_rank: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        qk_head_dim: int,
+        v_head_dim: int,
+        kv_b_proj: "ColumnParallelLinear",
+        indexer: object | None = None,
+        topk_indices_buffer: torch.Tensor | None = None,
+        q_pad_num_heads: int | None = None,
+    ) -> None:
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            kv_cache_dtype,
+            kv_lora_rank,
+            qk_nope_head_dim,
+            qk_rope_head_dim,
+            qk_head_dim,
+            v_head_dim,
+            kv_b_proj,
+        )
+
+        # The indexer carries the shared buffer for normal layers and tests;
+        # the explicitly-passed buffer covers backbone skip layers, whose
+        # indexer is not constructed (see deepseek_v2.py).
+        self.topk_indices_buffer: torch.Tensor | None = (
+            indexer.topk_indices_buffer  # type: ignore[attr-defined]
+            if indexer is not None
+            else topk_indices_buffer
+        )
+
+        self._use_flashinfer_concat_mla_k = (
+            has_flashinfer()
+            and which("ninja") is not None
+            and (self.num_heads == 128)
+            and (self.qk_nope_head_dim == 128)
+            and (self.qk_rope_head_dim == 64)
+        )
diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py
index 9559b05e91..2708e0ab1f 100644
--- a/vllm/v1/attention/backend.py
+++ b/vllm/v1/attention/backend.py
@@ -781,6 +781,10 @@ class AttentionImplBase(ABC, Generic[T]):
     method - subclasses define their own forward interfaces.
     """
 
+    # Whether this impl uses a sparse (top-k) attention path. Used by MLA to
+    # route between the dense-MHA prefill and sparse-MQA paths.
+    is_sparse: ClassVar[bool] = False
+
     # Required attributes that all impls should have
     num_heads: int
     head_size: int
@@ -976,7 +980,6 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
     ) -> None:
         raise NotImplementedError
 
-    @abstractmethod
     def forward_mha(
         self,
         q: torch.Tensor,
@@ -1038,86 +1041,6 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):
         )
 
 
-class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):
-    """Sparse MLA attention implementation with only forward_mqa method.
-
-    Sparse MLA implementations only support decode (MQA-style) attention.
-    They do not support prefill (MHA-style) attention.
-    """
-
-    def fused_output_quant_supported(self, quant_key: "QuantKey"):
-        """
-        Does this attention implementation support fused output quantization.
-        Since MLA quantization is done manually in forward_impl (common code),
-        all MLA backends support it by default.
-        """
-        return quant_key in (
-            kFp8StaticTensorSym,
-            kNvfp4Dynamic,
-            kFp8Dynamic128Sym,
-            kFp8Dynamic64Sym,
-        )
-
-    @abstractmethod
-    def __init__(
-        self,
-        num_heads: int,
-        head_size: int,
-        scale: float,
-        num_kv_heads: int,
-        alibi_slopes: list[float] | None,
-        sliding_window: int | None,
-        kv_cache_dtype: str,
-        logits_soft_cap: float | None,
-        attn_type: str,
-        kv_sharing_target_layer_name: str | None,
-        # MLA Specific Arguments
-        q_lora_rank: int | None,
-        kv_lora_rank: int,
-        qk_nope_head_dim: int,
-        qk_rope_head_dim: int,
-        qk_head_dim: int,
-        v_head_dim: int,
-        kv_b_proj: "ColumnParallelLinear",
-        indexer: object | None = None,
-        q_pad_num_heads: int | None = None,
-    ) -> None:
-        raise NotImplementedError
-
-    @abstractmethod
-    def forward_mqa(
-        self,
-        q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
-        kv_c_and_k_pe_cache: torch.Tensor,
-        attn_metadata: T,
-        layer: AttentionLayer,
-    ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """MQA-style decode forward pass."""
-        raise NotImplementedError
-
-    def do_kv_cache_update(
-        self,
-        kv_c_normed: torch.Tensor,
-        k_pe: torch.Tensor,
-        kv_cache: torch.Tensor,
-        slot_mapping: torch.Tensor,
-        kv_cache_dtype: str,
-        k_scale: torch.Tensor,
-    ) -> None:
-        if kv_cache.numel() == 0:
-            return
-        from vllm import _custom_ops as ops
-
-        ops.concat_and_cache_mla(
-            kv_c_normed,
-            k_pe.squeeze(1),
-            kv_cache,
-            slot_mapping.flatten(),
-            kv_cache_dtype=kv_cache_dtype,
-            scale=k_scale,
-        )
-
-
 def subclass_attention_backend(
     name_prefix: str,
     attention_backend_cls: type[AttentionBackend],
diff --git a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
index 664bd649fd..7c320b1ea3 100644
--- a/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
@@ -4,22 +4,23 @@
 from dataclasses import dataclass
 from typing import Any, ClassVar
 
-import numpy as np
 import torch
 
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
+from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonImpl,
+    SparseMLACommonMetadataBuilder,
+)
 from vllm.platforms.interface import DeviceCapability
-from vllm.utils.torch_utils import np_to_pinned_tensor
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
     AttentionLayer,
     AttentionMetadata,
-    AttentionMetadataBuilder,
-    CommonAttentionMetadata,
+    MLAAttentionImpl,
     MultipleOf,
-    SparseMLAAttentionImpl,
 )
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 from vllm.v1.attention.backends.mla.sparse_utils import (
@@ -50,7 +51,7 @@ class FlashAttnMLASparseBackend(AttentionBackend):
         return FlashAttnMLASparseMetadataBuilder
 
     @staticmethod
-    def get_impl_cls() -> type[SparseMLAAttentionImpl[Any]]:
+    def get_impl_cls() -> type[MLAAttentionImpl[Any]]:
         return FlashAttnMLASparseImpl
 
     @classmethod
@@ -125,13 +126,21 @@ class FlashAttnMLASparseMetadata(AttentionMetadata):
 
     block_table: torch.Tensor
     req_id_per_token: torch.Tensor
+    seq_lens: torch.Tensor
     block_size: int = 64
     topk_tokens: int = 2048
+    num_decodes: int = 0
+    num_prefills: int = 0
+    num_decode_tokens: int = 0
+    prefill_max_seq_len: int = 0
+    prefill: MLACommonPrefillMetadata | None = None
+    cp_kv_cache_interleave_size: int = 1
 
 
 class FlashAttnMLASparseMetadataBuilder(
-    AttentionMetadataBuilder[FlashAttnMLASparseMetadata]
+    SparseMLACommonMetadataBuilder[FlashAttnMLASparseMetadata]
 ):
+    metadata_cls = FlashAttnMLASparseMetadata
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
 
     def __init__(
@@ -141,55 +150,16 @@ class FlashAttnMLASparseMetadataBuilder(
         vllm_config: VllmConfig,
         device: torch.device,
     ) -> None:
-        self.vllm_config = vllm_config
-        self.layer_names = layer_names
-        self.kv_cache_spec = kv_cache_spec
-        self.model_config = vllm_config.model_config
-        self.device = device
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
 
-        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
-
-        self.topk_tokens = vllm_config.model_config.hf_config.index_topk
-        self.req_id_per_token_buffer = torch.empty(
-            (vllm_config.scheduler_config.max_num_batched_tokens,),
-            dtype=torch.int32,
-            device=device,
-        )
-
-    def build(
-        self,
-        common_prefix_len: int,
-        common_attn_metadata: CommonAttentionMetadata,
-        fast_build: bool = False,
-    ) -> FlashAttnMLASparseMetadata:
-        cm = common_attn_metadata
-        num_tokens = cm.num_actual_tokens
-        starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
-        seg_lengths = np.diff(starts)
-        req_id_per_token = np.repeat(
-            np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
-        )
-
-        self.req_id_per_token_buffer.fill_(0)
-        self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            np_to_pinned_tensor(req_id_per_token), non_blocking=True
-        )
-
-        return FlashAttnMLASparseMetadata(
-            num_reqs=cm.num_reqs,
-            max_query_len=cm.max_query_len,
-            max_seq_len=cm.max_seq_len,
-            num_actual_tokens=cm.num_actual_tokens,
-            query_start_loc=cm.query_start_loc,
-            slot_mapping=cm.slot_mapping,
-            block_table=cm.block_table_tensor,
-            req_id_per_token=self.req_id_per_token_buffer[:num_tokens],
-            block_size=self.kv_cache_spec.block_size,
-            topk_tokens=self.topk_tokens,
+        num_q_heads = self.model_config.get_num_attention_heads(
+            vllm_config.parallel_config
         )
+        threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)
+        self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True)
 
 
-class FlashAttnMLASparseImpl(SparseMLAAttentionImpl[FlashAttnMLASparseMetadata]):
+class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
     def __init__(
         self,
         num_heads: int,
@@ -217,22 +187,25 @@ class FlashAttnMLASparseImpl(SparseMLAAttentionImpl[FlashAttnMLASparseMetadata])
                 "FlashAttnMLASparseImpl currently supports only FP16/BF16 KV cache."
             )
 
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.scale = float(scale)
-        self.num_kv_heads = num_kv_heads
-        self.kv_cache_dtype = kv_cache_dtype
-        self.kv_lora_rank: int = mla_args["kv_lora_rank"]
-        self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            indexer=indexer,
+            topk_indices_buffer=topk_indices_buffer,
+            **mla_args,
         )
         assert self.topk_indices_buffer is not None, (
             "Indexer or topk_indices_buffer required for sparse MLA"
         )
         self.supports_quant_query_input = False
-        self.dcp_world_size = -1
-        self.q_pad_num_heads = None
 
     def forward_mqa(
         self,
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
index b66952cc74..a6b9f4abb0 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
@@ -5,37 +5,33 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, ClassVar
 
-import numpy as np
 import torch
 
 from vllm import envs
 from vllm.config import VllmConfig
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
-from vllm.model_executor.layers.attention.mla_attention import (
-    get_mla_dims,
+from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonImpl,
+    SparseMLACommonMetadataBuilder,
 )
 from vllm.platforms.interface import DeviceCapability
-from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
+from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
     AttentionLayer,
     AttentionMetadata,
-    AttentionMetadataBuilder,
     AttentionType,
-    CommonAttentionMetadata,
+    MLAAttentionImpl,
     MultipleOf,
-    SparseMLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
     triton_convert_req_index_to_global_index,
     triton_filter_and_convert_dcp_index,
 )
-from vllm.v1.attention.backends.utils import (
-    KVCacheLayoutType,
-    split_decodes_and_prefills,
-)
+from vllm.v1.attention.backends.utils import KVCacheLayoutType
 from vllm.v1.kv_cache_interface import AttentionSpec
 
 if TYPE_CHECKING:
@@ -85,7 +81,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):
         return [32, 64]
 
     @staticmethod
-    def get_impl_cls() -> type[SparseMLAAttentionImpl]:
+    def get_impl_cls() -> type[MLAAttentionImpl]:
         return FlashInferMLASparseImpl
 
     @classmethod
@@ -162,7 +158,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):
         return [64, 256]
 
     @staticmethod
-    def get_impl_cls() -> type[SparseMLAAttentionImpl]:
+    def get_impl_cls() -> type[MLAAttentionImpl]:
         from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (
             FlashInferMLASparseSM120Impl,
         )
@@ -256,7 +252,10 @@ class FlashInferMLASparseMetadata(AttentionMetadata):
     # Sequence lengths for all requests (context + query)
     seq_lens: torch.Tensor
     num_decodes: int
+    num_prefills: int
     num_decode_tokens: int
+    prefill_max_seq_len: int = 0
+    prefill: MLACommonPrefillMetadata | None = None
 
     # Sparse-specific
     block_size: int = 64
@@ -265,10 +264,11 @@ class FlashInferMLASparseMetadata(AttentionMetadata):
 
 
 class FlashInferMLASparseMetadataBuilder(
-    AttentionMetadataBuilder[FlashInferMLASparseMetadata]
+    SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata]
 ):
     """Builder for FlashInfer MLA Sparse attention metadata."""
 
+    metadata_cls = FlashInferMLASparseMetadata
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
 
     def __init__(
@@ -278,75 +278,20 @@ class FlashInferMLASparseMetadataBuilder(
         vllm_config: VllmConfig,
         device: torch.device,
     ) -> None:
-        self.vllm_config = vllm_config
-        self.layer_names = layer_names
-        self.kv_cache_spec = kv_cache_spec
-        self.model_config = vllm_config.model_config
-        self.device = device
-
-        self.mla_dims = get_mla_dims(self.model_config)
-        self.topk_tokens = vllm_config.model_config.hf_config.index_topk
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
 
+        num_q_heads = vllm_config.model_config.get_num_attention_heads(
+            vllm_config.parallel_config
+        )
+        threshold = {8: 128, 16: 128, 32: 128, 64: 256, 128: 1024}.get(
+            num_q_heads, 1024
+        )
         self._init_reorder_batch_threshold(
-            1,
+            threshold,
             supports_spec_as_decode=True,
             supports_dcp_with_varlen=True,
         )
 
-        self.req_id_per_token_buffer = torch.empty(
-            (vllm_config.scheduler_config.max_num_batched_tokens,),
-            dtype=torch.int32,
-            device=device,
-        )
-
-    def build(
-        self,
-        common_prefix_len: int,
-        common_attn_metadata: CommonAttentionMetadata,
-        fast_build: bool = False,
-    ) -> FlashInferMLASparseMetadata:
-        cm = common_attn_metadata
-        num_tokens = cm.num_actual_tokens
-        assert self.reorder_batch_threshold is not None
-        num_decodes, _, num_decode_tokens, _ = split_decodes_and_prefills(
-            cm,
-            decode_threshold=self.reorder_batch_threshold,
-            treat_short_extends_as_decodes=True,
-        )
-
-        # Build req_id_per_token mapping
-        starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
-        seg_lengths = np.diff(starts)
-        req_id_per_token = np.repeat(
-            np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
-        )
-
-        # Zero-fill for cudagraphs
-        self.req_id_per_token_buffer.fill_(0)
-        self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            np_to_pinned_tensor(req_id_per_token), non_blocking=True
-        )
-        req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens]
-
-        return FlashInferMLASparseMetadata(
-            num_reqs=cm.num_reqs,
-            max_query_len=cm.max_query_len,
-            max_seq_len=cm.max_seq_len,
-            num_actual_tokens=cm.num_actual_tokens,
-            query_start_loc=cm.query_start_loc,
-            slot_mapping=cm.slot_mapping,
-            block_table=cm.block_table_tensor,
-            req_id_per_token=req_id_per_token_tensor,
-            seq_lens=cm.seq_lens,
-            num_decodes=num_decodes,
-            num_decode_tokens=num_decode_tokens,
-            block_size=self.kv_cache_spec.block_size,
-            topk_tokens=self.topk_tokens,
-            cp_kv_cache_interleave_size=(
-                self.vllm_config.parallel_config.cp_kv_cache_interleave_size
-            ),
-        )
-
 
 # Global workspace buffer (lazily initialized)
 _fi_sparse_workspace: torch.Tensor | None = None
@@ -365,7 +310,7 @@ def _get_workspace_buffer(device: torch.device) -> torch.Tensor:
     return _fi_sparse_workspace
 
 
-class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
+class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):
     """FlashInfer MLA Sparse implementation.
 
     Uses the TRT-LLM MLA kernel with sparse_mla_top_k parameter for
@@ -407,22 +352,20 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata
                 "FlashInferMLASparseImpl"
             )
 
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.scale = float(scale)
-        self.num_kv_heads = num_kv_heads
-        self.kv_cache_dtype = kv_cache_dtype
-
-        # MLA-specific dimensions
-        self.kv_lora_rank: int = mla_args["kv_lora_rank"]
-        self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]
-        self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]
-
-        # The indexer carries the shared buffer for normal layers and tests;
-        # the explicitly-passed buffer covers backbone skip layers, whose
-        # indexer is not constructed (see deepseek_v2.py).
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            indexer=indexer,
+            topk_indices_buffer=topk_indices_buffer,
+            **mla_args,
         )
 
         self._workspace_buffer: torch.Tensor | None = None
diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
index d802f56883..9db696a39c 100644
--- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
+++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py
@@ -9,7 +9,7 @@ import torch
 from vllm.v1.attention.backend import (
     AttentionLayer,
     AttentionType,
-    SparseMLAAttentionImpl,
+    MLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (
     FlashInferMLASparseMetadata,
@@ -29,9 +29,11 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:
     return "pow2_fp32"
 
 
-class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):
+class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]):
     """SM120 FlashInfer sparse-MLA implementation."""
 
+    is_sparse = True
+
     def __init__(
         self,
         num_heads: int,
diff --git a/vllm/v1/attention/backends/mla/flashmla_sparse.py b/vllm/v1/attention/backends/mla/flashmla_sparse.py
index 19381efd73..4dea36a3de 100644
--- a/vllm/v1/attention/backends/mla/flashmla_sparse.py
+++ b/vllm/v1/attention/backends/mla/flashmla_sparse.py
@@ -1,31 +1,30 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, ClassVar
+from typing import TYPE_CHECKING, ClassVar
 
-import numpy as np
 import torch
 
 from vllm import _custom_ops as ops
 from vllm.config import VllmConfig, get_current_vllm_config
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
-from vllm.model_executor.layers.attention.mla_attention import (
-    get_mla_dims,
+from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
+    SparseMLACommonImpl,
+    SparseMLACommonMetadataBuilder,
 )
 from vllm.platforms import current_platform
 from vllm.platforms.interface import DeviceCapability
 from vllm.utils.platform_utils import num_compute_units
-from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
+from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backend import (
     AttentionBackend,
     AttentionCGSupport,
     AttentionLayer,
     AttentionMetadata,
-    AttentionMetadataBuilder,
     CommonAttentionMetadata,
     MultipleOf,
-    SparseMLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.sparse_utils import (
     triton_convert_req_index_to_global_index,
@@ -109,7 +108,7 @@ class FlashMLASparseBackend(AttentionBackend):
         return FlashMLASparseMetadataBuilder
 
     @staticmethod
-    def get_impl_cls() -> type[SparseMLAAttentionImpl[Any]]:
+    def get_impl_cls() -> type["FlashMLASparseImpl"]:
         return FlashMLASparseImpl
 
     @classmethod
@@ -159,6 +158,14 @@ class FlashMLASparseMetadata(AttentionMetadata):
     block_size: int = 64
     topk_tokens: int = 2048
 
+    num_decodes: int = 0
+    num_prefills: int = 0
+    num_decode_tokens: int = 0
+    seq_lens: torch.Tensor | None = None
+    prefill_max_seq_len: int = 0
+    prefill: MLACommonPrefillMetadata | None = None
+    cp_kv_cache_interleave_size: int = 1
+
     @dataclass
     class FP8KernelMetadata:
         scheduler_metadata: FlashMLASchedMeta
@@ -228,8 +235,11 @@ def get_prefill_workspace_size(max_model_len: int):
     return max_model_len * 5
 
 
-class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetadata]):
+class FlashMLASparseMetadataBuilder(
+    SparseMLACommonMetadataBuilder[FlashMLASparseMetadata]
+):
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
+    metadata_cls = FlashMLASparseMetadata
 
     def __init__(
         self,
@@ -238,29 +248,30 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
         vllm_config: VllmConfig,
         device: torch.device,
     ) -> None:
-        self.vllm_config = vllm_config
-        self.layer_names = layer_names
+        super().__init__(kv_cache_spec, layer_names, vllm_config, device)
         cache_config = vllm_config.cache_config
-        self.kv_cache_spec = kv_cache_spec
-        self.model_config = vllm_config.model_config
         parallel_config = vllm_config.parallel_config
-        self.device = device
 
-        # Classify single-token queries (plus num_speculative_tokens via
-        # supports_spec_as_decode=True) as decodes; longer queries go to
-        # prefill.
-        self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
+        num_q_heads = self.model_config.get_num_attention_heads(parallel_config)
+        if current_platform.is_device_capability_family(100):
+            threshold = {8: 128, 16: 128, 32: 128, 64: 256, 128: 1024}.get(
+                num_q_heads, 1024
+            )
+        else:
+            threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)
+        self._init_reorder_batch_threshold(
+            threshold,
+            supports_spec_as_decode=True,
+        )
 
         sm_count = num_compute_units(device.index)
 
         self.num_heads = self.model_config.get_num_attention_heads(parallel_config)
-        self.mla_dims = get_mla_dims(self.model_config)
         # FP8 decode kernel only supports h_q = 64 or 128, so we need to pad
         self.fp8_decode_padded_heads = (
             FlashMLASparseImpl._compute_fp8_decode_padded_heads(self.num_heads)
         )
 
-        self.topk_tokens = vllm_config.model_config.hf_config.index_topk
         self.use_fp8_kv_cache = cache_config.cache_dtype == "fp8_ds_mla"
         max_num_seqs = vllm_config.scheduler_config.max_num_seqs
         # Shape: [max_num_seqs], all elements = topk_tokens (constant for full-CG)
@@ -306,11 +317,6 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
             dtype=torch.int32,
             device=device,
         )
-        self.req_id_per_token_buffer = torch.empty(
-            (vllm_config.scheduler_config.max_num_batched_tokens,),
-            dtype=torch.int32,
-            device=device,
-        )
 
     def _build_fp8_mixed_decode_prefill(
         self,
@@ -493,51 +499,24 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad
         common_attn_metadata: CommonAttentionMetadata,
         fast_build: bool = False,
     ) -> FlashMLASparseMetadata:
-        cm = common_attn_metadata
-        num_tokens = cm.num_actual_tokens
-        starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)
-        seg_lengths = np.diff(starts)
-        req_id_per_token = np.repeat(
-            np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
-        )
-        # Zero-fill for cudagraphs
-        self.req_id_per_token_buffer.fill_(0)
-        self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
-            np_to_pinned_tensor(req_id_per_token), non_blocking=True
-        )
-        req_id_per_token = self.req_id_per_token_buffer[:num_tokens]
+        metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)
 
-        fp8_extra_metadata: (
-            FlashMLASparseMetadata.FP8SeparatePrefillDecode
-            | FlashMLASparseMetadata.FP8KernelMetadata
-            | None
-        ) = None
         fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL
+        metadata.fp8_use_mixed_batch = fp8_use_mixed_batch
         if self.use_fp8_kv_cache:
             if fp8_use_mixed_batch:
-                fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(cm)
+                metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(
+                    common_attn_metadata
+                )
             else:
-                fp8_extra_metadata = self._build_fp8_separate_prefill_decode(cm)
-
-        metadata = FlashMLASparseMetadata(
-            num_reqs=cm.num_reqs,
-            max_query_len=cm.max_query_len,
-            max_seq_len=cm.max_seq_len,
-            num_actual_tokens=cm.num_actual_tokens,
-            query_start_loc=cm.query_start_loc,
-            slot_mapping=cm.slot_mapping,
-            block_table=cm.block_table_tensor,
-            req_id_per_token=req_id_per_token,
-            block_size=self.kv_cache_spec.block_size,
-            topk_tokens=self.topk_tokens,
-            fp8_extra_metadata=fp8_extra_metadata,
-            fp8_use_mixed_batch=fp8_use_mixed_batch,
-        )
+                metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode(
+                    common_attn_metadata
+                )
 
         return metadata
 
 
-class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
+class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):
     @staticmethod
     def _compute_fp8_decode_padded_heads(num_heads: int) -> int:
         # FP8 decode kernel only supports h_q = 64 or 128
@@ -561,19 +540,22 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):
         indexer: "Indexer | None" = None,
         **mla_args,
     ) -> None:
-        self.num_heads = num_heads
-        self.head_size = head_size
-        self.scale = float(scale)
-        self.num_kv_heads = num_kv_heads
-        self.kv_cache_dtype = kv_cache_dtype
-        self.kv_lora_rank: int = mla_args["kv_lora_rank"]
-        self.softmax_scale = scale
-        # The indexer carries the shared buffer for normal layers and tests;
-        # the explicitly-passed buffer covers backbone skip layers, whose
-        # indexer is not constructed (see deepseek_v2.py).
-        self.topk_indices_buffer: torch.Tensor | None = (
-            indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer
+        super().__init__(
+            num_heads,
+            head_size,
+            scale,
+            num_kv_heads,
+            alibi_slopes,
+            sliding_window,
+            kv_cache_dtype,
+            logits_soft_cap,
+            attn_type,
+            kv_sharing_target_layer_name,
+            indexer=indexer,
+            topk_indices_buffer=topk_indices_buffer,
+            **mla_args,
         )
+        self.softmax_scale = scale
         # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell
         self.prefill_padding = (
             128 if current_platform.is_device_capability_family(100) else 64
diff --git a/vllm/v1/attention/backends/mla/indexer.py b/vllm/v1/attention/backends/mla/indexer.py
index b38445f9ff..2f578b4677 100644
--- a/vllm/v1/attention/backends/mla/indexer.py
+++ b/vllm/v1/attention/backends/mla/indexer.py
@@ -238,7 +238,9 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):
 
 
 class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
-    reorder_batch_threshold: int = 1
+    # The indexer opts out of the shared reorder-threshold vote (see __init__),
+    # so this is None; its own split uses self.decode_threshold.
+    reorder_batch_threshold: int | None = None
 
     @classmethod
     def get_cudagraph_support(
@@ -285,7 +287,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         )
 
         next_n = self.num_speculative_tokens + 1
-        self.reorder_batch_threshold += self.num_speculative_tokens
+        self.decode_threshold = next_n
+        self.reorder_batch_threshold = None
         # NOTE: SM100 datacenter GPUs support any next_n natively via the
         # multi-atom paged MQA logits kernels (FP8 and FP4 indexer
         # caches). Outside the SM100 family the FP8
@@ -565,7 +568,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
         num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
             split_decodes_and_prefills(
                 common_attn_metadata,
-                decode_threshold=self.reorder_batch_threshold,
+                decode_threshold=self.decode_threshold,
                 require_uniform=not self.use_flattening,
             )
         )
diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
index 70e1f41c81..4e35072acd 100644
--- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py
+++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py
@@ -3,7 +3,7 @@
 """FlashAttention backend for MLA prefill."""
 
 import functools
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, ClassVar
 
 import torch
 
@@ -25,7 +25,10 @@ from vllm.v1.attention.backends.fa_utils import (
     get_flash_attn_version,
     is_flash_attn_varlen_func_available,
 )
-from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend
+from vllm.v1.attention.backends.mla.prefill.base import (
+    MLADimensions,
+    MLAPrefillBackend,
+)
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -40,6 +43,14 @@ else:
 class FlashAttnPrefillBackend(MLAPrefillBackend):
     """FlashAttention backend for MLA prefill."""
 
+    supported_mla_dimensions: ClassVar[list[MLADimensions]] = [
+        MLADimensions(
+            qk_nope_head_dim=128,
+            qk_rope_head_dim=64,
+            v_head_dim=128,
+        ),
+    ]
+
     @staticmethod
     def get_name() -> str:
         return "FLASH_ATTN"
@@ -76,7 +87,9 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
         )
         qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
         self.flash_attn_varlen_func = flash_attn_varlen_func
-        self.vllm_flash_attn_version = get_flash_attn_version(head_size=qk_head_dim)
+        self.vllm_flash_attn_version = get_flash_attn_version(
+            head_size=qk_head_dim, head_size_v=v_head_dim
+        )
         if self.vllm_flash_attn_version is not None:
             self.flash_attn_varlen_func = functools.partial(
                 flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version
diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
index 110fcbdb71..bda63e1ae2 100644
--- a/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
@@ -24,8 +24,8 @@ from vllm.v1.attention.backend import (
     AttentionMetadata,
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
+    MLAAttentionImpl,
     MultipleOf,
-    SparseMLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.rocm_aiter_mla import (
     AiterMLAHelper,
@@ -626,7 +626,9 @@ def reference_mla_sparse_prefill(
     return (result, lse)
 
 
-class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata]):
+class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):
+    is_sparse = True
+
     def __init__(
         self,
         num_heads: int,
diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py
index 9e54b62a00..9faf56f639 100644
--- a/vllm/v1/attention/backends/mla/sparse_swa.py
+++ b/vllm/v1/attention/backends/mla/sparse_swa.py
@@ -292,8 +292,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
     - Chunked prefill (aligns with the indexer's chunking)
     """
 
-    # Base threshold: query_len <= 1 is decode
-    reorder_batch_threshold: int = 1
+    reorder_batch_threshold: int | None = None
     _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
 
     def __init__(self, *args, **kwargs):
@@ -315,14 +314,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):
         )
         # Decode can have query_len up to
         #   1 + (2 if parallel drafting else 1) * num_speculative_tokens.
-        # This MUST match the flashmla_sparse / indexer threshold so that
-        # all backends agree on the decode/prefill split.
+        # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take
+        # the prefill path and the decode/prefill split stays at that width.
         spec_mult = (
             2 if (spec_config is not None and spec_config.parallel_drafting) else 1
         )
-        self.decode_threshold = (
-            self.reorder_batch_threshold + spec_mult * self.num_speculative_tokens
-        )
+        self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens
+        self.reorder_batch_threshold = None
 
         hf_config = self.vllm_config.model_config.hf_config
         assert hasattr(hf_config, "sliding_window")
diff --git a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
index 9aad453210..d24147bfc3 100644
--- a/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+++ b/vllm/v1/attention/backends/mla/xpu_mla_sparse.py
@@ -21,7 +21,7 @@ from vllm.v1.attention.backend import (
     AttentionMetadata,
     AttentionMetadataBuilder,
     CommonAttentionMetadata,
-    SparseMLAAttentionImpl,
+    MLAAttentionImpl,
 )
 from vllm.v1.attention.backends.mla.flashmla_sparse import (
     triton_convert_req_index_to_global_index,
@@ -170,7 +170,9 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]
         return metadata
 
 
-class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):
+class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]):
+    is_sparse = True
+
     def __init__(
         self,
         num_heads: int,