test: unified GEMM/BMM fuzzer + convention auditor - #3539
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds a unified GEMM/BMM fuzzing and conformance test suite (deterministic adapters, quantization oracles, NaN-poison checks, autotune-cache handling), plus cuDNN backend gating and stricter bmm_mxfp8 size validation. ChangesUnified GEMM/BMM Fuzzing Framework
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces three comprehensive randomized fuzzing test suites targeting FP8 batched GEMM, low-precision FP4 GEMM/quantization, and fused MoE implementations to catch edge-case bugs, non-determinism, and device state corruption. The review feedback identifies three key issues: the use of Python's non-deterministic 'hash()' function for seeding in the FP4 fuzzer, an inefficient generator instantiation inside a list comprehension in the FP8 fuzzer that defeats seed randomization, and control flow issues in the MoE fuzzer that prematurely skip valuable determinism/device-state checks and silently ignore shape mismatches.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
YangXu1990uiuc
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/gemm/test_unified_gemm_fuzz.py (1)
769-791: 💤 Low valuePotential issue:
cfg.adapterpoints toads[0]butad.runuses a different adapter.The
cfgobject is created withadapter=ads[0], but inside the loop,ad.run(a, b, out, be, cfg)is called with potentially different adapters fromads. If any adapter'srunfunction usescfg.adapterinternally (e.g., to check properties likequant_modeorop_shape), it would get the wrong adapter's values.Looking at the run functions (
_run_bf16_mm,_run_fp8_bmm, etc.), they usecfg.fp8_idt,cfg.fp8_mdt, andcfg.use_8x4but notcfg.adapterdirectly. However, since all adapters in the same group share the same(op_shape, quant_mode), and the_canonicalfunction usescfg.adapter.quant_modeandcfg.adapter.op_shape, this could cause issues if adapters in the same convention group have differentquant_modevalues.In practice, by construction (grouping by quant_mode), this should be fine, but the code would be clearer if
cfgwere recreated per adapter.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/gemm/test_unified_gemm_fuzz.py` around lines 769 - 791, The cfg is built once with adapter=ads[0] but reused for other adapters, risking stale adapter-specific fields; update cfg.adapter for each adapter (or recreate cfg inside the loop) before calling ad.run(a, b, out, be, cfg) so that adapter-dependent values (cfg.adapter, cfg.quant_mode, cfg.op_shape, cfg.fp8_idt/fp8_mdt/use_8x4 used by _run_bf16_mm/_run_fp8_bmm and related runners) reflect the current ad; ensure _canonical is called with the correct cfg.adapter when constructing a,b,_ref for each ad if needed.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/gemm/test_unified_gemm_fuzz.py`:
- Around line 769-791: The cfg is built once with adapter=ads[0] but reused for
other adapters, risking stale adapter-specific fields; update cfg.adapter for
each adapter (or recreate cfg inside the loop) before calling ad.run(a, b, out,
be, cfg) so that adapter-dependent values (cfg.adapter, cfg.quant_mode,
cfg.op_shape, cfg.fp8_idt/fp8_mdt/use_8x4 used by _run_bf16_mm/_run_fp8_bmm and
related runners) reflect the current ad; ensure _canonical is called with the
correct cfg.adapter when constructing a,b,_ref for each ad if needed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 31def3b6-f837-4465-9a09-660f2c9588f7
📥 Commits
Reviewing files that changed from the base of the PR and between 0ba7a3f and 208e3b073a123b65a53eb1eeed56a75a47ae140c.
📒 Files selected for processing (1)
tests/gemm/test_unified_gemm_fuzz.py
|
@dhiraj113 @bkryu wanna take a look here? |
|
/bot run tests/gemm |
|
[FAILED] Pipeline #57269339: 7/20 passed |
…to latest main) Squash of the 24-commit PR flashinfer-ai#3539 branch onto current main (162 commits of drift; the intermediate MoE-fuzzer commits conflicted with the flashinfer-ai#3093 unified MoE fuzzer that landed on main independently, so history is collapsed to the net diff: tests/gemm/test_unified_gemm_fuzz.py + the two gemm_base.py guards). AI-assisted (squash-rebase by Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
9343f72 to
1df17e3
Compare
…t run Pipeline 57269339 (7/20): every red unit-test leg reduced to one of two causes, both in the fuzzer itself (the multi_gpu_test_b300[cu129] red is an unrelated pre-existing tests/comm/test_allreduce_unified_api.py trtllm-fusion failure): 1. all-zero-oracle FALSE POSITIVE (9 legs, same config+seed mm_nvfp4_cudnn_m7_n32_k32_float16_s65667983): with sparse inputs and a tiny 224-elem output the ORACLE itself is 223/224 zero and the output matched it, but the standalone '~all-zero output' invariant only looked at the oracle's max magnitude. The invariant is REMOVED as redundant: an all-zero/all-NaN output necessarily fails the tight numeric check, and the failure dump (now the worst 100 elements + stats) makes the pattern self-evident. Verified: the exact CI seed passes on SM100 with the same ratio (0.01355) CI logged. 2. ledgered configs could still CRASH (gb200[cu129]: 966 cascade failures): the flashinfer-ai#3604 bmm_mxfp8 b>1/M%128!=0 entry xfailed only at compare time, AFTER running the kernel; on the cu129 stack the same root cause is an illegal memory access (b=16 m=7 n=512 k=2688) that poisons the CUDA context for every later test. Ledger entries now carry a crash-capable flag: crash entries xfail UP FRONT (never launched) in the fuzz test and are skipped in the autotune-dynshape M-sequence; numeric-only entries still run and keep the xpass 'fixed -> remove me' signal. Verified: the exact CI IMA config now xfails in <1s without touching the device. Also found during re-validation (new numeric-only ledger entry): on cuDNN 9.23.0 + SM90 the AUTOTUNED bf16->bf16 tactic is garbage -- per-plan enumeration shows exactly the five eng7_k17=4_* plans (engine 7, CUDNN_KNOB_TYPE_SPLIT_K_SLC=4) miscompute (ratio ~1.39) while eng7 without split-k is correct; the tuner picks the broken one because split-k wins the timing race on tall-K shapes (mm_bf16 m63 n32 k2688). NOT tactic-index drift: profiled and executed against the same graph in-process. Verified fixed in 9.23.1/9.23.2 (same per-plan matrix all-correct), so the ledger gate is exactly ==92300. bf16-out sibling of the fp16-out 9.23.0 bug gemm_base.py hard-bans; the default tactic is correct so no product change. Plus the long_running marker per flashinfer-ai#3770 (merged into the existing pytestmark list -- a second bare assignment silently overwrote the first). AI-assisted (CI log triage + per-plan bisect by Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1df17e3 to
b4ec56f
Compare
…nobs - tests/test_helpers/fuzz_ledger.py: one mechanism for all quality fuzzers (gh #3605): quarantine=False entries run with tolerated wrong answers and flag xpass loudly; quarantine=True entries xfail up front (crash class). Every entry must reference a tracking issue (validated at construction). Same shape the scaled-GEMM fuzzer (PR #3539) evolved independently. - Migrate the unified MoE fuzzer to the shared ledger. - Debug knobs for gh #3957 bisection: FLASHINFER_UMOE_FUZZ_BACKENDS (backend-scoped sequences) and FLASHINFER_UMOE_FUZZ_NO_AUTOTUNE. - No quarantine entry for #3957: it is cumulative cross-call state corruption with a moving victim (config-predicate quarantine tried and refuted on hardware); the file staying red on SM100 is the signal. AI-assisted (bisect + validation on live SM100 hardware). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
/bot run tests/gemm |
|
Rebased onto latest main + both bot-run failure classes fixed. History is squashed to the net diff (the intermediate MoE-fuzzer commits conflicted with the unified MoE fuzzer that landed independently via #3093): one base commit + one fix commit. Triage of pipeline 57269339 (7/20)Every red unit-test leg was this fuzzer file itself (all 19 other 1. Nine legs, all the same single failure 2. The remaining red, Bonus finding from re-validationOn a cuDNN 9.23.0 box, the autotune-ON winner validation caught the autotuner-selected cuDNN tactic returning garbage for bf16→bf16 on SM90 ( Re-validation (all four local archs, rebased code)
Zero failures. The rebase also newly exercises the Internal CI has been re-triggered on the new head. Note for reading the result: |
…out (breadth -> unified fuzzers) GEMM (fuzzer default-on since #3539): bmm_fp8 3456->6 cases (~71 min/leg), mm_fp4 23760->12 + 4 auto (~59 min on Blackwell legs), mm_bf16 7560->13, mm_mxfp8 2688->6 (+large-dim 400->4, stats 18->3), bmm_mxfp8 576->5, bmm_bf16 240->8. Every kept case verified to actually run somewhere (none self-skip into no-ops). Kept in full: error-path / cache-behavior / invariant tests, the #3560 ragged-K anchor, and non-fuzzed paths (8x4 SF layout, bias/pdl epilogues, trtllm weight shuffle). MoE (rides the #3958 fuzzer default-on flip, assumes #3892 routing axes): shape fan-out compressed, the quant x routing x weight-layout kernel-selection matrix kept in FULL -- non-NVFP4 quant numerics are not fuzzed yet, so only shapes were cut, never modes: - renormalize trio: shared axis constants -> boundary tokens [8,3072] + intermediate [1024,384] (trio 4608 -> ~1500 collected) - sigmoid 576->192; deepseekv3 27216->4032 (intermediate axis = hitting set of every routing config's compatible_intermediate_size, each config still runs) - routed-parity 3456->144 (routing x quant x packed/unpacked in full); per-token 4over6 108->16; fp4 tactic sweep 72->24 (odd-token anchor kept); cute-dsl accuracy 144->16 and 24->8 Untouched: error-path / OA-param / routing-replay / kernel-tier / LoRA / CUDA-graph / EP / tactic-pairing tests; #3595 / #3067 anchors. Every touched file carries a header directing future coverage to the fuzzers. AI-assisted (grid analysis + curation by Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
3 participants
test: unified GEMM/BMM fuzzer + convention auditor
What this is
One harness —
tests/gemm/test_unified_gemm_fuzz.py— for flashinfer's scaled GEMM/BMM family:{op-shape: mm, bmm} × {quant: bf16, fp8 (e4m3+e5m2), nvfp4, mxfp4} × {backend}, driven through thinper-API adapters. It replaces the earlier per-op fuzzers (
test_mm_fp4_fuzz.py,test_bmm_fp8_fuzz.py) — their logic, plus the#2440quantize-root test, is folded in here — sothere is one strong, encapsulated tester instead of N scattered files.
Input model (same as the unified MoE fuzzer)
Sparse (~75% zero) + exactly-representable inputs, snapped to each quant mode's grid via a
per-mode round-trip (bf16 = identity, fp8 =
to_float8, nvfp4 =nvfp4_quantize+e2m1decode,mxfp4 =
mxfp4_quantize+mxfp4_dequantize). Because input quantization is then lossless:short sparse reductions instead of being averaged away;
atol = C·‖ref‖∞against the authoritative snapped-inputreference (C = the accumulation/requant floor), not a loose
cosine > 0.97. This catches bothstructural bugs and the sub-floor accuracy regressions a cosine oracle misses.
Magnitude regimes (tiny/large/…) are kept only in the standalone quantize-root test, where extreme
magnitudes are the point (
#2440: finite inputs must never yield non-finite scale factors).Oracles (per config)
no-spurious-NaN/Inf · tight numeric vs authoritative reference · not-(almost)-all-zero (
#3398/#3068)· output-buffer poison (NaN-fill → catches a kernel that doesn't fully write its output) ·
run-to-run determinism (
#2514) · device-state probe (a context-corrupting IMA → cleanfailure) · cross-arch by construction (run the same seed on each GPU, diff pass/fail).
Convention auditor (existing APIs unchanged)
Each scaled-GEMM API ships its own scale convention (per-tensor alpha vs block vs block+global alpha
vs none; A/B-scale roles; layout) — the surface where the fp4-vs-fp8 incompatibility was found. The
APIs cannot be changed (would break users), so this harness does not force them to agree:
never against a different convention (forcing cross-convention equality is a false-positive trap);
test_convention_conformancecross-checks backends that share a declared convention (a realcross-backend oracle), prints a convention matrix, and a
_CONVENTION_DIVERGENCESledger documentsknown cross-mode incompatibilities so they're tracked, not silently passing.
This is the enforcement hook for the future: if a unified GEMM API (or an incremental
convention-compat fix) makes two divergent APIs share a convention, move them into one conformance
group + drop the ledger entry → the test then enforces they agree.
Debuggability (every test)
from its seed → bit-reproducible.
stats (nan/inf/zero counts,
max|.|) + the worst ≤30 elements — so a CI log shows whether theoutput is all-zero / NaN / Inf / garbage without rerunning.
REPRO:line;FLASHINFER_GEMM_FUZZ_ONLY_SEED=<seed>reruns exactlythat one config.
Validation
Run on all four archs on the dev box — A100/SM80, L40S/SM89, H100/SM90, B200/SM100 — all clean
(fp4 adapters skip cleanly below SM100; fp8 below SM89). Tolerances calibrated on SM100 and verified
to hold cross-arch (fp8 ≤ 0.011, bf16 ≤ 0.0035). Default
FLASHINFER_GEMM_FUZZ_NUM_TESTS=1000(~10-min full sweep; tunable via env). The new dump immediately surfaced — and we fixed — a harness
B-layout bug (a
.contiguous()made fp8Brow-major, which cublas silently computes garbage fromwhile cudnn/cutlass reject it).
Follow-ups (tracked in-file TODOs)
autotune ON/OFF + cache-coherence oracle · non-contiguous-input axis (B1) ·
use_8x4_sf_layout=True/
#2861(C1, needs the matching SF layout + a trtllm backend) · grouped op-shape(
group_*/*deepgemm*, m_indptr) · more point APIs (mm_fp8low-latency,mm_mxfp8,bmm_mxfp8).Summary by CodeRabbit