Write numerical equivalence tests for a Spyre-aware kernel (spyre.py) vs the original, including distribution invariance across core counts. Produces tests/triton/test_<name>_spyre.py. Use when asked to test a Spyre kernel. For plain _td kernels, use td-test.
التثبيت
التثبيت باستخدام Codex أو Claude انسخ هذا Prompt والصقه في Codex أو Claude أو مساعد آخر ليراجع صفحة Skill ويثبّتها لك.
Write numerical equivalence tests for a Spyre-aware kernel (spyre.py) vs the original, including distribution invariance across core counts. Produces tests/triton/test_<name>_spyre.py. Use when asked to test a Spyre kernel. For plain _td kernels, use td-test.
Spyre Kernel Test Skill
Write numerical-equivalence tests for a Spyre-aware kernel against the original,
including distribution invariance (same result for any core count). Output:
tests/triton/test_<name>_spyre.py. For a _td kernel (no distribution loop),
use td-test.
Trigger
Use when asked to write/add tests for a spyre.py kernel.
Pre-flight
Run the KB consult in ../_shared/preflight.md
(precision notes for Spyre DL16 can shift tolerances).
Read kernels/<name>/spyre.py (signature), kernels/<name>/wrapper.py
(launch), and tests/triton/test_<name>.py if present (style).
Execution path (intended vs interim)
The right Spyre verification path is not "GPU triton":
Lower Triton → TTIR → KTIR via the mlir-frontend (the Triton repo's spyre
backend; dump_round_trip.py does this end-to-end).
Execute the KTIR on ktir-cpu (the test_ktir_examples.py flow).
Interim: until that path is wired into this repo's test harness, tests run
the kernel on GPU/CUDA (the @pytest.fixture device + torch.cuda pattern
below) and compare against the original kernel. Write tests against the
original-kernel reference so they port cleanly to the ktir-cpu path later.
Launching the kernel — reuse wrapper.py
Launch both kernels through the existing kernels/<name>/wrapper.py (which
dispatches to the Spyre kernel, typically via a kernel_fn= argument) — do not
copy its launch logic into the test. A forked launch path drifts from
production and must be updated twice when the signature changes.
Distribution-invariance tests need to vary num_cores, and the wrapper should
accept num_cores as a parameter so the test drives it through the same
launch path the production code uses. If the wrapper does not yet expose
num_cores (or another parameter a test needs), extend wrapper.py rather
than building a parallel launch helper in the test.
Test categories
1. Correctness vs original
Parametrize over representative shapes (include non-divisible), seed with
torch.manual_seed(42), random inputs, torch.testing.assert_close.
Tolerances — size to the output dtype's ULP, via a per-dtype TOL dict
(same rationale as td-test):
Loosen up only with a stated reason. Important — matmul / tl.dot
reductions:tl.dot and a NumPy @ reference accumulate the K dimension in
different orders, so a 1e-5 rtol is unrealistic — the project's
matmul/default fixture uses rtol≈1e-2. Align matmul-style tolerances with
that fixture, not the ULP table. Single-tile cases (size <= BLOCK_SIZE)
have no reduction-order divergence — assert atol=0, rtol=0.
2. Distribution invariance
Run the Spyre kernel with num_cores in [1, 4, 16, 32]; compare each against
the original. Also test: more cores than work items (idle cores) and
single core (sequential). This verifies the distribution loop for any
partitioning.
3. Edge cases
Non-divisible shapes; minimum size (< one tile); asymmetric shapes; zero input
(eps); large values; kernel-specific edges (topk=1, single batch, vocab <
BLOCK_SIZE).
Spyre wrapper conventions
Grid: fixed (num_cores,), not derived from problem size.
No stride args if the kernel assumes contiguous layout (assert it).
Explicit constexprs (no autotune).
Extra runtime args (num_requests, num_elements) that replace
grid-size-equals-problem-size patterns.
Checklist
Test at tests/triton/test_<name>_spyre.py
Both kernels launched via wrapper.py with configurable num_cores
(extend the wrapper if it doesn't expose it — no forked launch path)
Correctness parametrized over shapes; distribution tests over [1,4,16,32]
Edge cases: non-divisible, min size, asymmetric, more-cores-than-work,
single-core
Compare with torch.testing.assert_close against the original only
Per-dtype TOL; matmul-style tl.dot uses fixture rtol (~1e-2), not 1e-5
