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>
21 lines
776 B
Python
21 lines
776 B
Python
import torch
|
|
import deepinfer # noqa: sets vendor path
|
|
from deepinfer.model_executor.models.v4_modules import mhc_ops
|
|
import tile_kernels.torch as ref
|
|
torch.manual_seed(0)
|
|
t, h, M = 6, 256, 4
|
|
m3 = M * (M + 2)
|
|
res = torch.randn(1, t, M, h, dtype=torch.bfloat16)
|
|
fn = torch.randn(m3, M * h) * 0.02
|
|
w = torch.rand(M * h) + 0.5
|
|
print('n_splits', mhc_ops.default_n_splits())
|
|
for ns in (1, 16):
|
|
try:
|
|
out = mhc_ops.mhc_pre_norm_fn(res, fn, w, 1e-6, n_splits=ns)
|
|
r = ref.mhc_pre_norm_fn_ref(res, fn, w, 1e-6)
|
|
print(ns, 'max diff', (out - r).abs().max().item())
|
|
except Exception as e:
|
|
print(ns, 'ERR', type(e).__name__, str(e)[:300])
|
|
x = torch.randn(t, 2, 128, dtype=torch.bfloat16)
|
|
o = mhc_ops.expand_to_mhc(x, 4); print('expand ok', o.shape)
|