jburtoft's picture
v1.0.0: NKI kernels for NewBERT (gte) encoders on trn2
25e8424
|
Raw History Blame Contribute Delete
1.41 kB

Tests

Marker File Needs What it checks
sim test_sim.py nki, numpy, ml_dtypes each kernel under nki.simulate vs a NumPy fp32 reference (ragged padding, LNC 1/2, query split, padding isolation)
cpu test_pack_cpu.py torch, transformers host-side weight packing reproduces the Hugging Face NewLayer attention half in fp32 (GTE_TEST_MODEL picks the model)
sim + cpu test_patch_sim.py all of the above patch_model end to end on a 2-layer model with the kernels run by the simulator
neuron test_device.py trn2 + PyTorch Native full patched model vs CPU fp32 and vs stock PyTorch Native: score Pearson >= 0.9999, top-1 agreement >= stock, padding isolation bit-exact, fused ops actually used
python -m pytest -m sim                       # local, no device
python -m pytest -m cpu                       # local, downloads the model
NEURON_LOGICAL_NC_CONFIG=1 python -m pytest -m neuron -s tests/test_device.py   # on trn2, once per LNC
NEURON_LOGICAL_NC_CONFIG=2 python -m pytest -m neuron -s tests/test_device.py

A simulator pass is not device evidence. The simulator does not model SBUF/PSUM capacity, and some DMA access patterns behave differently on hardware (for example, one strided DMA spanning several 32-partition blocks is exact in the simulator but copies only the first block on trn2).