Skip to content

[AIROCMLIR-1077] Fix intermittent f32 attention tuning verification false-fail in Tune rocMLIR#2431

Draft
bogdan-petkovic wants to merge 2 commits into
developfrom
bogdan-petkovic/f32-attention-tuning-verify
Draft

[AIROCMLIR-1077] Fix intermittent f32 attention tuning verification false-fail in Tune rocMLIR#2431
bogdan-petkovic wants to merge 2 commits into
developfrom
bogdan-petkovic/f32-attention-tuning-verify

Conversation

@bogdan-petkovic

Copy link
Copy Markdown
Contributor

Motivation

The weekly Tune rocMLIR stage (tuningRunner.py --op attention) intermittently fails f32 attention configs on gfx90a and gfx942 during CPU verification. The verifier reports [1 1 0]: only the relDiff check fails while RMS and absDiff pass. The GPU output is numerically sound — the failure is a benign accumulation-order divergence between the GPU flash / online-softmax reduction and the sequential CPU reference, which the strict relDiff_threshold=1e-4 over-penalizes on a small fraction of elements. It surfaces intermittently because tuningRunner verifies only the fastest (winning) perf_config per problem, and that winner varies by architecture and from run to run.

The goal of this PR is to stop these false-fails without weakening real regression detection.

Technical Details

Root cause: tuningRunner.py's CPU verification emits -pv -relDiff_threshold=0.0001 -RMS_threshold=0.15 uniformly for every op. For f32 attention, a small fraction of output elements (~0.05–0.45%) exceed the 1e-4 relDiff threshold (observed maxRelDiff ranged from ~1.6e-3 up to ~1.5e-2) while RMS and average relDiff stay ~1e-3 or better. This is benign: the GPU computes genuine f32 throughout (amdgpu.mfma ... : f32, f32, vector<4xf32>, softmaxType = f32, no f16/bf16 downcast), and re-running the exact failing config with a relaxed relDiff yields [1 1 1] with identical GPU output. rocmlir-gen already disables relDiff for f16/bf16 and i8 attention uses integer verification, so only f32 attention is affected.

The verifier (mcpuVerify in mlir/lib/ExecutionEngine/conv-validation-wrappers.cpp) already supports an allclose-style gate: when -absDiff_threshold is passed, relDiff is ignored for elements whose absolute difference is within tolerance, while RMS and absDiff still gate correctness. This mechanism is already used for f32 attention in mlir/utils/performance/parameterSweeps.py and mlir/test/e2e/PrAttentionF32.toml, but it was never wired into tuningRunner.py.

This PR adds attention_cpu_verify_flags() in tuningRunner.py, which appends -absDiff_threshold 5e-2 for f32 attention CPU verification, invoked from verify_perfconfig (user-supplied --rocmlir-gen-flags still take precedence). RMS (0.15) remains the primary correctness metric and absDiff still bounds gross per-element errors. The 5e-2 tolerance covers the largest observed benign divergence (~1.5e-2 on gfx90a) with margin, while a genuinely wrong element (abs error on the order of 0.1–1 for bounded attention output) still fails absDiff and RMS. Only f32 attention is affected; all other dtypes and ops are unchanged. Configs that previously passed had maxAbsDiff ~1e-4 (implied by the strict relDiff), so they remain passing.

Affected configs reproduced during investigation (all f32, transK=true, split_kv=1): g=8 seq_len=14400 head_dim=40 with perf_config attn:v3:32,96,192,4,32,48,16,4,1,1,1,0,1 on gfx942; g=8 seq_len=3600 head_dim=80 with perf_config attn:v3:64,64,48,4,16,48,16,4,1,2,1,2,1 on gfx942; and g=1 seq_len=14400 head_dim=512 with perf_config attn:v3:64,64,48,16,16,48,16,4,1,1,2,0,1 on gfx90a.

Test Plan

  • pytest mlir/utils/performance/tests/test_tuningRunner.py, including 4 new tests for attention_cpu_verify_flags (f32 attention gets the gate; f16/bf16/i8 and non-attention ops do not; gpu/none verify modes do not).
  • flake8 (with the CI ignore list) and yapf --diff on the changed files.
  • Manual verification pipeline (rocmlir-gen | rocmlir-driver -c | mlir-runner) on gfx942 (MI300X) for the failing configs, comparing the strict threshold against the new absDiff gate.
  • End-to-end tuningRunner.py --op attention on a previously failing f32 config, confirming the emitted verification command and the tuning outcome.

Test Result

  • pytest: 86 passed.
  • flake8: clean; yapf: no diff.
  • End-to-end tuningRunner on -t f32 ... -g 8 -seq_len_q 3600 -seq_len_k 3600 -head_dim_qk 80 -head_dim_v 80 (gfx942): Tuning completed successfully. The generated verification command included -absDiff_threshold 5e-2, and the winning perf_config attn:v3:64,64,48,4,16,48,16,4,1,2,1,2,1 (the exact config that fails in CI) passed, producing a valid tuning DB entry.
  • Manual pipeline: both f32 gfx942 configs (seq_len=14400 head_dim=40 perf_config attn:v3:32,96,192,4,32,48,16,4,1,1,1,0,1, and seq_len=3600 head_dim=80) return [1 1 1] with the gate; they return [1 1 0] with the previous strict threshold. The gfx90a config (head_dim=512) is not runnable on gfx942 hardware but is covered by the same mechanism (observed abs diff ~1.5e-2 < 5e-2).
  • PR CI
  • Weekly CI

Submission Checklist

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant