Compare the module trees, not two compilations of them

tests/vllm/_granite4_fullsize_tests.py::test_logits_match[4.0-micro] fails on
the vllm20 CI leg and passes on vllm19, and has done so on unrelated PRs -- 117
(2f919be) is a docs-and-comments-only commit, which cannot move a logit. So the
assertion is measuring something other than the code under test.

What it measures is compilation. The test asserts atol=rtol=0 between two
separately launched vLLM engines: upstream GraniteMoeHybrid and GraniteSwitch
with num_adapters=0. Those are different graphs, so torch.compile and inductor
are free to fuse and tile them differently, and at production dimensions that
is enough to round one bf16 logit's last bit apart. vLLM logs its own per-engine
kernel choice: the failing engine reported IrOpPriorityConfig(rms_norm=
['native']) where its sibling reported ['vllm_c', 'native'] -- a different
RMSNorm kernel in 80 places.

The attribution rests on what does pass. 4.0-1b and 4.0-350m have the same 40
dense layers and are bit-exact on both legs; the same comparison on the HF
backend is bit-exact for all three; and real granite-4.0-micro weights are
bit-exact through vLLM in test_skinning_equivalence.py. A weight-transfer or
skinning bug would not spare all of those.

So run these eager. Bit-exactness is a property of the module tree, and eager is
where the module tree is what executes -- enforce_eager takes the compiler, the
CUDA-graph capture and vLLM's on-disk compile cache out of the comparison and
leaves the gate at atol=rtol=0 rather than loosening it. There was already
precedent for the flag in _granite4_mini_tests.py's _UPSTREAM_EAGER_CONFIGS. The
compiled path keeps the coverage it already had, from the two tests that are
deliberately tolerant of a flipped last bit: test_skinning_equivalence.py on
real weights and test_generation_equivalence.py on distributions.

Two supporting changes, because the failure was also unreadable:

  * The bit-exact branch now calls this repo's assert_close instead of
    torch.testing.assert_close. A msg= string handed to the latter *replaces*
    its whole diagnostic, which is why four CI runs recorded "logprobs should be
    bit-exact" and not one number -- there was no way to tell a last-bit drift
    from a real regression without re-running 4 GPU-hours. assert_close reports
    the count, the worst diff and its magnitude.

  * assert_close now treats matching non-finite entries as equal. Logprob
    tensors are padded with -inf for positions vLLM did not return, and
    -inf - -inf is nan, which fails every comparison including diff <= 0.0. That
    trap was latent while only the tolerance branch used the helper; the
    bit-exact branch would have hit it as a bogus worst=nan.

The margin is printed on a pass too, so a gate that has quietly started riding
at one ULP is visible before it trips.

Verified on CPU: 35 tests in tests/hf/test_granite4_mini.py and
test_granitemoe_mini.py (the existing assert_close callers) still pass, the three
TestGranite4FullSizeWeightTransfer cases pass, and the five assert_close
semantics above were checked directly. The GPU legs are the real gate.

Signed-off-by: itzikvaisman <yizhak.vaisman@ibm.com>

Author: ItzikVa

tests/shared/granite4_equivalence.py

@@ -31,11 +31,22 @@ def assert_close(actual, expected, *, atol, rtol, msg=""):
 
     This avoids the joint formula (atol + rtol * |b|) where the two
     tolerances inflate each other's budget.
+
+    ``atol=rtol=0.0`` is a valid bit-exact gate. Passing it here rather than
+    to ``torch.testing.assert_close`` is deliberate: that function's ``msg=``
+    string *replaces* its whole diagnostic, so a failure reports only the
+    message and none of the numbers needed to tell a last-bit drift from a
+    real regression.
+
+    Matching non-finite entries count as equal. Logprob tensors are padded
+    with ``-inf`` for positions vLLM did not return, and ``-inf - -inf`` is
+    ``nan``, which fails every comparison — including ``diff <= 0.0`` — and
+    would surface as a bogus ``worst=nan``.
     """
     diff = (actual - expected).abs()
     abs_ok = diff <= atol
     rel_ok = diff <= rtol * expected.abs()
-    ok = abs_ok | rel_ok
+    ok = abs_ok | rel_ok | (actual == expected)
 
     if not ok.all():
         num_bad = (~ok).sum().item()
@@ -256,6 +267,21 @@ def get_tolerances(layer_types, long_sequence=False, has_kv_hidden=False):
        position while upstream embeds the original control id. Visible
        positions attending to the control position pick up that delta.
 
+    Case 1's bit-exactness holds for the *eager* graph only. Under
+    torch.compile the switch graph and the upstream graph are different
+    graphs, so inductor is free to fuse and tile them differently; at
+    production dimensions that is enough to round one bf16 logit's last bit
+    apart. It surfaced as ``test_logits_match[4.0-micro]`` failing on vLLM
+    0.20 while passing on 0.19 — with 4.0-1b and 4.0-350m (same 40 layers,
+    same dense stack) passing on both, and the same test bit-exact on the HF
+    backend, which is why the divergence is attributed to compilation rather
+    than to the weight transfer. vLLM's per-engine kernel choice is visible
+    in its own logs: the failing engine reported
+    ``IrOpPriorityConfig(rms_norm=['native'])`` where its sibling reported
+    ``['vllm_c', 'native']``. The vLLM callers that want a bit-exact gate
+    therefore pass ``enforce_eager=True``; see
+    ``tests/vllm/_granite4_fullsize_tests.py``.
+
     Args:
         layer_types: list of "attention" strings
         long_sequence: unused (kept for API compatibility)

tests/vllm/_granite4_fullsize_tests.py

@@ -3,6 +3,10 @@
 
 Full-size Granite 4 equivalence tests via vllm.LLM.
 Requires CUDA GPU and vLLM installed.
+
+Runs eager. These are the bit-exact tests, and bit-exactness is a property of
+the module tree, not of whatever inductor decided to fuse in each of the two
+graphs — see the comment on ``enforce_eager`` below.
 """
 
 import pytest
@@ -53,22 +57,39 @@ def test_logits_match(self, model_name, tmp_path):
             tmpdir=tmp_path,
             max_model_len=64,
             gpu_memory_utilization=0.4,
+            # Compare the two module trees, not two compilations of them.
+            # torch.compile/inductor pick fusions and kernels per graph, and
+            # the switch graph is not the upstream graph, so at production
+            # dimensions the compiled paths can land a bf16 last bit
+            # differently while the eager math is identical. That is what
+            # broke 4.0-micro under vLLM 0.20 and not 0.19 (see the note in
+            # get_tolerances). enforce_eager removes the compiler, the
+            # CUDA-graph capture and vLLM's on-disk compile cache from the
+            # comparison, which restores a meaningful bit-exact gate.
+            # The compiled path is still covered, by the tests that are
+            # tolerant of a flipped last bit:
+            # tests/composer/test_skinning_equivalence.py (real weights) and
+            # tests/vllm/test_generation_equivalence.py (distribution).
+            enforce_eager=True,
         )
 
         tol = get_tolerances(layer_types)
-        if tol is None:
-            torch.testing.assert_close(
-                switch,
-                upstream,
-                atol=0.0,
-                rtol=0.0,
-                msg=f"{model_name}: logprobs should be bit-exact",
-            )
-        else:
-            assert_close(
-                switch,
-                upstream,
-                atol=tol[0],
-                rtol=tol[1],
-                msg=f"{model_name}: full-size logprobs diverge",
-            )
+        atol, rtol = (0.0, 0.0) if tol is None else tol
+        assert_close(
+            switch,
+            upstream,
+            atol=atol,
+            rtol=rtol,
+            msg=f"{model_name}: full-size logprobs diverge",
+        )
+
+        # Print the margin on a pass, not only into the failure message. A
+        # bit-exact gate that has quietly started riding at one ULP is a
+        # regression in progress, and this is the only place it shows.
+        delta = (switch - upstream).abs()
+        finite = delta[delta.isfinite()]
+        print(
+            f"{model_name}: {finite.numel()} finite logprob entries, "
+            f"max |delta| = {finite.max().item():.3e} "
+            f"(atol={atol:.1e} rtol={rtol:.1e})"
+        )