ROCm / ROCm/aiter

[Test] Expand FP8 MHA harness coverage beyond 128x128

Open
#5,000 0 comments 0 reactions 0 assignees View on GitHub

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.py
  • op_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_QK is the last dimension of Q and K.
  • D_V is 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:

  • 192x128
  • 128x64
  • 256x256

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_func baseline 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_ref on the dequantized FP8 Q/K/V tensors, require the same -inf mask, 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_perftest measurements for both the quantized path and BF16 baseline (quant_fwd_us and fwd_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 at 128x128.
  • #3732, #4682, and #4733 added/fixed/tuned the gfx950 FP8 HD256 implementation and used manual HD256 harness commands; they did not add continuous 256x256 matrix coverage.
  • #4645 has specialized gfx942 192x128 Gluon 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.py and op_tests/test_mha_varlen_fp8.py retain 128x128 and include 192x128, 128x64, and 256x256 in 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.055 max-absolute-error gate without tolerance changes; optional LSE checks retain the < 0.01 gate and -inf mask 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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.