[MLA] Deterministic single-split decode option for reproducible serving (fixes #4364)

[MohitAMD]

Motivation

Closes #4364.

The MLA decode split-KV reduce is fp-non-associative, so a request's decode
output depends on how many KV splits it is given. The per-request split count is
drawn from the global budget min(cu_num, max_split_per_batch * batch_size)
(see csrc/kernels/mla/metadata/v1_2_device.cuh), which varies with batch
composition. Two identical requests at temperature 0 can therefore land on
different split counts on different steps and produce slightly different logits.

On GLM-5.1-FP8 (DSA, qh16 fp8 decode) this surfaces as run-to-run
nondeterminism — ±1–2 retrieved NIAH needles at mid context (20k–35k), while the
extremes (8k, 96k) are stable (#4364).

Reproducer

op_tests/test_mla_qh16_decode_determinism.py (added here) isolates the mechanism
on a single MI300X (gfx942), nhead=16 fp8/fp8 persistent decode, batch=1, ctx=8192:

[det] split=1 run-to-run  max|d| = 0.000e+00   (single-split reduce is a no-op)
[det] split=8 run-to-run  max|d| = 0.000e+00   (kernel is deterministic for FIXED metadata)
[det] split=1 vs split=8 max|d| = 2.441e-03   (the split COUNT changes the numerics)

So the kernel is bit-stable for a fixed split count, but the result depends on
the split count — and production picks that count from the batch-dependent
budget, which is why identical requests drift run-to-run.

Workaround (this PR)

Opt-in env knobs, honored in get_mla_metadata_v1:

• AITER_MLA_DECODE_DETERMINISTIC=1 — force a single split per request
  (shorthand for the cap below = 1).
• AITER_MLA_DECODE_MAX_SPLIT_PER_BATCH=<n> — pin the per-batch split cap to a
  fixed n (≥1), independent of batch size.

Forcing a single split makes the reduce a no-op, so a request's output no longer
depends on batch composition and is reproducible. Clamping the cap down is
always buffer-safe (it can only reduce the number of partial-reduce entries
relative to what the caller already allocated).

Validation (same reproducer, AITER_MLA_DECODE_DETERMINISTIC=1):

[det] split=1 vs split=8 max|d| = 0.000e+00   (clamp collapses split=8 onto split=1)
[det] PASS: deterministic single-split path.

Trade-off: forcing a single split removes decode split-K parallelism, so long
context / low batch decode loses some latency headroom. It is therefore opt-in
and intended for serving that requires reproducible outputs. The durable fix is
an order-stable reduce in the kernel (fixed split count + deterministic
accumulation order), which would keep both determinism and split-K performance;
happy to help test such a kernel change.

Test plan

• op_tests/test_mla_qh16_decode_determinism.py on MI300X (gfx942): baseline
  shows split-count-dependent numerics (2.44e-3) with each fixed split count
  bit-stable; AITER_MLA_DECODE_DETERMINISTIC=1 collapses split=8 onto split=1
  (max|d|=0) and PASSES.
• Broader shapes (varying nhead/ctx/batch) and a perf note on the single-split
  path at long context — reviewer/maintainer guidance welcome.

Env: ROCm 7.2.3, MI300X (gfx942), AITER v0.1.18, PyTorch 2.11, GLM-5.1-FP8.

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 4378 --add-label <label>

[Copilot]

Pull request overview

Adds an opt-in deterministic MLA decode mode by clamping the per-batch KV split cap via environment variables inside get_mla_metadata_v1, plus a standalone op_tests/ reproducer for issue #4364 to demonstrate split-count-dependent numerics and validate the workaround.

Changes:

• Add env-var-controlled clamping of max_split_per_batch in aiter/ops/attention.py to force single-split (or a fixed cap) for reproducible serving.
• Add op_tests/test_mla_qh16_decode_determinism.py to reproduce and validate the determinism workaround on qh16 fp8/fp8 decode.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
op_tests/test_mla_qh16_decode_determinism.py	Adds a minimal standalone repro/validator for split-count-driven nondeterminism and the deterministic single-split workaround.
aiter/ops/attention.py	Adds env knobs to clamp decode split cap inside get_mla_metadata_v1 to stabilize numerics across batch composition.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.