[Test] Expand FP8 MHA harness coverage beyond 128x128
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 565
- Forks
- 585
- Avg merge
- 3d 4h
- Merged PRs (30d)
- 366
Description
Background
The fixed-shape and variable-length per-tensor FP8 MHA harnesses currently parameterize only (d_qk, d_v) = (128, 128):
op_tests/test_mha_fp8.pyop_tests/test_mha_varlen_fp8.py
Both harnesses already expose independent --d_qk and --d_v CLI arguments. In particular, #1555 used 192,128 with both files, while the HD256 work in #3732, #4682, and #4733 used the same harnesses through manual CLI invocations. Those paths are therefore testable, but the pytest matrix discovered by standard CI still continuously covers only 128,128.
For this issue, each pair is (D_QK, D_V):
D_QKis the last dimension of Q and K.D_Vis the last dimension of V and the output.
Requested coverage
Keep the existing 128x128 case and add these pairs to the pytest matrices in both harnesses:
192x128128x64256x256
Use the existing fixed-shape and varlen public FP8 dispatch paths. This is harness coverage, not a request to force one architecture-specific backend.
Baseline and correctness
Preserve the current comparison contract:
- Quantize the generated BF16 Q/K/V tensors with the existing per-tensor FP8 quantization path.
- Compare the FP8 output with the existing BF16
flash_attn_func/flash_attn_varlen_funcbaseline generated from the same source tensors. - Require output max absolute difference
< 0.055; do not widen this tolerance to make a new shape pass. - When LSE validation is requested, use
attention_refon the dequantized FP8 Q/K/V tensors, require the same-infmask, and require finite-value max absolute difference< 0.01.
Keep the existing NRMS and timing output so failures and numerical drift remain diagnosable.
Run and performance requirements
- Exercise the new pairs in both fixed-shape and varlen harnesses across the existing causal/non-causal, head-count, batch, and sequence-length parameterization.
- Keep
run_perftestmeasurements for both the quantized path and BF16 baseline (quant_fwd_usandfwd_us, plus the derived throughput metrics) for every executed pair. - Run through the existing standard Aiter test matrix on MI35X and MI300X. If a pair is intentionally unsupported on an architecture, report an explicit pytest skip with the capability reason instead of removing or commenting out the pair.
- Performance is report-only for this issue. There is no evidence-backed absolute speedup or regression threshold in the current harness, so this issue should not invent one.
Related work and non-duplication
- #1555 added separate QK/V dimensions and manually validated
192x128, but left both pytest matrices at128x128. - #3732, #4682, and #4733 added/fixed/tuned the gfx950 FP8 HD256 implementation and used manual HD256 harness commands; they did not add continuous
256x256matrix coverage. - #4645 has specialized gfx942
192x128Gluon coverage in its own test file; it does not expand the two generic FP8 MHA harnesses.
No existing issue or PR was found that adds all three requested pairs to both generic FP8 MHA harnesses.
Acceptance criteria
- Both
op_tests/test_mha_fp8.pyandop_tests/test_mha_varlen_fp8.pyretain128x128and include192x128,128x64, and256x256in their pytest dimension matrices. - Each new pair executes the fixed-shape and varlen FP8 public paths and the existing BF16 baseline for both causal and non-causal coverage, subject only to explicit architecture capability skips.
- Output correctness passes the existing
< 0.055max-absolute-error gate without tolerance changes; optional LSE checks retain the< 0.01gate and-infmask equality. - The existing MI35X and MI300X standard test jobs pass, with no requested pair silently omitted.
- Quantized and BF16 baseline timing metrics are emitted for every executed pair; no new performance threshold is added.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in op_tests/test_mha_fp8.py and op_tests/test_mha_varlen_fp8.py by comparing their existing pytest dimension matrices and parameterization. Run the standard Aiter test matrix on MI35X and MI300X, checking fixed-shape and varlen paths, correctness gates, LSE checks, and timing output. Done means both files cover 128x128, 192x128, 128x64, and 256x256, with explicit capability skips only where required.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, testing-qa
- Issue type
- Feature
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 68/100