fix(h20): let the reference backend sync, and stop the reference raising under the guard

Running end to end surfaced two things the per-kernel A/B could not.

The engine wraps the forward pass in `forbid_cuda_sync()`, which sets torch's sync debug mode to
"error". The compiled kernels never sync; the torch references do, and not by oversight:

* `tile_kernels/torch/cast.py` built `torch.tensor(-1.0, device="cuda")` (and three more like it)
  for use as a `torch.where` operand. Constructing a device tensor from a Python scalar is a
  host-to-device copy, so it syncs -- and it is unnecessary, because `torch.where` accepts a plain
  scalar and passes it to the kernel. Patched via `patch_torch_ref_sync.py` (idempotent,
  reversible), which is a fix worth having on its own: the same cast raises for any caller that
  holds the device clean.
* `get_mapping_from_psum` reads the expert prefix-sum to the host to compute alignment ranges in
  Python, and `moe_topk_gate_forward` compacts rows with `nonzero()` and branches on `.any()`.
  Those cannot be made sync-free without rewriting how the reference addresses, which is the wrong
  trade for a fallback whose job is being runnable. So each compat module wraps its reference
  module once and every call through it runs inside `allow_cuda_sync()`.

Not done, deliberately: `set_sync_debug_mode("default")` process-wide. That would also stop the
guard catching accidental syncs in the compiled path, which is what it is for.

Also fixed: my own `bool(force_random.any())` in the gate adapter was itself a sync, outside the
reference, so it raised before the adapter could decide anything.

End to end now completes: exit=0, all 8 DP ranks agree, coherent answers. It is not token-identical
to the compiled path -- 16/16 samples differ, with a control run confirming the harness is
deterministic (two cuda runs are bit-identical), so the difference is real. Recorded in the doc as
the next thing to attribute rather than papered over.

deepinfer/tk_compat/_sync.py

@@ -0,0 +1,83 @@
+"""Let the reference implementations sync, without turning the guard off everywhere.
+
+The engine wraps the forward pass in `forbid_cuda_sync()`, which sets torch's sync debug mode to
+"error"; one blocking copy or one `nonzero()` raises. The compiled kernels never sync, so the guard
+is a real property of the production path and worth keeping.
+
+The torch references cannot honour it, and not because of an oversight that could be patched away:
+
+* `get_mapping_from_psum` reads `psum_num_tokens_per_expert` to the host to compute per-expert
+  alignment ranges in Python.
+* `moe_topk_gate_forward` compacts valid rows with `active.nonzero()` and branches on
+  `normal_mask.any()`, both of which need the result on the host.
+
+Making those sync-free means rewriting how the reference does its addressing. That is the wrong
+trade: the reference exists to make the model runnable and checkable, and it is not the path whose
+performance matters.
+
+So each compat module wraps its reference module once:
+
+    _ref = reference_module(tile_kernels.torch)
+
+and every call through it runs inside `allow_cuda_sync()`. The wrapper is at the module boundary
+rather than in each branch, so no call site knows about it and a newly used reference entry point
+is covered automatically.
+
+Deliberately *not* done: `torch.cuda.set_sync_debug_mode("default")` for the whole process. That
+would also stop the guard from catching accidental syncs in the compiled path, which is exactly
+what it is for.
+"""
+
+from __future__ import annotations
+
+import functools
+from collections.abc import Callable
+from typing import Any
+
+try:  # The guard lives in deepinfer; a bare install of this package would not have it.
+    from deepinfer.utils import allow_cuda_sync
+
+    _HAVE_GUARD = True
+except Exception:  # pragma: no cover
+    _HAVE_GUARD = False
+
+
+def _wrap(fn: Callable[..., Any]) -> Callable[..., Any]:
+    if not _HAVE_GUARD:
+        return fn
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        with allow_cuda_sync():
+            return fn(*args, **kwargs)
+
+    return wrapper
+
+
+class _ReferenceProxy:
+    """Attribute access that grants sync permission for every callable it returns."""
+
+    __slots__ = ("_module",)
+
+    def __init__(self, module: Any) -> None:
+        object.__setattr__(self, "_module", module)
+
+    def __getattr__(self, name: str) -> Any:
+        value = getattr(object.__getattribute__(self, "_module"), name)
+        return _wrap(value) if callable(value) else value
+
+    def __repr__(self) -> str:
+        return f"reference_module({object.__getattribute__(self, '_module')!r})"
+
+
+def reference_module(module: Any) -> Any:
+    """Wrap a `tile_kernels.torch` submodule so its calls may sync."""
+    return _ReferenceProxy(module)
+
+
+if _HAVE_GUARD:
+    allow_sync = allow_cuda_sync
+else:  # pragma: no cover - no guard installed, nothing to lift
+    import contextlib
+
+    allow_sync = contextlib.nullcontext

deepinfer/tk_compat/engram.py

@@ -17,7 +17,12 @@ from deepinfer.tk_compat import use_references
 if not use_references():
     from tile_kernels import engram as _compiled
 
-from tile_kernels import torch as _ref
+from tile_kernels import torch as _ref_raw
+
+from deepinfer.tk_compat._sync import reference_module
+
+# Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design.
+_ref = reference_module(_ref_raw)
 
 
 def engram_hash(

deepinfer/tk_compat/mhc.py

@@ -32,7 +32,12 @@ from deepinfer.tk_compat import use_references
 if not use_references():
     from tile_kernels import mhc as _compiled
 
-from tile_kernels import torch as _ref
+from tile_kernels import torch as _ref_raw
+
+from deepinfer.tk_compat._sync import reference_module
+
+# Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design.
+_ref = reference_module(_ref_raw)
 
 
 def expand_to_mhc_fwd(x: torch.Tensor, o: torch.Tensor) -> None:

deepinfer/tk_compat/moe.py

@@ -17,7 +17,12 @@ from deepinfer.tk_compat import use_references
 if not use_references():
     from tile_kernels import moe as _compiled
 
-from tile_kernels import torch as _ref
+from tile_kernels import torch as _ref_raw
+
+from deepinfer.tk_compat._sync import reference_module
+
+# Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design.
+_ref = reference_module(_ref_raw)
 
 
 def moe_topk_gate(
@@ -60,12 +65,20 @@ def moe_topk_gate(
             unmapped_topk_idx=unmapped_topk_idx,
             out=out,
         )
-    if force_random is not None and bool(force_random.any()):
-        raise NotImplementedError(
-            "force_random has no torch reference: it replaces the gate's choice with a random "
-            "expert for load balancing, so a reference run cannot reproduce it. Clear "
-            "args.force_load_balance (its usual default) to use the reference path."
-        )
+    # `force_random` is a device tensor, so testing whether it asks for anything needs a sync.
+    # Do it under the guard's permission: the engine forbids syncs in forward, and this one is
+    # ours, not the reference's. It is a single reduced-boolean read per layer.
+    if force_random is not None:
+        from deepinfer.tk_compat._sync import allow_sync
+
+        with allow_sync():
+            forced = bool(force_random.any())
+        if forced:
+            raise NotImplementedError(
+                "force_random has no torch reference: it replaces the gate's choice with a random "
+                "expert for load balancing, so a reference run cannot reproduce it. Clear "
+                "args.force_load_balance (its usual default) to use the reference path."
+            )
     result = _ref.moe_topk_gate_forward(
         logits,
         num_topk=num_topk,

deepinfer/tk_compat/quant.py

@@ -29,7 +29,12 @@ from deepinfer.tk_compat import use_references
 if not use_references():
     from tile_kernels import quant as _compiled
 
-from tile_kernels import torch as _ref
+from tile_kernels import torch as _ref_raw
+
+from deepinfer.tk_compat._sync import reference_module
+
+# Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design.
+_ref = reference_module(_ref_raw)
 
 
 def per_token_cast(

deepinfer/tk_compat/transform.py

@@ -23,7 +23,12 @@ from deepinfer.tk_compat import use_references
 if not use_references():
     from tile_kernels import transform as _compiled
 
-from tile_kernels import torch as _ref
+from tile_kernels import torch as _ref_raw
+
+from deepinfer.tk_compat._sync import reference_module
+
+# Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design.
+_ref = reference_module(_ref_raw)
 
 
 def apply_rotary(

scripts/tile_kernels_patch/patch_torch_ref_sync.py

@@ -0,0 +1,104 @@
+"""Make `tile_kernels.torch.cast` free of CPU-GPU syncs.
+
+The engine guards the forward pass with `forbid_cuda_sync()`, which turns torch's sync debug mode
+to `"error"`; a single blocking copy raises `RuntimeError: called a synchronizing CUDA operation`.
+The compiled kernels honour it. The reference does not, in four places:
+
+    torch.where(mask, abs_f, torch.tensor(-1.0, device=device, dtype=...))
+    torch.where(cond,     torch.tensor(0.0,  device=device), ...)
+
+Building a tensor *on the device from a Python scalar* is a host-to-device copy, so it syncs. It is
+also unnecessary: `torch.where` takes a plain scalar as its value operand and passes it to the
+kernel as an argument, with no copy and no sync.
+
+This is a defect in the reference independent of my bridge -- any caller that holds the device
+clean across a cast hits it -- so the right fix is here rather than a guard lifted in the compat
+layer. Re-running is a no-op.
+
+    python patch_torch_ref_sync.py <tile_kernels-dir>
+    python patch_torch_ref_sync.py <tile_kernels-dir> --revert
+"""
+
+from __future__ import annotations
+
+import argparse
+import shutil
+import sys
+from pathlib import Path
+
+# Each replacement is (old, new, why). The scalar form is not a micro-optimisation: it is the
+# difference between raising and not raising under the engine's sync guard.
+REPLACEMENTS = [
+    (
+        "torch.where(reshaped_mask, abs_f, torch.tensor(-1.0, device=device, dtype=abs_f.dtype))",
+        "torch.where(reshaped_mask, abs_f, -1.0)",
+        "pad sentinel as a scalar operand",
+    ),
+    (
+        "torch.where(dequant_sf_rounded == 0, torch.tensor(0.0, device=device), 1.0 / dequant_sf_rounded)",
+        "torch.where(dequant_sf_rounded == 0, 0.0, 1.0 / dequant_sf_rounded)",
+        "zero guard as a scalar operand",
+    ),
+    (
+        "torch.where(ds_int_rounded == 0, torch.tensor(0.0, device=device), 1.0 / dequant_sf_e4m3)",
+        "torch.where(ds_int_rounded == 0, 0.0, 1.0 / dequant_sf_e4m3)",
+        "zero guard as a scalar operand (e4m3 sf)",
+    ),
+    (
+        "torch.where(ds_int_rounded == 0, torch.tensor(0.0, device=device), max_quant_val_expanded / max_val)",
+        "torch.where(ds_int_rounded == 0, 0.0, max_quant_val_expanded / max_val)",
+        "zero guard as a scalar operand (sqr variant)",
+    ),
+]
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("target", type=Path, help="the installed tile_kernels package directory")
+    parser.add_argument("--revert", action="store_true")
+    args = parser.parse_args()
+
+    path = args.target / "torch" / "cast.py"
+    if not path.is_file():
+        print(f"not found: {path}", file=sys.stderr)
+        return 1
+    backup = path.with_suffix(".py.orig")
+
+    if args.revert:
+        if not backup.is_file():
+            print(f"no backup for {path}", file=sys.stderr)
+            return 1
+        shutil.move(str(backup), str(path))
+        print(f"reverted {path.name}")
+        return 0
+
+    text = path.read_text()
+    applied = 0
+    for old, new, why in REPLACEMENTS:
+        if old in text:
+            text = text.replace(old, new)
+            print(f"applied: {why}")
+            applied += 1
+
+    if applied == 0:
+        # Either already patched, or the file changed upstream. Distinguish the two, because the
+        # first is fine and the second needs a human.
+        if "torch.tensor(-1.0, device=device" not in text and "torch.tensor(0.0, device=device)" not in text:
+            print("already patched; nothing to do")
+            return 0
+        print(
+            "the device-scalar pattern is still present but the exact lines changed; "
+            "update REPLACEMENTS against the installed version",
+            file=sys.stderr,
+        )
+        return 1
+
+    if not backup.is_file():
+        shutil.copy2(path, backup)
+    path.write_text(text)
+    print(f"\npatched {path}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())