Add NKI partition equivalence and hardware scaling tests

[scttfrdmn]

Phase 4 PR 3 — simulator and hardware tests for counter-partitioned multi-chip RNG

Closes portion of #20 (hardware validation items).

What this adds

tests/test_nki_partitioning.py — all tests skipped automatically on dev hosts
without neuronxcc. Two gate markers:

• @pytest.mark.nki_simulator — requires TRNRAND_USE_SIMULATOR=1 + nki>=0.3.0
• @pytest.mark.neuron — trn1.32xlarge hardware only, run manually

TestPartitionEquivalencePerTile (nki_simulator)

Bit-exact partition equivalence for uniform, normal, exponential via the
per-tile dispatch path (threefry_uniform_nki / threefry_normal_nki), P=2 and P=4.

Concatenating P-chip draws in rank order must equal the 1-chip draw — torch.equal.

TestPartitionEquivalenceStreaming (nki_simulator)

Same invariant tested via threefry_stream_normal / threefry_stream_uniform
directly, so counter_offset arithmetic in streaming dispatch is independently
validated without going through GeneratorProgram.

Counter offset per chip: r × ceil(n / 16384) × _PROGRAM_TILES.

TestPartitionEquivalenceProgram (nki_simulator)

Validates GeneratorProgram.stream_into applies partition offsets correctly.
Includes a multi-distribution program test (normal + uniform in the same program, P=2).

TestAdvanceResumeEquivalence (nki_simulator)

advance(N_skip) followed by continued generation matches the tail of an
uninterrupted N_skip + N_draw run. Also validates advance_to(N) == advance(N).

TestHardwarePartition (neuron — trn1.32xlarge, manual)

• test_zero_cross_chip_coordination — smoke test; profiler confirms no collective ops during kernel (NEURON_PPROF_SINK=/tmp/profile)
• test_near_linear_strong_scaling — 1-chip baseline for ≥87% efficiency (≥28× 1-chip) target at 32 chips
• test_partition_equivalence_P32_hardware — chip 0 of 32 matches first slice of 1-chip output (structural validation of counter arithmetic at P=32 without requiring 32 physical chips)
• test_neff_cache_reuse_with_partition — second stream_into with partition args completes in <100ms (NEFF cache hit)

Partition equivalence boundary

Equivalence holds when N is a multiple of 512 × P (per-tile) or 16,384 × P (streaming).
Test sizes are chosen to satisfy this. Partial-batch behavior is already covered in
test_partitioning.py.

Running

# Dev host (all marked tests skip automatically)
pytest tests/test_nki_partitioning.py -v

# NKI simulator
TRNRAND_USE_SIMULATOR=1 pytest tests/test_nki_partitioning.py -m nki_simulator -v

# Hardware
pytest tests/test_nki_partitioning.py -m neuron -v