flash_mla v2.0.0

aarch64 wheel: narrow the gencode matrix, serialize nvcc

The job was OOM-killed ~50 minutes in -- 'The hosted runner lost communication
with the server', no logs uploaded, which is the signature of a terminated
runner rather than a failed compile. The hosted ubuntu-24.04-arm runner is
4 cores / 16 GB and the full 7-gencode matrix with parallel nvcc does not fit.

No aarch64 NVIDIA host ships an sm_80/86/89/120 GPU, so this wheel only needs
Grace-class parts: sm_90 (GH200), sm_100 (GB200), sm_121 (GB10/DGX Spark). The
x86_64 wheels keep the full matrix. MAX_JOBS and NVCC_THREADS drop to 1.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_014PzQ5tQQaJazSBhiMuR6KL

Nightly 20260721-052127 (8465684)

nightly-84656849d05681ae3457bfceb3b0f4354bf60eca

test: bound int8 prefill peak memory by the selection, not the pool

Nightly 20260718-151414 (fdba9cd)

nightly-fdba9cdf39590f260d531305863c543f8f1b5b82

Remove pre-harmonization kernel name aliases; breaking rename

Nightly 20260718-151211 (b9f8478)

nightly-b9f847847be56202bd5638d5d141d00173b3cfb5

Harmonize kernel names to appmana roles: sparse_mla_prefill/decode_fp…

Nightly 20260718-151627 (8ec3de6)

nightly-8ec3de601c18205a146d5496af2a8a1338bd56f7

Bump version to 2.0.0 for the breaking kernel rename

Nightly 20260713-194816 (eb855ae)

nightly-eb855ae645b615aba4bbca726c6fe0787cbbbe4b

decode: native int8_ds_mla entry (fwd_sparse_int8_decode_mla) via the…

Nightly 20260713-192125 (de88e03)

nightly-de88e030ee0b30d2c77a0a86c90774cfe21de0cb

wheels: trigger on appmana/vllm-ampere (branch renamed)

Nightly 20260710-003535 (a9738dd)

Remove the legacy staged sparse-MLA prefill kernel

The staged two-kernel gather+attention prefill path was dead code: the fused
tensor-core kernel (csrc/flash_sparse_mla_prefill_fused_sm80.cu) had already
replaced it as the default, and no production caller reached the staged path
(the vLLM fork calls flash_sparse_mla_prefill without use_staged_prefill, which
always selected fused). It was reported to abort at runtime on torch 2.11 via a
stable-ABI torch_call_dispatcher "aten::new_empty" failure. It is also much
slower at the true 16k footprint, so it is removed outright.

Removed:
- csrc/flash_sparse_mla_prefill_staged_sm80.cu (888 lines) and its setup.py entry.
- The FLASH_MLA_PREFILL_FUSED=0 kill-switch and sparse_mla_prefill_fused_enabled();
  run_sparse_mla_prefill now always runs the fused kernel (fp8 and int8).
- The use_staged_prefill / staged_chunk_tokens kwargs and the Triton/decode
  fallbacks in flash_sparse_mla_prefill (signature otherwise unchanged; no caller
  passed the kwarg -- confirmed in the vLLM fork).
- benchmarks/profile_prefill_stages.py (staged-only) and the staged bench rows.

Renamed to drop the now-misleading "staged" from the prefill ABI:
  torch op  fwd_sparse_prefill_staged_mla       -> fwd_sparse_prefill_mla
  torch op  fwd_sparse_int8_prefill_staged_mla  -> fwd_sparse_int8_prefill_mla
  struct    Sparse_mla_prefill_staged_params    -> Sparse_mla_prefill_params
  python    flash_sparse_mla_prefill_native_staged -> (folded into flash_sparse_mla_prefill)
  python    sparse_int8_mla_prefill_native_staged  -> sparse_int8_mla_prefill

The decode kill-switch FLASH_MLA_DECODE_FUSED=0 is KEPT: the legacy decode kernel
runs correctly on torch 2.11 (tested), so it stays as a useful A/B. Only the
prefill staged path was broken/removed.

vLLM fork caller to update (not touched here): docker/Dockerfile.ampere-flashmla
imports sparse_int8_mla_prefill_native_staged -> sparse_int8_mla_prefill.

Test evidence (RTX A5000 sm_86, torch 2.11.0+cu130, CUDA_VISIBLE_DEVICES=1):
  before (pristine tag build):
    34 passed in 8.25s
  after (staged removed):
    30 passed in 6.52s
  Delta of 4 == the staged-only parametrizations deleted (use_staged True/False):
    test_prefill_adversarial_parity  -2, test_prefill_swa_only_parity -1,
    test_sparse_mla_prefill_two_stream_parity -1. All fused cases stay green.

Fused prefill unchanged at the true 16k footprint (T=1024, ctx 16384, H=64):
  fp8  width 512  : 2.998 ms (2.93 us/tok)
  fp8  width 1024 : 5.258 ms (5.13 us/tok)
  int8 width 1024 : 5.280 ms
Fused vs removed staged, head-to-head (this box):
  width 1024: fused 5.3 ms vs staged 22.2 ms (4.2x)
  width 2048: fused 9.9 ms vs staged 44.8 ms (4.5x)

Nightly 20260709-220002 (a2e2473)

nightly-a2e2473c230667669bb6786b909129118f55bd3d

decode H6: warp-wide combine -- mlse one-split-per-lane + branchless …

Nightly 20260709-211933 (3d34783)

nightly-3d347835d6f2301725077065f4689f086a6a82c6

docs: decode 16k-footprint sanity check — no microbench artifact, ALU…