Aligned shared rows for pipelined NAX on CUDA

The explicit iron_steel_gemm_splitk_nax_bt_direct_32x64_pipelined_ld40 kernel lowers repeated BF16 complete-GEMM GPU time on GB10 by 38.63% / 41.18% at M96/N512/K2048 and 42.20% / 41.68% at M256/N512/K2048, in forward/reverse process order. The candidate uses inline PTX. Every primary round beats both the fastest measured current-dev configuration and the fastest held stride-36 pipeline configuration: 36/36 comparisons against each.

Source commit 3b3fd9ecbbd3070635bd58d5609d2887c24c12b1, tree 6597b9aa7b5f79f30d17b50053e777457ee66d25; actual dev baseline e6c3af82eee11596b51a209e70945f261f392eb7, including #342 and #345. Issue #347. The 186-line kernel body matches PR #299 at b29387a07689fce140b1fad5cd616e757093d092 exactly. All existing kernel-file content, all 690 other existing crate files, and all 5,527 old test rows and tolerances are preserved. The compiler, runtime and automatic selection do not change.

Scope and mechanism

The candidate uses the existing software pipeline with 40-element shared rows. Those rows make the F16 staging pitch a multiple of 16 bytes, allowing existing inline-PTX lowering to use direct aligned shared views. For BF16 M96/N512/K2048 at four splits, the generated shared-memory requirement is 15,872 bytes per workgroup, versus 27,392 for both the current-dev wide kernel and the held stride-36 pipeline under inline PTX. This is static code-generation evidence; hardware occupancy was not separately measured.

The claim concerns repeated complete GEMM with prepared weights and the explicit inline-PTX configuration. Initial CUDA partial-pass screening finds software regressions and no WMMA improvement. Metal screening generally loses to the held stride-36 pipeline. Those results remain in the evidence; no default-strategy, automatic-dispatch, Metal, H100-runtime or whole-model gain is claimed. The held candidate in #346 remains held under its failed Metal gate.

Contract and correctness

Inputs are row-major A[M,K] and B[N,K], in F32, F16 or BF16. The kernel writes F32 partials[splits,M,N]; an existing accumulation pass produces the final output. Dispatch exactly 128 threads with grid [N/64,M/32,splits]. M, K and K-per-split are positive multiples of 32; N is a positive multiple of 64. Every split starts below K, total split extent covers K, and the final split clamps to K. Inputs have exact extents and separate output storage without aliasing. Metal NAX requires Apple10 and Metal 4. BF16 stages through F16 on each backend, so callers must account for F16’s range.

Thirty new tests cover ten shapes and three dtypes, including a single tile, short final splits, odd tile counts, long K and larger matrices. They check every partial slab against independent F64 products of the actual dtype-rounded inputs, with a zero first input row, a larger final row, poisoned partials and 64 trailing guards. All initial partial-output tolerances remain 1e-4. A legal numerical negative replaces only the candidate’s final store value with zero: all 30 fail, and the restored body passes all 30. The final body matches that restored body exactly.

All 30 pass on M5 and under both CUDA memcheck and synccheck in all three CUDA strategies, giving 180 passing sanitizer executions with no skipped cases or errors. Complete-operation checks compare every final output with an independent F64 reference, poison outputs before both checks, and verify guards on both final and partial allocations. Their unchanged tolerances are F32 1e-4, F16 2e-3 and BF16 2e-2. Initial full-chain F32/F16/BF16 sweeps pass 99 configurations each; both full-chain BF16 sanitizers also pass all 99. The final optimized host example passes another 99 BF16 checks and emits all 183 generated CUDA files byte-for-byte identically to the validated prototype. On Spark1, all 30 inline-PTX cases and all 99 full-chain BF16 configurations pass both sanitizers again, with no skipped cases or errors. Its source files, frozen binaries and all 183 generated CUDA files match Spark2 exactly, without a rebuild.

Measurement protocol

One process contains the geometrically compatible current-dev dense, fused NAX, Steel and split-K families across software, WMMA and inline PTX. The sampled split grid is 1, 2, 4, 8, 16 for BT/wide/LD40 and the held pipeline, and 2, 4, 8, 16 for original NAX/Steel split families, filtered for compatible geometry. Comparisons use the fastest measured dev configuration, not an exhaustive search of every valid split count. Both weight layouts are uploaded once and shared among all configurations; dev receives its preferred layout without conversion cost. Every split-K arm includes the same existing elementwise accumulation kernel. Device events measure a captured complete-operation graph. Allocation, upload, layout conversion, source compilation and F64 reference calculation occur outside timing.

The primary candidate is fixed at four splits for M96 and two splits for M256, both BF16/N512/K2048, using inline PTX. The archived plan was frozen before confirmation at SHA256 2c93ae0f9f13af89cee05ac27feaf66c7417e795c74ec3596949b07f078f0ca1. It requires at least 10% lower time in both process-median and paired-round-median comparisons, in both directions, against both current dev and the held pipeline, plus every primary round faster. No primary, direction, losing round or threshold was changed after sampling.

Confirmation contains 730,080 valid samples in 20 processes, with 4,056 pre/post complete-output F64 and allocation-guard checks. Each process has nine rounds, 20 warmups and 40 recorded samples per round and configuration. Arm order rotates every repetition. The second direction reverses both case and arm order. Process comparisons use the lowest dev process median; paired comparisons use the lowest dev median within each round and then the median of nine percentage reductions. The held control is treated the same way. Controls include F16, larger N, short K, tiny matrices, F32, and pack=8 repeats of both primary shapes. Controls cannot rescue a failed primary gate.

Case M×N×K	Splits / pack	Dev → candidate µs, forward	Dev → candidate µs, reverse	Dev process reduction	Dev paired-round reduction	Dev round wins	Held process reduction	Held round wins
BF16 96×512×2048	4 / 1	37.440 → 22.976	38.464 → 22.624	38.63% / 41.18%	39.07% / 41.47%	18/18	38.53% / 38.04%	18/18
BF16 256×512×2048	2 / 1	87.968 → 50.848	88.224 → 51.456	42.20% / 41.68%	42.14% / 41.71%	18/18	40.58% / 39.09%	18/18
F16 96×512×2048	4 / 1	38.464 → 23.072	38.432 → 22.368	40.02% / 41.80%	39.97% / 41.80%	18/18	38.27% / 40.05%	18/18
F16 256×512×2048	2 / 1	88.288 → 51.088	87.888 → 51.232	42.13% / 41.71%	41.99% / 41.78%	18/18	40.85% / 40.13%	18/18
BF16 256×1024×2048	2 / 1	167.328 → 92.512	167.488 → 92.384	44.71% / 44.84%	44.44% / 44.82%	18/18	43.73% / 43.07%	18/18
BF16 96×512×512	4 / 1	16.560 → 12.240	16.064 → 11.968	26.09% / 25.50%	25.63% / 25.55%	18/18	22.88% / 25.35%	18/18
BF16 32×64×32	1 / 1	4.432 → 4.000	4.576 → 4.416	9.75% / 3.50%	6.72% / -0.37%	13/18	30.94% / 24.18%	18/18
F32 96×512×2048	4 / 1	67.552 → 67.568	67.632 → 68.960	-0.02% / -1.96%	-0.12% / -1.37%	3/18	-5.95% / -7.64%	0/18
BF16 96×512×2048	4 / 8	32.954 → 17.680	33.284 → 17.768	46.35% / 46.62%	46.39% / 46.39%	18/18	44.63% / 44.90%	18/18
BF16 256×512×2048	2 / 8	83.676 → 45.876	83.724 → 46.060	45.17% / 44.99%	45.16% / 45.01%	18/18	42.57% / 41.79%	18/18

Only the first two rows define the promotion gate. Negative reductions are regressions and are retained. Pack=8 reports device-event time divided by eight, checking sensitivity to launch/event overhead. Available clocks, temperature and power are recorded before and after each confirmation process; these snapshots do not establish constant clocks during every sample. No absolute timing level is discarded.

The earlier optimized Spark2 screen contains 97,920 samples and 3,264 checks in all 16 planned processes. It excludes and preserves 2 whole processes that overlapped observed Cargo/rustc activity. The earlier prototype screen was retired after seven valid processes (41,580 samples and 1,386 checks) and three overlapping attempts; all data are retained separately. Only our waiting supervisor was stopped. Compiling the host example with opt-level=3 reduced CPU setup time; GPU source remained identical. Every exploratory case was then repeated using the final binary.

Spark2 was occupied when independent confirmation was first prepared, so a pre-sampling hardware record selected Spark1. An initial Spark1 screen was retired wholesale when another GPU benchmark was observed running despite our shared lock. Only our own screen supervisors and active benchmark were interrupted; every completed or partial attempt is archived separately. A frozen monitoring record added direct NVIDIA compute-process checks to enforce the original quiet-host requirement. Spark1 remained busy and its replacement screen never started a process. Once Spark2 became idle, a new pre-sampling hardware record moved the entire final campaign back there, with no confirmation samples yet collected.

The unchanged frozen binary repeated the complete monitored screen on Spark2: 97,920 samples and 3,264 checks in all 16 processes, with 0 whole-process compilation or foreign-GPU overlap exclusions preserved. All 20 independent confirmation processes also use Spark2. Host results, binary versions, retired screens, exploratory samples and confirmation samples are kept separate. The cases, sampling, directions, fixed splits and original promotion gate did not change.

All 20 confirmation processes were accepted, with no observed compilation or foreign-GPU overlap. Before each screen or confirmation process, the driver requires 30 consecutive seconds with no Cargo/rustc or other GPU process. During execution it checks both resources in every polling iteration with a 250 ms sleep, excluding only GPU processes descended from our own launch. The shared /tmp/iron-gpu.lock is also held. Other agents’ processes were not interrupted. All accepted final-screen and confirmation processes have no observed compilation or foreign-GPU overlap.

Repository validation and provenance

• Full workspace: 1,244 passed, zero failed, 59 ignored in 134 result groups, including the complete registered kernel harness.
• Inventory witness: 5,557 cases, FNV a2997536b0ed4cae; all 5,527 old rows and tolerances remain unchanged.
• Workspace/all-target/all-feature Clippy, final Linux CUDA all-target Clippy, format, inventory, typos and the 664-kernel binding audit pass.
• All 696 final source and dependency files match local and CUDA snapshots, including unchanged Cargo.lock SHA256 2427939dd3932e3270c5d4b3119e23b3be2c67442cd6edd6a6b2251f672c8b0c. No dependencies were added. Cargo.lock is ignored by this repository; its frozen copy is retained in the source archive.
• Confirmation and final screen: GB10 Spark2, GPU GPU-18624310-3601-3bc3-9904-3618f57acd14; additional correctness validation and retired screen: GB10 Spark1, GPU GPU-e6b83176-1b83-cbdb-3446-0548daa571bd. Both use driver 580.173.02, CUDA 13.0.88, Rust 1.98.0. The host example uses opt-level=3 with no fast-math flags; dependencies use the ordinary dev profile with debug information and incremental compilation disabled.
• CUDA confirmation executable SHA256 cd1ee47a6313845d23ecf0cb281d6bf70049640197f0c57193880daea69edc3c; CUDA runner SHA256 7c265b098ec8b30d0f8c39791005a4448c7f53943426b6c1dfaaa124fd8b8e53. Both remain frozen throughout their runs.
• M5 Max, 128 GiB, macOS 27.0 build 26A5353q. Final Metal runner SHA256 a477d9954d40c3ebfae76a3cffdaafd036b3d0f8820dc91ad886755d88201e6d.

The CUDA builds and Clippy succeeded; a subsequent binary-copy step exhausted disk space. Only our completed earlier integration build cache was removed after preserving and verifying its runner. The incomplete copy was repaired and both final binaries and all source files were reverified without rebuilding. The failed-copy log and recovery receipt are retained.

Six native Metal benchmarks pass with 20 warmups and 40 runs, without observed compilation overlap. They measure only the partial pass, independently of the complete-operation performance claim. The repository has no make bench target.

Partial-pass shape	Dtype	Minimum µs	Mean µs
partial pass m256 n512 k2048 splits2	f32	54.125	59.815
partial pass m256 n512 k2048 splits2	f16	17.250	24.221
partial pass m256 n512 k2048 splits2	bf16	17.500	21.644
partial pass m96 n512 k2048 splits4	f32	23.250	23.620
partial pass m96 n512 k2048 splits4	f16	9.417	9.727
partial pass m96 n512 k2048 splits4	bf16	9.625	10.044

Reproduction

The final-source archive contains the exact 696 audited files and frozen Cargo.lock. On Linux CUDA, build wh-iron-std with the cuda feature and the __iron_runner binary. Build the nax_ld40_cuda_chain example using cargo rustc –locked -p wh-iron-std –features cuda –example nax_ld40_cuda_chain – -C opt-level=3. The archived build script records the complete environment.

For a primary run set GEMM_DTYPE=bf16, GEMM_M=96, GEMM_N=512, GEMM_K=2048, GEMM_ROUNDS=9, GEMM_WARMUP=20, GEMM_SAMPLES=40, GEMM_PACK=1 and GEMM_REVERSE=0, then repeat with reverse=1. M256 is the other primary. Leave GEMM_CONFIGS, GEMM_ARMS and GEMM_STRATEGIES unset to include all configurations. GEMM_CHECK_ONLY=1 checks correctness without timing, and GEMM_EMIT_DIR records generated CUDA. The archived protocol pins every control and split count and selects the fixed candidate during analysis.

The report directory contains one evidence.tar.gz archive. Its root SHA256SUMS verifies the archive; receipts/SHA256SUMS inside it verifies every individual receipt. Extract the archive to obtain the source snapshots, raw rows and drivers.

Drivers pin the published hashes. Rebuilding on another host requires a new provenance record while retaining this evidence. SHA256SUMS verifies every packaged receipt; raw samples, exclusions, oracles, numerical control, source snapshots, commands and the predeclared plan remain available. AI assistance was used for extraction, cleanup, benchmarking and validation.