refactor(mla): redesign batch MLA backend internals

benchmarks/README.md

@@ -25,8 +25,21 @@ Currently supports testing attention, gemm, fused MOE, normalization, quantizati
     - `BatchPrefillWithRaggedKVCacheWrapper` - Prefill attention with ragged KV cache.
         - Also supports computationally similar `cudnn_batch_prefill_with_kv_cache` (cudnn-native) and  `trtllm_ragged_attention_deepseek`.
     - `BatchMLAPagedAttentionWrapper` - MLA attention proposed in DeepSeek series of models.
-        - Also supports computationally similar `trtllm_batch_decode_with_kv_cache_mla` (trtllm-native) and CuTe DSL MLA decode kernel (cute-dsl, SM100+).
-    - All four attention routines accept `--backends prims-ts` on SM100/SM103 to benchmark the experimental task-scheduled attention implementation. `prims_ts` is accepted as an alias.
+        - Benchmarks the stateful MLA wrapper for `fa2`, `fa3`, `cutlass`,
+          `trtllm-gen`, `cute-dsl-monolithic`, `cute-dsl-modular`, `xqa`, or
+          deterministic wrapper `auto`. Auto promotion is architecture-primary
+          product policy, not a total benchmark-derived performance order;
+          backend planners still determine exact support.
+        - On SM100/SM103, `prims-ts` is available as an adapter-backed
+          comparative candidate. It is not a public MLA wrapper backend.
+    - `batch_mla_paged_attention` - Dense functional MLA decode.
+        - Benchmarks the public functional API separately for `auto`, `xqa`,
+          `trtllm-gen`, `cute-dsl`, `fa2`, `fa3`, and `cutlass`. It never
+          dispatches through the wrapper benchmark and does not include sparse MLA.
+    - Stateful attention routines accept `--backends prims-ts` on SM100/SM103
+      to benchmark the experimental task-scheduled attention implementation.
+      `prims_ts` is accepted as an alias. The functional
+      `batch_mla_paged_attention` routine does not accept this adapter.
 - GEMM:
     - `gemm_fp8_nt_groupwise` - GEMM with FP8 data types using groupwise scaling.
     - `group_gemm_fp8_nt_groupwise` - Group GEMM with FP8 data types using groupwise scaling.
@@ -173,14 +186,153 @@ $ python3 flashinfer_benchmark.py --routine mxfp8_quantize --m 2048 --k 8192 --i
 [PERF] cuda           :: median time 0.016 ms; std 0.000 ms; achieved tflops 3.118 TFLOPs/sec; achieved tb_per_sec 3.150 TB/sec
 ```
 
+### MLA wrapper benchmark
+
+`BatchMLAPagedAttentionWrapper` uses one two-phase lifecycle for every
+requested backend. It first constructs, plans, and captures one synchronized
+output from every candidate, then computes the independent PyTorch reference
+from the same inputs. Each candidate whose output and optional LSE match that
+reference is admitted to warm, cold-L2, and optional CUDA-graph timing. If the
+reference cannot be computed, no candidate receives timing evidence. Candidate
+attempt order is randomized from `--random_seed`, but cannot affect the
+correctness reference. Incorrect outputs, planner rejections, and runtime
+failures have no timing evidence. The planner remains the source of truth for
+support.
+
+The runtime `auto` mode promotes FA3 on SM90, TRTLLM-GEN on SM100/SM103, and
+XQA on SM120/SM121; SM80/SM89 and unrecognized architectures retain a
+conservative FA2-first complete candidate list. This is product policy, not a
+total benchmark-derived performance order. Future measurements may justify
+architecture-local ordering changes; each backend planner remains authoritative
+for exact support.
+
+The MLA-only contract flags mirror wrapper planning:
+
+```bash
+python3 flashinfer_benchmark.py --routine BatchMLAPagedAttentionWrapper \
+  --backends fa2 fa3 cutlass trtllm-gen cute-dsl-monolithic \
+  cute-dsl-modular xqa auto --batch_size 8 --s_qo 2 --s_kv 1024 \
+  --page_size 64 --num_qo_heads 128 --num_kv_heads 1 \
+  --head_dim_ckv 512 --head_dim_kpe 64 --mla-qk-nope-head-dim 128 \
+  --mla-q-lengths 2,2,2,2,2,2,2,2 \
+  --mla-kv-lengths 1024,1024,1024,1024,1024,1024,1024,1024 \
+  --mla-metadata-form dual --mla-enable-pdl default \
+  --mla-lse-mode none --mla-kv-layout combined \
+  --mla-output-scale none --mla-scale-mode default
+```
+
+`--mla-lse-mode`, `--mla-kv-layout`, `--mla-output-scale`,
+`--mla-scale-mode`, and `--mla-skip-softmax` expose the five plan/run contract
+groups. `--mla-qk-nope-head-dim` declares one logical QK NoPE profile
+unchanged to every concrete backend and `auto`; the physical query remains
+`--head_dim_ckv` wide. Its value must be positive, or `none`; omitting it also
+declares `None`, which may make
+`trtllm-gen` unsupported, while `128` lets `auto` consider that backend and
+may make other planners unsupported. Those planner decisions remain output
+rows. The underscore spelling `--qk_nope_head_dim` remains an alias.
+
+The harness does not tune or update `backend="auto"` during a benchmark run.
+
+Exact `--mla-q-lengths` and `--mla-kv-lengths` avoid substituting a requested
+shape for a randomly observed ragged shape.
+`--mla-metadata-form` controls whether CSR, dense, or both equivalent forms
+reach every planner. `--mla-enable-pdl default` preserves a `None` plan
+declaration; `true` and `false` are explicit, while the shared `--enable_pdl`
+flag remains a true alias. Rows serialize this field uniformly as
+`default|true|false`. `--causal` is also forwarded to every MLA planner.
+`--autotune` is intentionally rejected for this routine; a future explicit
+`backend="autotune"` mode will own online tuning. MLA output rows include
+requested/resolved backend, explicit correctness, the declared QK NoPE width,
+structured status/reason, phase summaries with explicit sample counts
+(`cuda_graph_*` timings map to the `graph_replay` objective), workspace bytes,
+and
+`peak_memory_delta_bytes` when it is reliably available. Every public-wrapper
+candidate that survives capture retains a distinct 128 MiB workspace so later
+timing cannot alias another wrapper's state. PrimTS owns its internal workspace
+and does not receive this extra buffer. The peak field is candidate-relative: it is the
+larger of the capture and timing CUDA allocator peaks after subtracting the
+bytes allocated at the start of that phase, rather than the process-wide peak.
+Only `_BackendPlanUnsupportedError` raised by `plan()` is an unsupported
+result; the same typed exception from construction, first run, output cloning,
+trace access, or timing is an error. Planning facts already produced by a
+successful plan, including `resolved_backend` and auto rejection trace, remain
+in that error row.
+
+CSV output uses the Python CSV writer for both headers and rows. Missing output
+columns are backfilled from the parsed routine arguments; values emitted by a
+routine are preserved. Both bare `--help` and selected-routine help exit
+successfully. A local run is not performance evidence.
+
+The correctness gate compares every candidate against the independent PyTorch
+reference. It does not infer correctness from agreement between concrete
+backends or from `auto` resolving to the same concrete implementation.
+
+### Functional MLA benchmark
+
+The functional routine benchmarks only calls through
+`flashinfer.mla.batch_mla_paged_attention`. It uses a separate
+name because functional `auto` and wrapper `auto` have different lifecycles
+and backend sets:
+
+```bash
+python3 flashinfer_benchmark.py \
+  --routine batch_mla_paged_attention \
+  --backends trtllm-gen cute-dsl fa2 fa3 cutlass auto \
+  --batch_size 2 --s_qo 1 --s_kv 512 --page_size 64 \
+  --num_qo_heads 128 --head_dim_ckv 512 --head_dim_kpe 64 \
+  --mla-qk-nope-head-dim 128 \
+  --mla-q-lengths 1,1 --mla-kv-lengths 512,384 \
+  --mla-is-var-seq false --mla-cute-dsl-impl auto \
+  --mla-enable-pdl default
+```
+
+Inputs are deterministic for `--random_seed`. Every requested backend is
+captured before any candidate is timed, and timing is published only after its
+output and optional LSE match the independent PyTorch reference. Unsupported,
+incorrect, reference-unavailable, and runtime-error rows retain their status
+but have no timing evidence.
+
+The routine preserves the public functional controls for variable query
+metadata, shared or duplicated page indices, PDL, LSE, sinks, skip-softmax
+threshold scaling, BMM scaling, and CuTe DSL implementation selection.
+`--autotune` is supported only when functional `auto` is requested: it profiles
+once under `flashinfer.autotune(True, cache=...)`, then captures and measures
+under `flashinfer.autotune(False, cache=...)`. Functional `auto` leaves
+`resolved_backend` empty because the public API does not expose an
+authoritative resolution trace.
+
+The public functional entrypoint performs backend planning on each call.
+Functional benchmarking therefore defaults to eager first-run, warm, and
+cold-L2 measurements. `--cuda-graph` explicitly requests the optional graph
+phase. If graph capture is unavailable, the row preserves correctness and
+eager timing, reports zero graph repetitions, and records the graph failure in
+`mla_reason`. Use the stateful wrapper benchmark when graph-replay performance
+is required.
+
+Functional rows map the warm median and standard deviation into the shared
+`median_time` and `std_time` columns while retaining detailed phase columns.
+The CSV header defines the expanded output contract.
+
+This is the supported replacement for the removed mixed functional/wrapper
+benchmark. Legacy command lines, CSV rows, and performance baselines are not
+reproduced. Packed sparse MLA remains out of scope and requires a separate
+benchmark.
+
 ### Batch Testing
 
-Run multiple tests from a file and save results:
+Run multiple tests from a file and write generated results outside the source
+tree:
+
 ```bash
-python3 flashinfer_benchmark.py --testlist samples/sample_testlist.txt --output_path samples/sample_testlist_output.csv
+python3 flashinfer_benchmark.py \
+  --testlist samples/sample_testlist.txt \
+  --output_path /tmp/flashinfer-benchmark-output.csv
 ```
 
-See `samples/sample_testlist.txt` for an example stdout output from the above command; `samples/sample_testlist_output.csv` for csv output from the same run.
+`samples/sample_testlist.txt` is an executable command example, not a
+performance baseline. Generated CSV and stdout depend on the GPU architecture,
+CUDA/PyTorch/FlashInfer versions, and timing provider and are intentionally not
+checked in.
 
 The output CSV will contain detailed metrics including:
 - Median execution time
@@ -207,7 +359,7 @@ The output CSV will contain detailed metrics including:
 | `--verbose`, `-v`        | Print additional information (can be used multiple times for more verbosity, e.g. `-vv`)                   |
 | `--case_tag`              | Optional tag for the test case, useful for annotating or filtering results in the output CSV.              |
 | `--generate_repro_command`| If set, prints a reproducer command for the test case and stores it in the output CSV.                     |
-| `--backends`             | Space-separated list of backends to test, e.g. fa2, fa2_tc, fa3, auto, cudnn, cudnn-native, cutlass, trtllm, trtllm-gen, trtllm-native, prims-ts, cute-dsl, cublas, trtllm_low_latency. (`prims_ts` aliases `prims-ts`; `auto` support is routine-dependent.)|
+| `--backends`             | Space-separated list of backends to test, e.g. fa2, fa2_tc, fa3, auto, cudnn, cudnn-native, cutlass, trtllm, trtllm-gen, trtllm-native, prims-ts, cute-dsl, cute-dsl-monolithic, cute-dsl-modular, xqa, cublas, trtllm_low_latency. (`prims_ts` aliases `prims-ts`; `auto` support is routine-dependent and includes the deterministic architecture-preferred `BatchMLAPagedAttentionWrapper` routine.)|
 
 ### Attention Flags
 | Flag                     | Description                                                                                                 |
@@ -491,7 +643,11 @@ Notes:
 - Prefill pre-L2-normalizes k and calls the kernel with `use_qk_l2norm_in_kernel=False` so the kernel and reference see identical inputs.
 
 ## `flashinfer_benchmark.py` Routine & Backend Support Matrix
-The following table summarizes the support surface of each routine & backend's on various [CUDA Compute Capabilities](https://developer.nvidia.com/cuda-gpus).
+The following table summarizes each routine's declared backend capability on
+various [CUDA Compute Capabilities](https://developer.nvidia.com/cuda-gpus).
+It is not measured performance evidence: backend planners remain authoritative,
+and SM103/SM121 entries are unvalidated unless accompanied by separate runtime
+evidence.
 
 Each column represents a compute capability. Backends inside cells represent supported backends. A blank cell means no backend is supported for that routine at that compute capability.
 
@@ -513,7 +669,8 @@ Legend:
 | **BatchDecodeWithPagedKVCacheWrapper** | fa2 | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn | fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native, prims-ts | fa2, fa2_tc, cudnn, trtllm-gen, trtllm-native, prims-ts | fa2, fa2_tc, cudnn |
 | **BatchPrefillWithPagedKVCacheWrapper** |  | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, fa3, cudnn, cudnn-native | fa2, cudnn, cudnn-native, trtllm-gen, trtllm-native, prims-ts | fa2, cudnn, cudnn-native, trtllm-gen, trtllm-native, prims-ts | fa2, cudnn, cudnn-native |
 | **BatchPrefillWithRaggedKVCacheWrapper** |  | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, cudnn, cudnn-native | fa2, fa3, cudnn, cudnn-native | fa2, cudnn, cudnn-native, cutlass, trtllm-native, prims-ts | fa2, cudnn, cudnn-native, cutlass, trtllm-native, prims-ts | fa2, cudnn, cudnn-native |
-| **BatchMLAPagedAttentionWrapper** |  | fa2 | fa2 | fa2 | fa2, fa3 | fa2, cutlass, trtllm-native, cute-dsl, prims-ts | fa2, cutlass, trtllm-native, prims-ts | fa2 |
+| **BatchMLAPagedAttentionWrapper** |  | fa2, auto | fa2, auto | fa2, auto | fa2, fa3, auto | fa2, cutlass, trtllm-gen, cute-dsl-monolithic, cute-dsl-modular, prims-ts, auto | fa2, cutlass, trtllm-gen, cute-dsl-monolithic, cute-dsl-modular, prims-ts, auto | fa2, xqa, auto |
+| **batch_mla_paged_attention** |  | fa2 | fa2 | fa2 | fa2, fa3 | fa2, cutlass, trtllm-gen, cute-dsl, auto | fa2, cutlass, trtllm-gen, cute-dsl, auto | fa2, xqa, auto |
 | **gemm_fp8_nt_groupwise** |  |  |  |  |  | cutlass | cutlass |  |
 | **group_gemm_fp8_nt_groupwise** |  |  |  |  |  | cutlass | cutlass |  |
 | **bmm_fp8** |  |  |  | cudnn, cublas | cudnn, cublas | cudnn, cublas, cutlass | cudnn, cublas, cutlass | cudnn, cublas |

benchmarks/bench_trtllm_gen_mla.py

@@ -106,7 +106,7 @@ def bench_trtllm_mla(
 
     # Run decode-MLA
     # warmup
-    flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(**common_kwargs)
+    flashinfer.mla.batch_mla_paged_attention(**common_kwargs)
     # benchmark
     # NOTE: cold_l2_cache=True is requested but silently degrades to warm-L2
     # because the inputs are captured in the lambda's closure rather than
@@ -116,9 +116,7 @@ def bench_trtllm_mla(
     # cross-backend comparisons remain fair, only absolute GB/s numbers
     # are optimistic vs. a real cold-cache serving workload.
     measurements = bench_gpu_time(
-        lambda: flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
-            **common_kwargs
-        ),
+        lambda: flashinfer.mla.batch_mla_paged_attention(**common_kwargs),
         dry_run_iters=5,
         repeat_iters=30,
         enable_cupti=False,

benchmarks/flashinfer_benchmark.py

@@ -1,4 +1,5 @@
 import argparse
+import csv
 import sys
 
 # Only import utilities at module level - routine modules are imported lazily
@@ -19,7 +20,11 @@ def run_test(args):
 
     ## Depending on routine type, route to corresponding test routine
     ## Imports are done lazily to avoid loading unnecessary dependencies
-    if args.routine in benchmark_apis["attention"]:
+    if args.routine in benchmark_apis["mla"]:
+        from routines.mla import run_mla_test
+
+        res = run_mla_test(args)
+    elif args.routine in benchmark_apis["attention"]:
         from routines.attention import run_attention_test
 
         res = run_attention_test(args)
@@ -80,7 +85,8 @@ def run_test(args):
 
     # Write results to output file if specified
     if args.output_path is not None:
-        with open(args.output_path, "a") as fout:
+        with open(args.output_path, "a", encoding="utf-8", newline="") as fout:
+            writer = csv.writer(fout)
             for cur_res in res:
                 for key in full_output_columns:
                     # Backfill every output column the routine didn't set: from
@@ -89,11 +95,7 @@ def run_test(args):
                     # KeyError below.  Routine-set values are preserved.
                     if key not in cur_res or cur_res[key] == "":
                         cur_res[key] = getattr(args, key, "")
-
-                output_line = ",".join(
-                    [str(cur_res[col]) for col in full_output_columns]
-                )
-                fout.write(output_line + "\n")
+                writer.writerow([cur_res[col] for col in full_output_columns])
             fout.flush()
     return
 
@@ -111,13 +113,16 @@ def parse_args(line=sys.argv[1:]):
     """
 
     ## Shared arguments
-    parser = argparse.ArgumentParser()
+    # Delay installing --help until after routine dispatch. Otherwise the
+    # preliminary parse exits early and hides the selected routine's options.
+    parser = argparse.ArgumentParser(add_help=False)
     parser.add_argument(
         "--routine",
         "-R",
         type=str,
-        required=True,
-        choices=list(benchmark_apis["attention"])
+        required=False,
+        choices=list(benchmark_apis["mla"])
+        + list(benchmark_apis["attention"])
         + list(benchmark_apis["gemm"])
         + list(benchmark_apis["moe"])
         + list(benchmark_apis["moe_comm"])
@@ -134,6 +139,10 @@ def parse_args(line=sys.argv[1:]):
     )
     args, _ = parser.parse_known_args(line[:])
 
+    parser.add_argument(
+        "-h", "--help", action="help", help="show this help message and exit"
+    )
+
     parser.add_argument(
         "--no_cuda_graph",
         action="store_true",
@@ -237,9 +246,20 @@ def parse_args(line=sys.argv[1:]):
         ),
     )
 
+    if args.routine is None:
+        # `--help` exits successfully after all shared options are installed.
+        # Other routine-less invocations retain the required-routine error.
+        if "-h" in line or "--help" in line:
+            parser.parse_args(line)
+        parser.error("the following arguments are required: --routine/-R")
+
     ## Check routine and pass on to routine-specific argument parser
     ## Imports are done lazily to avoid loading unnecessary dependencies
-    if args.routine in benchmark_apis["attention"]:
+    if args.routine in benchmark_apis["mla"]:
+        from routines.mla import parse_mla_args
+
+        args = parse_mla_args(line, parser, args.routine)
+    elif args.routine in benchmark_apis["attention"]:
         from routines.attention import parse_attention_args
 
         args = parse_attention_args(line, parser)
@@ -334,8 +354,8 @@ def parse_args(line=sys.argv[1:]):
 
     # Setup output file if specified
     if testlist_args.output_path is not None:
-        with open(testlist_args.output_path, "w") as fout:
-            fout.write(",".join(full_output_columns) + "\n")
+        with open(testlist_args.output_path, "w", encoding="utf-8", newline="") as fout:
+            csv.writer(fout).writerow(full_output_columns)
 
     # Process tests either from testlist file or command line arguments
     if testlist_args.testlist is not None:

benchmarks/mla/__init__.py

@@ -0,0 +1 @@
+"""Independent reference helpers for MLA correctness benchmarks."""