mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-02 09:46:28 +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.
122 lines
4.4 KiB
Python
122 lines
4.4 KiB
Python
"""`tile_kernels.moe` -- expert routing.
|
|
|
|
Known gap, measured rather than assumed: the compiled `moe_topk_gate` accepts `force_random`
|
|
(driven by `args.force_load_balance`), and the torch reference has no such parameter at all. So
|
|
`ref` mode cannot reproduce forced load balancing. It raises instead of dropping the flag, because
|
|
dropping it would produce valid-looking but differently routed tokens.
|
|
|
|
`fix_routing_mask` *is* in the reference and is forwarded.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
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_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(
|
|
logits: torch.Tensor,
|
|
num_topk: int,
|
|
use_shared_as_routed: bool,
|
|
num_shared_experts: int,
|
|
routed_scaling_factor: float,
|
|
ep_rank: int,
|
|
scoring_func: str = "sqrtsoftplus",
|
|
mask: torch.Tensor | None = None,
|
|
bias: torch.Tensor | None = None,
|
|
image_bias: torch.Tensor | None = None,
|
|
image_token_mask: torch.Tensor | None = None,
|
|
fix_routing_mask: torch.Tensor | None = None,
|
|
to_physical_map: torch.Tensor | None = None,
|
|
logical_count: torch.Tensor | None = None,
|
|
force_random: torch.Tensor | None = None,
|
|
unmapped_topk_idx: torch.Tensor | None = None,
|
|
out: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Top-k expert routing with sqrtsoftplus scoring."""
|
|
if not use_references():
|
|
return _compiled.moe_topk_gate(
|
|
logits,
|
|
num_topk=num_topk,
|
|
use_shared_as_routed=use_shared_as_routed,
|
|
num_shared_experts=num_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
ep_rank=ep_rank,
|
|
scoring_func=scoring_func,
|
|
mask=mask,
|
|
bias=bias,
|
|
image_bias=image_bias,
|
|
image_token_mask=image_token_mask,
|
|
fix_routing_mask=fix_routing_mask,
|
|
to_physical_map=to_physical_map,
|
|
logical_count=logical_count,
|
|
force_random=force_random,
|
|
unmapped_topk_idx=unmapped_topk_idx,
|
|
out=out,
|
|
)
|
|
# `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,
|
|
use_shared_as_routed=use_shared_as_routed,
|
|
num_shared_experts=num_shared_experts,
|
|
routed_scaling_factor=routed_scaling_factor,
|
|
ep_rank=ep_rank,
|
|
scoring_func=scoring_func,
|
|
mask=mask,
|
|
bias=bias,
|
|
image_bias=image_bias,
|
|
image_token_mask=image_token_mask,
|
|
fix_routing_mask=fix_routing_mask,
|
|
to_physical_map=to_physical_map,
|
|
logical_count=logical_count,
|
|
unmapped_topk_idx=unmapped_topk_idx,
|
|
out=out,
|
|
)
|
|
return result
|
|
|
|
|
|
def normalize_weight(
|
|
topk_weights: torch.Tensor,
|
|
psum_num_tokens_per_expert: torch.Tensor | None = None,
|
|
topk_ids: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Normalise routing weights. Same signature on both paths."""
|
|
if not use_references():
|
|
return _compiled.normalize_weight(topk_weights, psum_num_tokens_per_expert, topk_ids)
|
|
return _ref.normalize_weight(topk_weights, psum_num_tokens_per_expert, topk_ids)
|
|
|
|
|
|
def mask_indices_by_tp(x: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor:
|
|
if not use_references():
|
|
return _compiled.mask_indices_by_tp(x, *args, **kwargs)
|
|
return _ref.mask_indices_by_tp(x, *args, **kwargs)
|
|
|
|
|
|
def backend() -> str:
|
|
return "ref" if use_references() else "cuda"
|