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.
209 lines
7.0 KiB
Python
209 lines
7.0 KiB
Python
"""`tile_kernels.quant` -- FP8/FP4 casting, fused SwiGLU, scale-factor layout.
|
|
|
|
The compiled functions take `num_per_channels` and write into caller-provided `out` buffers; the
|
|
references take `block_size` and return fresh tensors. The signatures here are the compiled ones,
|
|
because that is what the 4 call sites in the model are written against.
|
|
|
|
The adapters are thin but not vacuous -- three things genuinely differ:
|
|
|
|
* `num_per_channels=n` means a per-token block of `(1, n)`, not `(n, n)`.
|
|
* `swiglu_forward_and_per_token_cast` fuses activation and cast; the reference has them separate,
|
|
so it is activation followed by cast.
|
|
* `cast_back`'s `fmt` is the *output* dtype (`'bf16'`/`'fp32'`), which is easy to misread as the
|
|
quantised format.
|
|
|
|
Verified against the compiled kernels: `per_token_cast`, `cast_back` and
|
|
`swiglu_forward_and_per_token_cast` are bitwise identical on the SM90 recipe
|
|
(`scripts/tile_kernels_patch/compare_with_cuda.py`).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Sequence
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from deepinfer.tk_compat import use_references
|
|
|
|
if not use_references():
|
|
from tile_kernels import quant 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 per_token_cast(
|
|
x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
|
|
fmt: str,
|
|
num_per_channels: int,
|
|
x_block_size: tuple[int, int] | None = None,
|
|
use_tma_aligned_col_major_sf: bool = False,
|
|
round_sf: bool = False,
|
|
use_packed_ue8m0: bool = False,
|
|
use_e4m3_sf: bool = False,
|
|
stochastic_cast: bool = False,
|
|
out: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Cast to FP8/FP4 with per-token scaling factors."""
|
|
if not use_references():
|
|
return _compiled.per_token_cast(
|
|
x,
|
|
fmt,
|
|
num_per_channels,
|
|
x_block_size,
|
|
use_tma_aligned_col_major_sf,
|
|
round_sf,
|
|
use_packed_ue8m0,
|
|
use_e4m3_sf,
|
|
stochastic_cast,
|
|
out,
|
|
)
|
|
if stochastic_cast:
|
|
raise NotImplementedError(
|
|
"stochastic_cast has no torch reference; it changes rounding, so silently using the "
|
|
"deterministic path would produce plausible but different numbers"
|
|
)
|
|
result = _ref.cast(
|
|
x,
|
|
fmt,
|
|
block_size=(1, num_per_channels),
|
|
x_block_size=x_block_size,
|
|
round_sf=round_sf,
|
|
use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
|
|
use_packed_ue8m0=use_packed_ue8m0,
|
|
use_e4m3_sf=use_e4m3_sf,
|
|
)
|
|
if out is not None:
|
|
# Honour the out-buffer contract: the call sites allocate these and may hold views.
|
|
out[0].copy_(result[0])
|
|
out[1].copy_(result[1])
|
|
return out[0], out[1]
|
|
return result
|
|
|
|
|
|
def cast_back(
|
|
x: tuple[torch.Tensor, torch.Tensor],
|
|
fmt: str,
|
|
x_block_size: tuple[int, int],
|
|
x_special_fmt: str | None = None,
|
|
) -> torch.Tensor:
|
|
"""Dequantise an FP8/FP4 tensor. `fmt` is the output dtype, not the quantised format."""
|
|
if not use_references():
|
|
return _compiled.cast_back(x, fmt, x_block_size, x_special_fmt)
|
|
if x_special_fmt is not None:
|
|
raise NotImplementedError(f"x_special_fmt={x_special_fmt!r} has no torch reference")
|
|
return _ref.cast_back(x, fmt, block_size=x_block_size)
|
|
|
|
|
|
def swiglu_forward(
|
|
x: torch.Tensor,
|
|
fmt: str,
|
|
psum_num_tokens_per_expert: torch.Tensor | None = None,
|
|
topk_weights: torch.Tensor | None = None,
|
|
routed_scaling_factor: float | None = None,
|
|
clamp_value: float | None = None,
|
|
clamped_count: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
"""SwiGLU activation. `fmt` is unused by the compiled path here; kept for signature parity."""
|
|
if not use_references():
|
|
return _compiled.swiglu_forward(
|
|
x,
|
|
fmt,
|
|
psum_num_tokens_per_expert,
|
|
topk_weights,
|
|
routed_scaling_factor,
|
|
clamp_value,
|
|
clamped_count,
|
|
)
|
|
return _ref.swiglu_forward(
|
|
x,
|
|
psum_num_tokens_per_expert,
|
|
topk_weights,
|
|
routed_scaling_factor,
|
|
clamp_value,
|
|
clamped_count,
|
|
)
|
|
|
|
|
|
def swiglu_forward_and_per_token_cast(
|
|
x: torch.Tensor,
|
|
fmt: str,
|
|
num_per_channels: int,
|
|
psum_num_tokens_per_expert: torch.Tensor | None = None,
|
|
topk_weights: torch.Tensor | None = None,
|
|
routed_scaling_factor: float | None = None,
|
|
use_tma_aligned_col_major_sf: bool = False,
|
|
round_sf: bool = False,
|
|
use_packed_ue8m0: bool = False,
|
|
clamp_value: float | None = None,
|
|
clamped_count: torch.Tensor | None = None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Fused SwiGLU + per-token cast."""
|
|
if not use_references():
|
|
return _compiled.swiglu_forward_and_per_token_cast(
|
|
x,
|
|
fmt,
|
|
num_per_channels,
|
|
psum_num_tokens_per_expert,
|
|
topk_weights,
|
|
routed_scaling_factor,
|
|
use_tma_aligned_col_major_sf,
|
|
round_sf,
|
|
use_packed_ue8m0,
|
|
clamp_value,
|
|
clamped_count,
|
|
)
|
|
activated = _ref.swiglu_forward(
|
|
x,
|
|
psum_num_tokens_per_expert,
|
|
topk_weights,
|
|
routed_scaling_factor,
|
|
clamp_value,
|
|
clamped_count,
|
|
)
|
|
return _ref.cast(
|
|
activated,
|
|
fmt,
|
|
block_size=(1, num_per_channels),
|
|
round_sf=round_sf,
|
|
use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
|
|
use_packed_ue8m0=use_packed_ue8m0,
|
|
)
|
|
|
|
|
|
def batched_transpose_weight_sf(
|
|
sf: Sequence[torch.Tensor],
|
|
block_size: int,
|
|
transpose_sf1s: Sequence[bool],
|
|
outs: Sequence[tuple[torch.Tensor, torch.Tensor] | None],
|
|
) -> list[tuple[torch.Tensor, torch.Tensor]]:
|
|
"""Transpose the scale factors of several weights in one launch."""
|
|
if not use_references():
|
|
return _compiled.batched_transpose_weight_sf(sf, block_size, transpose_sf1s, outs)
|
|
# The reference allocates its outputs, so the `outs` slots cannot be filled in place. Check
|
|
# rather than ignore: a caller that pre-allocated would read stale memory.
|
|
if any(o is not None for o in outs):
|
|
raise NotImplementedError("the torch reference allocates its own outputs; pre-allocated `outs` cannot be filled")
|
|
return _ref.batched_transpose_weight_sf(sf, block_size, transpose_sf1s)
|
|
|
|
|
|
def unpack_from_e2m1fn_x2(x: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
|
|
"""Decode packed FP4 pairs. Same on both paths -- it is host-side bit manipulation."""
|
|
return _ref.unpack_from_e2m1fn_x2(x, out_dtype)
|
|
|
|
|
|
def backend() -> str:
|
|
return "ref" if use_references() else "cuda"
|
|
|
|
|
|
def _passthrough(name: str) -> Any:
|
|
"""Resolve a name this module does not adapt from whichever backend is active."""
|
|
if not use_references():
|
|
return getattr(_compiled, name)
|
|
return getattr(_ref, name)
|