[KMCompiler][Test] Check an override against the reference, not the generic kernel

Author: cheersluvs

tests/test_fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert.py

@@ -641,9 +641,23 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
     k_cache_compare(k_cache, k_cache_ref, block_size, rtol=1e-2, atol=1e-2)
 
 
-# ── Test 5: a backend override must not change results ───────────────────────
-# A vendor override is a performance change. It runs only where one is
+# ── Test 5: a backend override must still satisfy the contract ───────────────
+# A vendor override is a performance change, so it is worth checking on shapes
+# the tests above do not cover -- in particular ones that are not multiples of
+# whatever tiling the override uses. It runs only where an override is
 # registered, so nothing below executes on the generic path.
+#
+# It compares against the torch reference, not against the generic kernel. An
+# earlier version asserted the two Triton implementations agreed byte for byte,
+# which is stricter than anything the operator promises and is not even stable:
+# on Hygon BW1000 they are bit-identical under torch 2.10.0 / FlagTree
+# 0.6.1a1+hcu3.6, while under the combination backends.yaml pins for that card
+# (torch 2.9.0 / FlagTree 0.5.1+hcu3.1) q differs on 5 of 16744448 elements by up
+# to 2 bf16 ULP at 511 tokens and one FP8 cache byte differs at 1000 tokens. A
+# 2-D tile assigns lanes differently than a 1-D block, so `variance` rounds
+# differently and how that lowers depends on the compiler. The contract each
+# implementation owes is agreement with the reference, at the tolerances the
+# tests above use, and that is what this asserts.
 _OVERRIDE_ACTIVE = (
     flag_gems.fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert is not _generic_impl
 )
@@ -661,7 +675,7 @@ def test_combined_q_and_kv(num_tokens: int, n_heads: int, block_size: int):
 )
 @pytest.mark.parametrize("n_heads", [64, 128])
 @pytest.mark.parametrize("block_size", [16, 64])
-def test_backend_override_matches_generic(
+def test_backend_override_matches_reference(
     num_tokens: int, n_heads: int, block_size: int
 ):
     torch.manual_seed(3)
@@ -679,32 +693,19 @@ def test_backend_override_matches_generic(
     positions = torch.arange(num_tokens, dtype=torch.int64, device=device)
     cos_sin_cache = make_cos_sin_cache(max_pos, ROPE_DIM, torch.float32, device)
 
-    q_gen, k_cache_gen = q.clone(), k_cache.clone()
-    _generic_impl(
-        q_gen, kv, k_cache_gen, slot_mapping, positions, cos_sin_cache, eps, block_size
+    q_ref, kv_ref = q.clone(), kv.clone()
+    k_cache_ref = k_cache.clone()
+    ref_impl_chunked(
+        q_ref,
+        kv_ref,
+        k_cache_ref,
+        slot_mapping.clone(),
+        positions.clone(),
+        cos_sin_cache.clone(),
+        eps,
+        block_size,
     )
     fused_impl(q, kv, k_cache, slot_mapping, positions, cos_sin_cache, eps, block_size)
 
-    # The cache is derived from kv through per-64-element reductions, identical
-    # in any sane decomposition, so it must match byte for byte -- an override
-    # that quantizes differently has changed behaviour.
-    assert torch.equal(k_cache, k_cache_gen), (
-        f"override wrote a different FP8 cache at num_tokens={num_tokens}, "
-        f"n_heads={n_heads}, block_size={block_size}: "
-        f"{int((k_cache != k_cache_gen).sum().item())} of {k_cache.numel()} "
-        "bytes differ"
-    )
-    # q goes through a 512-element RMSNorm reduction whose tree may legitimately
-    # differ between decompositions: a 2-D [TPP, 512] tile assigns lanes
-    # differently than a 1-D 512 block, so `variance` rounds differently, and
-    # that propagates through rsqrt into the stored bf16. Measured on a Hygon
-    # BW1000 at num_tokens=511: 5 of 16744448 elements differ, by up to 2 bf16
-    # ULP (7.8e-3 absolute, 7.8e-3 relative).
-    #
-    # The tolerance is the operator's own -- the same rtol/atol the accuracy
-    # tests above use against the torch reference. Asserting that an override
-    # tracks the generic kernel more closely than either tracks the reference
-    # would be asserting something the operator never promised. (It is in fact
-    # bit-identical on MetaX C550, but that is a property of one backend's lane
-    # assignment, not a contract.)
-    torch.testing.assert_close(q, q_gen, rtol=1e-2, atol=1e-2)
+    k_cache_compare(k_cache, k_cache_ref, block_size, rtol=1e-2, atol=1e-2)
+    assert_close_chunked(q, q_ref, rtol=1e-2, atol=1e-2)