mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-02 17:56:22 +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.
119 lines
4.1 KiB
Python
119 lines
4.1 KiB
Python
"""`tile_kernels.engram` -- n-gram hash embedding indices.
|
|
|
|
Signature is the same on both paths, and the compiled `engram_hash` was measured bitwise
|
|
identical to `engram_hash_ref`. The dtypes are not obvious and differ per argument, so they are
|
|
worth stating: ids int32; multipliers int64 shaped ``(layers, max_ngram_size)``; vocab_sizes int32
|
|
shaped ``(layers, max_ngram_size - 1, tables)``; offsets int32.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
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_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(
|
|
ngram_token_ids: torch.Tensor,
|
|
multipliers: torch.Tensor,
|
|
vocab_sizes: torch.Tensor,
|
|
offsets: torch.Tensor,
|
|
image_token_mask: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""Hash n-grams to embedding-pool indices."""
|
|
if not use_references():
|
|
return _compiled.engram_hash(ngram_token_ids, multipliers, vocab_sizes, offsets, image_token_mask)
|
|
return _ref.engram_hash_ref(ngram_token_ids, multipliers, vocab_sizes, offsets, image_token_mask)
|
|
|
|
|
|
def fused_weight(weight_hidden: torch.Tensor, weight_embed: torch.Tensor) -> Any:
|
|
"""Pre-compute the product of the two RMSNorm weights.
|
|
|
|
On the compiled path this is a fused `weight_hidden * weight_embed` in float32, and the gate
|
|
consumes only the product.
|
|
|
|
On the reference path the product is not enough: `engram_gate_ref` applies the two weights as
|
|
*separate* RMSNorm scales, one to the hidden states and one to the key, so recovering them from
|
|
their product is impossible. So this returns the pair and `engram_gate_fwd` unpacks it. The
|
|
return type therefore differs between the two paths -- a tuple instead of a tensor -- which is
|
|
safe because the call site only stores and forwards it, never inspects it.
|
|
"""
|
|
if not use_references():
|
|
return _compiled.fused_weight(weight_hidden, weight_embed)
|
|
return (weight_hidden, weight_embed)
|
|
|
|
|
|
def engram_gate_fwd(
|
|
hidden_states: torch.Tensor,
|
|
kv: torch.Tensor,
|
|
weight_fused: Any,
|
|
eps: float,
|
|
clamp_value: float,
|
|
save_for_backward: bool = True,
|
|
image_token_mask: torch.Tensor | None = None,
|
|
inplace: bool = False,
|
|
) -> tuple[torch.Tensor, ...]:
|
|
"""Engram gate. Returns a tuple whose first element is the output, like the compiled kernel."""
|
|
if not use_references():
|
|
return _compiled.engram_gate_fwd(
|
|
hidden_states,
|
|
kv,
|
|
weight_fused,
|
|
eps,
|
|
clamp_value,
|
|
save_for_backward,
|
|
image_token_mask,
|
|
inplace,
|
|
)
|
|
|
|
# `fused_weight` handed back the pair on this path; see its docstring.
|
|
if not isinstance(weight_fused, tuple):
|
|
raise TypeError(
|
|
"ref mode expects the (weight_hidden, weight_embed) pair from fused_weight, got "
|
|
f"{type(weight_fused).__name__}. Mixing a compiled-mode product with the reference "
|
|
"gate cannot work: the reference needs the two weights separately."
|
|
)
|
|
weight_hidden, weight_embed = weight_fused
|
|
|
|
# Argument order differs: the compiled gate takes (eps, clamp_value), the reference takes
|
|
# (clamp_value, eps).
|
|
result = _ref.engram_gate_ref(
|
|
hidden_states,
|
|
kv,
|
|
weight_hidden,
|
|
weight_embed,
|
|
clamp_value,
|
|
eps,
|
|
save_for_backward=save_for_backward,
|
|
image_token_mask=image_token_mask,
|
|
)
|
|
if save_for_backward:
|
|
output = result[0]
|
|
rest = tuple(result[1:])
|
|
else:
|
|
output = result
|
|
rest = ()
|
|
|
|
if inplace:
|
|
# The compiled kernel writes the result over `hidden_states`; call sites pass inplace=True
|
|
# and rely on the returned tensor, but keep the aliasing so behaviour matches.
|
|
hidden_states.copy_(output)
|
|
output = hidden_states
|
|
return (output, *rest)
|
|
|
|
|
|
def backend() -> str:
|
|
return "ref" if use_references() else "cuda"
|