mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 07:17:08 +08:00
760e3cadc0
- tests/neuron/setup_trn2.sh: one-shot setup of a fresh Trn2 (venv layered
on the DLAMI torch-neuronx venv, C++ build, official reference + tokenizer,
tiny checkpoint, CPU golden).
- tests/neuron/run_all.sh cpu|device|all [stage...]: vendor kernel tests,
tilelang_ops interface tests (basic-mode stubs reported as skips via
nie_skip), deep_ep at world size 1/2/4 (gloo) and 2/4/8 on the neuron
process group, device primitive checks, end-to-end vs golden on CPU and
on a NeuronCore, and the device diff; PASS/FAIL summary.
- tests/neuron/device/test_device_primitives.py: torch-neuronx behaviours
the port depends on or works around (fp8 strided-view indexing, uint64,
1-rank all_to_all, inference-mode version counters, ...).
- examples/neuron/device_diff.py: re-runs every kernel call of a NeuronCore
run on host copies and ranks kernels by device/host difference.
- Unit tests written during bring-up moved from scratch dirs into
tests/neuron/{deep_gemm,flash_mla,deep_ep,tile_kernels,sampling}.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
53 lines
1.9 KiB
Python
53 lines
1.9 KiB
Python
"""Cross-check: tile_kernels.quant.per_token_cast output feeds deep_gemm's FP8 GEMM.
|
|
|
|
This is the exact ReplicatedLinear.apply_weight_fp8 path
|
|
(deepinfer/model_executor/layers/linear.py) with CPU tensors.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
|
|
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
|
|
|
|
import torch
|
|
|
|
import deep_gemm
|
|
import tile_kernels
|
|
|
|
CAST_PARAMS = dict(fmt="e4m3", num_per_channels=32, round_sf=True,
|
|
use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
|
|
|
|
|
|
def calc_diff(x, y):
|
|
x, y = x.double(), y.double()
|
|
denom = (x * x + y * y).sum().item()
|
|
return 0.0 if denom == 0 else 1 - 2 * (x * y).sum().item() / denom
|
|
|
|
|
|
m, n, k = 48, 128, 256
|
|
x = torch.randn(m, k, dtype=torch.bfloat16)
|
|
w = torch.randn(n, k, dtype=torch.bfloat16)
|
|
|
|
xq, x_sf = tile_kernels.quant.per_token_cast(x, **CAST_PARAMS)
|
|
wq, w_sf = tile_kernels.quant.per_token_cast(w, **CAST_PARAMS)
|
|
print(f"{xq.dtype=} {x_sf.dtype=} {x_sf.shape=} {x_sf.stride()=}")
|
|
print(f"{wq.dtype=} {w_sf.dtype=} {w_sf.shape=} {w_sf.stride()=}")
|
|
|
|
out = torch.empty(m, n, dtype=torch.bfloat16)
|
|
deep_gemm.fp8_fp4_gemm_nt((xq, x_sf), (wq, w_sf), out, None, recipe=(1, 1, 32))
|
|
|
|
ref = x.float() @ w.float().T
|
|
diff = calc_diff(out.float(), ref)
|
|
print(f"diff vs bf16 reference: {diff:.3e}")
|
|
assert diff < 2e-3, diff
|
|
|
|
# The dequantized operands must agree with the GEMM result exactly.
|
|
from deep_gemm._quant import dequantize # noqa: E402
|
|
|
|
x_deq = dequantize(xq, deep_gemm.transform_sf_into_required_layout(x_sf, m, k, (1, 1, 32), is_sfa=True))
|
|
w_deq = dequantize(wq, deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 1, 32), is_sfa=False))
|
|
exact = (x_deq @ w_deq.T).to(torch.bfloat16)
|
|
print("exact match with dequantized matmul:", torch.equal(out, exact))
|
|
print(f"quantization error (dequant vs bf16): {calc_diff(exact.float(), ref):.3e}")
|
|
print("OK")
|