mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +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.
105 lines
3.9 KiB
Python
Executable File
105 lines
3.9 KiB
Python
Executable File
"""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())
|