From bc2afd760c971e1cccf5eaf59edd592899266fd2 Mon Sep 17 00:00:00 2001 From: liyjiang Date: Mon, 28 Sep 2026 02:04:30 -0500 Subject: [PATCH 1/2] [MLA v4 nm] Fix _run_one_point reading packed BF16 as FP32 partials Since #4311 the v4 nm dispatcher derives out_16_nosplit from num_kv_splits and ignores the caller's value, so a single-split launch writes packed BF16 into the logits buffer. _run_one_point still read logits_buf[:, 0] as FP32 for num_kv_splits == 1, which compared every other BF16 element (plus the never-written tail of the buffer) against the reference and printed spurious `fp8_dequant_ref vs asm` checkAllclose failures in the script-mode sweep. Read output_buf, which holds the final result for every split count. Co-authored-by: Cursor --- op_tests/test_mla_v4_nm.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/op_tests/test_mla_v4_nm.py b/op_tests/test_mla_v4_nm.py index 3254de5e74a..1deeb2d8439 100644 --- a/op_tests/test_mla_v4_nm.py +++ b/op_tests/test_mla_v4_nm.py @@ -944,20 +944,13 @@ def _run_one_point( num_rotate_args=1, ) - # Resolve the asm output to compare against. Three cases, all reading the - # buffer the wrapper actually populated (the 2b call above): - # out_16_nosplit=1 -> kernel writes packed-BF16 into the logits region; - # the wrapper unpacks it into output_buf (see - # mla_decode_fwd_v4_nm). Read output_buf directly. - # single-pass (fp32) -> kernel writes one FP32 partial to logits[:, 0], - # no stage2; cast it to BF16. - # multi-pass -> stage2 merge wrote merged BF16 to output_buf. - if out_16_nosplit != 0: - out_asm = output_buf # wrapper unpacked packed-BF16 here - elif num_kv_splits == 1: - out_asm = logits_buf[:, 0].to(dtypes.bf16) # [total_q, num_heads, dv] - else: - out_asm = output_buf # already [total_q, num_heads, dv] BF16 + # output_buf is the authoritative result for every split count: single-pass + # writes packed BF16 straight into it, multi-pass gets it from the stage2 + # merge. logits_buf must not be read here -- the dispatcher derives + # out_16_nosplit from num_kv_splits (ignoring the caller's value), so for a + # single split the raw kernel call above fills logits_buf with packed BF16, + # not FP32 partials. + out_asm = output_buf # [total_q, num_heads, dv] BF16 # ---- accuracy ---- # Two comparisons, run for BOTH single- and multi-split (split-kv is a perf From 9979c04950f1db3c915b2b37de47e8f5b1ce6afd Mon Sep 17 00:00:00 2001 From: liyjiang Date: Mon, 28 Sep 2026 02:13:45 -0500 Subject: [PATCH 2/2] [MLA v4 nm] Gate accuracy checks so a failed! fails the test checkAllclose only raises on a catastrophic delta; otherwise it logs `failed!` and returns the mismatch fraction. Four of the six accuracy checks in test_mla_v4_nm.py dropped that return value, which is how the packed-BF16 readback bug fixed in the previous commit passed both pytest and the script-mode CI run. Route them through _gated_allclose, which asserts the mismatch fraction against the same tol_err_ratio checkAllclose uses for `failed!`. The script-mode sweep keeps going past a failing shape, prints a summary, and exits non-zero so aiter_test.sh reports the file as failed. Co-authored-by: Cursor --- op_tests/test_mla_v4_nm.py | 79 +++++++++++++++++++++++--------------- 1 file changed, 48 insertions(+), 31 deletions(-) diff --git a/op_tests/test_mla_v4_nm.py b/op_tests/test_mla_v4_nm.py index 1deeb2d8439..e8c5b759926 100644 --- a/op_tests/test_mla_v4_nm.py +++ b/op_tests/test_mla_v4_nm.py @@ -703,6 +703,22 @@ def _build_bf16_inputs( } +def _gated_allclose(a, b, msg, rtol=3e-2, atol=3e-2, tol_err_ratio=0.02): + """checkAllclose that fails the test on a `failed!` result. + + checkAllclose only raises on a catastrophic delta; otherwise it logs + `failed!` and returns the mismatch fraction, so an ungated call lets a bad + accuracy result pass silently in both pytest and the script-mode sweep. + """ + err = checkAllclose( + a, b, rtol=rtol, atol=atol, tol_err_ratio=tol_err_ratio, msg=msg + ) + assert err <= tol_err_ratio, ( + f"{msg}: {err:.1%} of elements outside rtol={rtol} atol={atol} " + f"(limit {tol_err_ratio:.0%})" + ) + + def _run_one_point( batch=2, kv_seq_lens=64, @@ -962,23 +978,17 @@ def _run_one_point( f"\n[v4 nm accuracy] batch={batch} kv_seq_lens={kv_seq_lens} " f"q_seq_logical={q_seq_logical} num_kv_splits={num_kv_splits} seed={seed}" ) - # Per-element check at checkAllclose's default 1% tolerance (rtol=atol=1e-2). - # checkAllclose prints pass/warning/failed with the offending-element ratio + - # max delta (it does not raise). - checkAllclose( + # Both are gated: golden vs fp8_ref catches a broken quant pipeline (e.g. + # all-zero e8m0 scales) that fp8_ref vs asm alone would miss, since the + # kernel and fp8_ref consume the same quantized bytes. + _gated_allclose( out_golden.float(), out_fp8_ref.float(), - rtol=3e-2, - atol=3e-2, - tol_err_ratio=0.02, msg="mla_v4_nm [golden_bf16 vs fp8_ref]", ) - checkAllclose( + _gated_allclose( out_fp8_ref.float(), out_asm.float(), - rtol=3e-2, - atol=3e-2, - tol_err_ratio=0.02, msg="mla_v4_nm [fp8_dequant_ref vs asm]", ) @@ -1104,12 +1114,9 @@ def _run_varlen_point(kv_lens, gqa_ratio=128, seed=0, attn_sink=True): f"\n[v4 nm varlen] gqa={gqa} kv_lens={kv_lens} total_kv={total_kv} " f"resolved_splits={resolved}" ) - checkAllclose( + _gated_allclose( out_ref.float(), out_asm, - rtol=3e-2, - atol=3e-2, - tol_err_ratio=0.02, msg=f"mla_v4_nm varlen [fp8_dequant_ref vs asm] kv_lens={kv_lens}", ) @@ -1488,12 +1495,9 @@ def _run_pad_poison_point(fill, poison_q, poison_kv, gqa_ratio, kv_seq_lens, bat f"kernel is reading past the {_QUANT_NUM_SCALE_BYTES}-byte scale field " f"at offset {_QUANT_D_NOPE} into the padding at {pad_off}." ) - checkAllclose( + _gated_allclose( out_ref.float(), out_asm, - rtol=3e-2, - atol=3e-2, - tol_err_ratio=0.02, msg=f"mla_v4_nm {who} pad=0x{fill:02X} [fp8_dequant_ref vs asm]", ) @@ -2549,6 +2553,7 @@ def test_v4_nm_kv_tail_not_tile_multiple(): args = parser.parse_args() perf_rows = [] + failures = [] for batch, kv_seq_lens, q_seq_logical in itertools.product( args.batch, args.kv_seq_lens, args.q_seq_logical ): @@ -2556,18 +2561,24 @@ def test_v4_nm_kv_tail_not_tile_multiple(): f"\n========== batch={batch} kv_seq_lens={kv_seq_lens} " f"q_seq_logical={q_seq_logical} ==========" ) - us_asm, us_ref = _run_one_point( - batch=batch, - kv_seq_lens=kv_seq_lens, - q_seq_logical=q_seq_logical, - seed=args.seed, - num_iters=args.iters, - num_warmup=args.warmup, - num_kv_splits=args.split_kv, - gqa_ratio=args.gqa_ratio, - attn_sink=args.attn_sink, - out_16_nosplit=args.out_16_nosplit, - ) + try: + us_asm, us_ref = _run_one_point( + batch=batch, + kv_seq_lens=kv_seq_lens, + q_seq_logical=q_seq_logical, + seed=args.seed, + num_iters=args.iters, + num_warmup=args.warmup, + num_kv_splits=args.split_kv, + gqa_ratio=args.gqa_ratio, + attn_sink=args.attn_sink, + out_16_nosplit=args.out_16_nosplit, + ) + except AssertionError as e: + # Keep sweeping so one bad shape doesn't hide the rest; the exit + # code below is what fails CI (aiter_test.sh runs this as a script). + failures.append(((batch, kv_seq_lens, q_seq_logical), str(e))) + continue perf_rows.append((batch, kv_seq_lens, q_seq_logical, us_asm, us_ref)) print("\n[v4 nm perf summary] (us; speedup = fp8_ref / asm_kernel)") @@ -2577,3 +2588,9 @@ def test_v4_nm_kv_tail_not_tile_multiple(): ) for b, k, q, ua, ur in perf_rows: print(f" {b:>6d} {k:>8d} {q:>6d} {ua:>10.2f} {ur:>12.2f} {ur / ua:>8.1f}x") + + if failures: + print(f"\n[v4 nm accuracy] {len(failures)} shape(s) FAILED:") + for (b, k, q), reason in failures: + print(f" batch={b} kv_seq_lens={k} q_seq_logical={q}: {reason}") + sys.exit(1)