mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 09:16:20 +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.
103 lines
3.5 KiB
Python
103 lines
3.5 KiB
Python
"""`tile_kernels.transform` -- rotary position embedding.
|
|
|
|
This is the widest gap between the two signatures, and the one worth reading carefully:
|
|
|
|
* the compiled kernel takes `cos_sin_cache`, a single tensor of ``cat([cos, sin], -1)``, and
|
|
**rotates in place**, returning ``None``;
|
|
* the reference takes `cos` and `sin` as separate tensors and **returns a new tensor**.
|
|
|
|
So the adapter has to split the cache, call the reference, and write the result back into the
|
|
caller's tensor -- because the call site relies on the in-place contract
|
|
(`rotary_embedding.py` documents it and discards the return value).
|
|
|
|
Measured against the compiled kernel: identical except one bf16 last-bit tie per 65536 elements,
|
|
which is rounding, not a logic difference.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
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_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(
|
|
query: torch.Tensor,
|
|
cos_sin_cache: torch.Tensor,
|
|
key: torch.Tensor | None = None,
|
|
positions: torch.Tensor | None = None,
|
|
interleaved: bool = False,
|
|
conjugate: bool = False,
|
|
seqlen_offset: int = 0,
|
|
) -> None:
|
|
"""Rotate `query` (and `key`) in place. Returns None, like the compiled kernel."""
|
|
if not use_references():
|
|
return _compiled.apply_rotary(
|
|
query,
|
|
cos_sin_cache,
|
|
key,
|
|
positions,
|
|
interleaved,
|
|
conjugate,
|
|
seqlen_offset,
|
|
)
|
|
|
|
# The cache stores cos and sin halves side by side; rot_dim is their total.
|
|
rot_dim = cos_sin_cache.shape[-1]
|
|
if rot_dim % 2:
|
|
raise ValueError(f"cos_sin_cache last dim must be even, got {rot_dim}")
|
|
half = rot_dim // 2
|
|
cos = cos_sin_cache[..., :half]
|
|
sin = cos_sin_cache[..., half:]
|
|
|
|
def rotate(x: torch.Tensor) -> None:
|
|
# The reference wants (seqlen, heads, rot_dim) or 4-D with a batch axis; the caller passes
|
|
# (num_tokens, num_heads, rot_dim). Cast for the reshape and back after.
|
|
original_shape = x.shape
|
|
if x.dim() == 3:
|
|
seqlen, heads, dim = original_shape
|
|
view = x.view(seqlen, heads, dim)
|
|
elif x.dim() == 4:
|
|
batch, seqlen, heads, dim = original_shape
|
|
view = x.view(batch * seqlen, heads, dim)
|
|
else:
|
|
raise ValueError(f"apply_rotary expects 3-D or 4-D input, got {x.dim()}-D")
|
|
|
|
# `positions` is per-token; a flattened batch axis would need them tiled, which the
|
|
# compiled kernel also does not accept.
|
|
if x.dim() == 4 and positions is not None:
|
|
raise NotImplementedError(
|
|
"4-D query with explicit positions has no reference equivalent: the reference "
|
|
"indexes positions along its batch axis. Pass None to use the default ordering."
|
|
)
|
|
|
|
rotated = _ref.rotary_embedding_ref(
|
|
view,
|
|
cos,
|
|
sin,
|
|
positions,
|
|
seqlen_offsets=seqlen_offset,
|
|
interleaved=interleaved,
|
|
conjugate=conjugate,
|
|
)
|
|
# In-place contract: the caller keeps using its own tensor.
|
|
x.copy_(rotated.view(original_shape))
|
|
|
|
rotate(query)
|
|
if key is not None:
|
|
rotate(key)
|
|
|
|
|
|
def backend() -> str:
|
|
return "ref" if use_references() else "cuda"
|