Add SM80 (Ampere/A100) dense MLA decode kernel

[bzantium]

Summary

Implements a CUDA kernel for the dense MLA decode path on SM80 GPUs (Ampere / A100). The current README excludes Ampere; this PR enables A100 deployment without forcing migration to Hopper.

• Drop-in: dense_decode_fwd API and the combine kernel are unchanged. Only the decode kernel itself is new.
• Correctness verified against a PyTorch eager BMM reference across BF16/FP16, multi-batch, multi-KV-head, causal mask, and split-K.
• Build flag: FLASH_MLA_DISABLE_SM100=1 FLASH_MLA_DISABLE_SM90=1 pip install -v . for A100-only nodes.

Why

The SM90 kernel relies on TMA, WGMMA, thread block clusters, and mbarrier async barriers — all sm90+ only. SM80 must use cp.async + mma.m16n8k16 + ldmatrix + cutlass::arch::NamedBarrier. This is a separate kernel rather than a SM90 fallback path.

Design

• 256 threads / CTA = 4 warpgroups × 1 warp × 32 lanes
• BLOCK_M=16 (one warp covers M via mma.m16n8k16); each wg owns HEAD_DIM_V/4 = 128 V columns
• All wgs compute QK^T independently (no cross-wg P transfer); compute is not the bottleneck
• 2× sK SMEM buffer + cross-block cp.async prefetch
• XOR swizzle (Swizzle<3,3,3>-equivalent) for SMEM bank-conflict-free ldmatrix
• cp.async.cg (L1 bypass) for K loading; cp.async.ca for Q (reused across K iterations)
• SMEM 162 KB / CTA = 18 KB sQ + 2 × 72 KB sK, fits within the 164 KB opt-in cap

Detailed design rationale and the SMEM budget tradeoffs are in docs/sm80-dense-decode-design.md.

Performance

Idle A100-SXM4-80GB, CUDA 12.9, exclusive node. Achievable copy bandwidth measured on
the same node is 1747 GB/s, which is what the percentages below are against; the 2039 GB/s
datasheet figure is not reachable by any kernel.

Both versions were built in one allocation and run in five alternating rounds. Median of
five, range in brackets.

Config	base	this PR	delta	ranges disjoint
b=64 sk=4096	490.8 GB/s	632.9 GB/s	+29.0%	yes
b=1 hq=64 sk=4096	37.9	45.1	+19.0%	yes
b=64 sk=1024	318.4	351.3	+10.3%	yes
b=1 sk=16384	135.3	141.6	+4.7%	yes
b=1 sk=256	4.4	4.6	+4.5%	yes
b=1 sk=1024	16.4	17.0	+3.7%	yes
b=1 sk=4096	50.9	52.4	+2.9%	yes
b=1 sk=65536	270.0	285.1	+5.6%	no
b=4 sk=1024	63.4	65.6	+3.5%	no
b=4 sk=4096	160.4	169.2	+5.5%	no
b=16 sk=1024	153.7	159.1	+3.5%	no
b=16 sk=4096	356.2	364.6	+2.4%	no

The last five move in the right direction but their ranges overlap across rounds, so treat
them as unchanged. b=64 sk=4096 goes from 28% to 36% of achievable bandwidth.

Registers drop from 255, the SM80 ceiling, to 206, and the 252/348 byte spill disappears.
DRAM throughput goes from 22.1% to 30.7% of peak.

How the occupancy work went

Nsight put the original kernel at 6.25% occupancy, Block Limit Shared Mem = 1, with no
hardware unit above 37% of peak: latency bound with one CTA per SM and nothing to run while
it waits on __syncthreads().

Staging a 64-token KV page twice costs 162 KB. Two CTAs need roughly 81.5 KB each, and sQ
alone is 18 KB, so K staging had to fall from 128 tokens to at most 56. A 16-token tile with
three stages is 54 KB and also deepens the prefetch. The page size stays 64 and each page is
walked as four sub-tiles.

Shared memory was not the only place assuming one resident CTA. The split count came from
the SM count, so the grid still supplied 108 CTAs for 108 SMs and achieved occupancy stayed
at 6.25% after the theoretical limit doubled. It now scales with the CTAs that actually fit,
derived from the kernel's SMEM footprint. The combine kernel's split ladder stopped at 160
and asserted past it, so it gains rungs to 256.

The extra partitions need work to fill them: each reloads the 18 KB Q tile, runs a prologue
and epilogue, and adds a combine row. Doubling unconditionally cost 19.7% at b=16 sk=1024,
where 216 partitions share 256 KV pages. Two gates, both empirical and marked as such in the
code, keep the scaling to cases the sweep shows it helps.

Full sweep and the earlier milestones are in docs/sm80-benchmark-2026-05-01.md.

Limitations / future work

36% of achievable bandwidth is still short of the SM90 path's ~80% on H800, but the
profile does not blame the missing Hopper instructions for most of it. Warp stall cycles at
b=64 sk=4096, as ratios per issued instruction:

reason	cycles	share of stall
wait (fixed-latency execution dependency)	1.80	33%
short_scoreboard (MIO, i.e. shared memory)	1.67	30%
long_scoreboard (global memory)	0.59	11%
mio_throttle	0.56	10%
barrier	0.34	6%
dispatch, lg_throttle, math pipe	0.47	8%

Global memory is not the constraint. DRAM reads 304 MB against a 302 MB KV cache, so the
kernel already streams it exactly once with no redundancy, and global dependency accounts
for 11% of the stall. The 40% in the MIO reasons and the 33% in arithmetic dependency are
what is left.

• Shared-memory traffic is the largest addressable item. All four warpgroups compute QK^T
  independently, so each ldmatrix over a K tile happens four times. That duplication is
  deliberate and documented, and it was measured against an 8-warpgroup variant back when
  shared memory was 162 KB and only one CTA fit. At 72 KB the trade is worth re-testing.
• The wait share is the online-softmax rescale chain, which runs every tile and now runs
  four times as often since the tile shrank from 64 tokens to 16. Rescaling less frequently
  would shorten it.
• Bank conflicts are 19.7k against 63.4M instructions, so the swizzle still holds after the
  tile change and is not worth touching.
• TMA and WGMMA genuinely cannot be used here: ptxas rejects cp.async.bulk.tensor,
  mbarrier::complete_tx::bytes and wgmma.mma_async for .target sm_80. Where TMA would
  have helped is visible in the counters: L1 requests 19.2M sectors, 615 MB, against 304 MB
  from DRAM, because cp.async moves at most 16 bytes per instruction and two of them land
  in the same 32-byte sector. That is a request-rate cost, not a bandwidth one.
• The ~80% figure on H800 is also not a like-for-like target. That part has 228 KB of shared
  memory per SM and far more HBM bandwidth, so occupancy and staging arithmetic both come
  out differently.
• Occupancy is now 12.5%, limited equally by shared memory and by registers at 206 per
  thread. A third CTA per SM needs both below their thresholds, not just one.
• The two partition-scaling gates are thresholds read off the A100 sweep. They are not
  derived, and another SM80 part may want different ones.
• SM80+SM90 combined builds currently fail because the sm90 sources use __launch_bounds__(N, M, K) (third arg is sm90 cluster). I left those untouched so this PR is minimal; happy to follow up with portable launch_bounds macros if maintainers prefer combined builds.
• A CuTe-based rewrite was attempted and hit cute's "Stride Divisibility Condition" for the (head_dim=576 non-pow2, BLOCK_M=16, dynamic stride) combo — this version uses raw PTX. CuTe path is a candidate for follow-up if maintainers prefer that style.
• No FP8 KV (sm89/sm90+ only). No sparse decode/prefill on SM80.

Modifications outside csrc/sm80/

File	Change
csrc/api/api.cpp, common.h, dense_decode.h	SM80 arch dispatch + Arch::is_sm80() helper. SM80 build only exposes dense_decode_fwd
csrc/smxx/decode/combine/combine.cu	#if __CUDA_ARCH__ >= 900 guard around cudaGridDependencySynchronize() (PDL is sm90+); split ladder extended from 160 to 256 rungs, needed once the SM80 path asks for more partitions. Existing callers resolve to the same constant as before
csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu	drop the sm90-only third arg from __launch_bounds__
setup.py	FLASH_MLA_DISABLE_SM80 flag, sm80 source list, plus include paths from pip-installed nvidia-* wheels (system CUDA can lack cusparse.h)

Test plan

• benchmark/bench_sm80_decode.py --check — all 12 sweep configurations match the PyTorch eager reference. Note that --check is silently ignored when --no-torch-baseline is also passed, since the comparison lives inside the baseline branch; worth tightening separately
• Sweep across (batch, sq, hq, hk, sk) configs (b=1..128, sk=256..65536, hk=1..4) — all PASS
• BF16 and FP16 instantiations
• is_causal=true for multi-token Q
• Split-K via combine kernel (num_sm_parts > 1)
• benchmark/profile_decode_step.py — DeepSeek-V3-shape attention block profile

Happy to make any code style / structure changes maintainers want before merge. If upstream prefers to keep the support matrix at SM90/SM100 only, this can also live as a maintained fork — see docs/sm80-distribution-strategy.md for that path.