[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>
2026-07-14 00:29:56 +00:00
parent 0b0ef8d7eb
commit c4f5cd60da
21 changed files with 1990 additions and 646 deletions
@@ -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
@@ -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)


@@ -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"
@@ -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,
    )
@@ -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 |
@@ -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
    )
@@ -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}"
@@ -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",
    [
@@ -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,
@@ -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,
@@ -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)
+        )
@@ -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],
@@ -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,
@@ -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
@@ -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,
@@ -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
@@ -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,
            )
        )
@@ -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
@@ -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,
@@ -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")
@@ -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,