mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +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>
12 lines
722 B
Python
12 lines
722 B
Python
import torch
|
|
from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import philox4x32_10
|
|
def t(v): return torch.tensor([v], dtype=torch.int64)
|
|
for ctr, key, exp in [((0,0,0,0),(0,0),(0x6627e8d5,0xe169c58d,0xbc57ac4c,0x9b00dbd8)),
|
|
((0xffffffff,)*4,(0xffffffff,)*2,(0x408f276d,0x41c83b0e,0xa20bc7c6,0x6d5451fd)),
|
|
((0x243f6a88,0x85a308d3,0x13198a2e,0x03707344),(0xa4093822,0x299f31d0),(0xd16cfe09,0x94fdcceb,0x5001e420,0x24126ea1))]:
|
|
out = philox4x32_10(*map(t, ctr), *map(t, key))
|
|
got = tuple(int(o) for o in out)
|
|
print([hex(g) for g in got], got == exp)
|
|
from deepinfer.model_executor.layers.tilelang_ops import ops
|
|
print(ops.random_sampling_tl, ops.fused_sample_tl, ops.map_lookup_tl)
|