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>
39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
"""Checks the tilelang stub provides exactly what platform-independent DeepInfer code uses."""
|
|
|
|
import sys
|
|
|
|
import torch
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
|
|
FAILS = []
|
|
|
|
|
|
def chk(name, cond):
|
|
print(f"[{'ok ' if cond else 'FAIL'}] {name}")
|
|
if not cond:
|
|
FAILS.append(name)
|
|
|
|
|
|
chk("PassConfigKey.TL_ENABLE_FAST_MATH", isinstance(tilelang.PassConfigKey.TL_ENABLE_FAST_MATH, str))
|
|
chk("T.float8_e4m3fn", T.float8_e4m3fn == "float8_e4m3fn")
|
|
chk("dtype map", tilelang.TL_DTYPE_TO_TORCH[T.float8_e4m3fn] is torch.float8_e4m3fn)
|
|
chk("norm.py dtype strings", T.float32 == "float" and T.bfloat16 == "bfloat16")
|
|
chk("max_value e4m3", T.max_value(T.float8_e4m3fn) == 448.0)
|
|
chk("min_value fp32", T.min_value(T.float32) == torch.finfo(torch.float32).min)
|
|
|
|
# the pass-config dict pattern of tilelang_ops/compress.py
|
|
cfg = {}
|
|
cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = True
|
|
chk("pass_configs assignment", cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] is True)
|
|
|
|
try:
|
|
T.some_dsl_primitive
|
|
FAILS.append("unknown attr should raise")
|
|
except AttributeError as e:
|
|
chk("unknown T attribute raises AttributeError", "Neuron stub" in str(e))
|
|
|
|
print("\nFAILURES:", FAILS if FAILS else "none")
|
|
sys.exit(1 if FAILS else 0)
|