From ec5f37821d7d173853d866ff2c3a24d510d1028b Mon Sep 17 00:00:00 2001 From: loopyt Date: Sat, 24 Oct 2026 19:10:49 +0800 Subject: [PATCH] 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 | 83 ++++++++++++++ deepinfer/tk_compat/engram.py | 7 +- deepinfer/tk_compat/mhc.py | 7 +- deepinfer/tk_compat/moe.py | 27 +++-- deepinfer/tk_compat/quant.py | 7 +- deepinfer/tk_compat/transform.py | 7 +- .../patch_torch_ref_sync.py | 104 ++++++++++++++++++ 7 files changed, 231 insertions(+), 11 deletions(-) create mode 100644 deepinfer/tk_compat/_sync.py create mode 100755 scripts/tile_kernels_patch/patch_torch_ref_sync.py diff --git a/deepinfer/tk_compat/_sync.py b/deepinfer/tk_compat/_sync.py new file mode 100644 index 0000000..75abe82 --- /dev/null +++ b/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 diff --git a/deepinfer/tk_compat/engram.py b/deepinfer/tk_compat/engram.py index c838b26..a770f0a 100644 --- a/deepinfer/tk_compat/engram.py +++ b/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( diff --git a/deepinfer/tk_compat/mhc.py b/deepinfer/tk_compat/mhc.py index 9efb9a8..150f373 100644 --- a/deepinfer/tk_compat/mhc.py +++ b/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: diff --git a/deepinfer/tk_compat/moe.py b/deepinfer/tk_compat/moe.py index 85df4b7..367ea94 100644 --- a/deepinfer/tk_compat/moe.py +++ b/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, diff --git a/deepinfer/tk_compat/quant.py b/deepinfer/tk_compat/quant.py index 9ed4048..2e06c43 100644 --- a/deepinfer/tk_compat/quant.py +++ b/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( diff --git a/deepinfer/tk_compat/transform.py b/deepinfer/tk_compat/transform.py index d2f1ed9..b4cc365 100644 --- a/deepinfer/tk_compat/transform.py +++ b/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( diff --git a/scripts/tile_kernels_patch/patch_torch_ref_sync.py b/scripts/tile_kernels_patch/patch_torch_ref_sync.py new file mode 100755 index 0000000..9d33b39 --- /dev/null +++ b/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 + python patch_torch_ref_sync.py --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())