mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 17:26:25 +08:00
ec5f37821d
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.
84 lines
2.9 KiB
Python
84 lines
2.9 KiB
Python
"""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
|