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>
37 lines
2.0 KiB
Python
37 lines
2.0 KiB
Python
import os
|
|
import torch
|
|
from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import (
|
|
two_level_sample, quantize_weights, acc_slots_for, rng_first_two, rng_sequence)
|
|
from deepinfer.model_executor.layers.tilelang_ops.random_sampling import random_sampling
|
|
from deepinfer.model_executor.layers.tilelang_ops.sampler import prepare_verification_metadata, prepare_verification_metadata_ref
|
|
|
|
for V in [8, 100, 1024, 4032, 8192, 129280]:
|
|
Q = 64
|
|
p = torch.rand(Q, V) ** 8; p[:, ::3] = 0; p = p / p.sum(1, keepdim=True)
|
|
seeds = torch.randint(-2**62, 2**62, (Q,), dtype=torch.int64)
|
|
pos = torch.randint(-5, 10**9, (Q,), dtype=torch.int32)
|
|
out = two_level_sample(p, seeds, pos, 1, V)
|
|
# independent: inverse CDF over columns ordered by (col % AS, col // AS)
|
|
AS = acc_slots_for(V)
|
|
fx = quantize_weights(p).to(torch.int64)
|
|
cols = torch.arange(V)
|
|
order = torch.argsort((cols % AS) * (V + 1) + cols // AS)
|
|
cdf = fx[:, order].cumsum(1)
|
|
total = cdf[:, -1]
|
|
lo, hi = rng_first_two(seeds, rng_sequence(pos, 1))
|
|
r = torch.tensor([((int(h) << 32) | int(l)) % int(t) for h, l, t in zip(hi, lo, total)])
|
|
exp = order[(cdf <= r[:, None]).sum(1)]
|
|
assert torch.equal(out.long(), exp), V
|
|
assert (p.gather(1, out.long()[:, None]) > 0).all()
|
|
print("two-level == slot-ordered inverse CDF: OK")
|
|
# strided probs rows
|
|
p = torch.softmax(torch.randn(16, 2048), -1); big = torch.zeros(16, 4096); big[:, :2048] = p
|
|
a = random_sampling(big[:, :2048], torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
|
|
b = random_sampling(p.contiguous(), torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
|
|
assert torch.equal(a, b); print("strided rows OK")
|
|
# nextn=0 metadata path (num_steps=0) and seqlen=1 rows
|
|
seqlen = torch.ones(37, dtype=torch.int32); ftp = torch.arange(37, dtype=torch.int32)
|
|
for got, ref in zip(prepare_verification_metadata(ftp, seqlen, 37, 0), prepare_verification_metadata_ref(ftp, seqlen, 37)):
|
|
assert torch.equal(got, ref)
|
|
print("prepare_verification_metadata num_steps=0 OK")
|