ATen full HW matrix CI (mlen×blen×batch×seq) + genuine GQA validation + native CLM-60M MLEN≥256 NaN fix

[booth-algo]

Summary

Makes the standalone ATen testbenches correct and genuinely validated across the full hardware tile matrix (MLEN/BLEN, batch, sequence length), replaces the old per-MLEN sweep with a fast per-PR quick check + an on-demand full matrix, and fixes the native CLM-60M MLEN≥256 all-NaN bug, the FFN down-projection K-split bug, and sub-64 MLEN support. Merged with main (PR #62's DMA-byte-primitive refactor) with the sub-64 support re-ported onto the new dma.rs. Companion PRs (all merged): AICrossSim/PLENA_Compiler#49, #50; AICrossSim/PLENA_Tools#4.

Unified [batch_size, seq_len, hidden] interface

All ATen tests take --batch-size and --seq-len and resolve rows = batch_size × seq_len via a shared resolve_rows(args, default_seq) helper in configurable.py. The per-row-independent ops (linear, softmax, rms/layer-norm, ffn, embedding-add, rope) flatten to [rows, hidden]; attention uses true [batch, seq, heads, head_dim]. No-arg defaults reproduce prior single-tensor behavior exactly.

CI: fast per-PR quick check + manual full matrix

• aten-tests.yml — "ATen Tests (quick)" (automatic, push/PR): a single job running ~12 crucial cells at mlen≤64 — all 8 ops @ mlen64/blen8/b1/seq64, a sub-64 RTL-tile path (linear + flash-attention @ mlen16/blen4/seq16), and a batch>1 path (mlen64/blen16/b2/seq32). Targets under ~10 min.

• aten-matrix-full.yml — "ATen Tests (full matrix, manual)" (workflow_dispatch only): the exhaustive sweep 8 ops × mlen{16,64,256} × blen{4,8,16,32} × batch{1,2} × seq{4,16,64,128,256} (960 logical cells), matrixed over (mlen, batch) = 6 jobs, bash-looping test × blen × seq with pruning rules:

  1. blen > mlen
  2. mlen % blen != 0
  3. (batch×seq) % blen != 0
  4. flash-attention seq_len > mlen (packed-GQA one-tile limit)
  5. softmax rows > mlen (single [rows, mlen] score-tile accumulator)

  ~45–50 min (the mlen=256 jobs dominate); run it before merges / DMA / tile-size changes.

Both share the cargo cache key. A "Free disk space" step reclaims ~20–30 GB so the torch+cu126 install + nix store + cargo target fit on the ~14 GB hosted runner (the parallel jobs were OOM-ing on disk). There are no tracked skips — every non-pruned cell passes with real numerical agreement.

GQA flash-attention: genuine validation (was passing vacuously)

The standalone GQA test was producing all-zero simulated output in every configuration (Mean Relative Error = 1.0, K/V matrix SRAM all-zero) and only "passing" because small golden values stayed within the loose atol=rtol=0.2 band vs zero. Root cause: the test never supplied tensor_layouts, so it inherited a stale build-dir tensor_layouts.json from a prior larger-MLEN run → K/V strided wrong and loaded zeroed into MRAM → QK^T=0 → uniform softmax → ~0 output (NaN at mlen=256, where an additional HBM tile-alignment gap also applies). The native CLM-60M decoder GQA works because the frontend builds correct layouts every run.

Fix: pass explicit per-tensor tensor_layouts (like the MHA/RoPE tests) + hbm_addrs (like the linear test). GQA now validates with real per-element agreement across mlen 16/64/256, batch 1/2, seq 4→mlen (MRE ~0.003–0.02, no NaN). The previous mlen=256 NaN skip and short-sequence prune are removed.

vram_preload bf16 encoding + GQA softmax scale

• create_sim_env: write vram_preload as bfloat16 (the VRAM native type), not float16 — fp16 bytes were silently reinterpreted as bf16, corrupting the prestaged Q tensor (~170× magnitude shrink). Affects the attention tests (the only vram_preload consumers).
• flash_attention_gqa_test: fp_preload[1] = scale/0.25 — the codegen pre-divides by the fixed bmm_scale (0.25), so the test's softmax-scale slot must match (the old value under-scaled QK^T by 4× at mlen=16/256).

Emulator HBM out-of-bounds read

MemoryBacked::read returns zeros for out-of-capacity addresses instead of panicking. An H_PREFETCH that requests more rows than a tensor has (blen > seq_len) over-reads past the tensor end; those bytes land in unused padding VRAM and are never compared. Fixes the prior embedding-add @ mlen=64/blen=16 index-OOB crash (now passes in the matrix). write is left strict.

configurable.py — BEHAVIOR.CONFIG mirror

PlenaCompiler reads tile dims + HBM prefetch counts from BEHAVIOR.CONFIG; write_toml mirrors MLEN/VLEN/BLEN/HLEN/BROADCAST_AMOUNT + HBM_M/V_Prefetch_Amount + HBM_V_Writeback_Amount there. Fixed two silent-corruption bugs: HLEN defaulting to MLEN (broke packed attention), and the preload-stride vs prefetch-count mismatch that produced NaN in the first mlen/8 softmax rows at MLEN≥192.

MLEN≥256 (HBM tile-align)

At MLEN≥256 the compiler tile-aligns HBM bases; passing the compiler's real hbm_addrs to the writer fixes the zero-weights → NaN bug.

MLEN<64 (relax the 64-byte HBM burst model)

Emulator + HBM writer now handle tightly-packed sub-64 rows; no compiler changes, all no-ops at MLEN≥64:

• main.rs transfer_mx_from_hbm: relax is_multiple_of(8*64)→is_multiple_of(8); element read covers [addr, +len) via aligned-64-byte-word reads + byte-granular extraction (drops the addr.is_multiple_of(64) assert).
• main.rs transfer_mx_to_hbm: relax store assert + read-modify-write aligned words (no neighbour clobber).
• sim_env_utils.py: pack rows tight (hbm_row_bytes = element_row_bytes) and drop the per-tensor 64-byte tail pad. MemoryBacked capacity zero-covers the final aligned-block read.
• configurable.py: remove the MLEN%64 guards.

Native CLM-60M MLEN≥256 — NaN fix + column-block comparison

The native 1-layer decoder produced all-NaN output at MLEN≥256. Root cause: the score-mask/softmax-max constant fp_preload[2] was -inf, and the packed-attention column-mask buffer VRAM-aliases O_proj's tile-align padding rows; the RMSNorm over physical rows did (-inf)*0 = NaN. Fixed in companion PRs + one testbench change here:

• PLENA_Compiler#50: fp_preload[2] -inf → -6e4 (finite, float16-safe; exp(score-6e4)=0 still masks), and comparison_params.num_rows spans column blocks at the physical-row stride.
• PLENA_Tools#4: reorder_stride_mode col_block_stride/physical_rows + logical-width slice + compact golden-parser support.
• emulator_runner.py: thread physical_rows from comparison_params into compare_vram_with_golden.
• smolvlm2_256m.yaml: sub-64 RTL presets (native_16x16x4_b1, native_32x32x8_b1, native_32x32x4_b1).

quant shim → symlink

Replaced the hand-written quant/ shim with a symlink to PLENA_Tools/plena_quant (single source of truth); companion plena_quant re-export in PLENA_Tools#3 (merged).

Test plan

• Quick CI green (automatic): all 12 crucial cells pass on the post-merge emulator. Full matrix green (manual run): all 6 mlen{16,64,256}×batch{1,2} jobs + lint + rust-lint + build-and-test passed, with no tracked skips.
• Merge with main (refactor: decompose HBM↔SRAM DMA into memory byte-primitives + MX wrapper #62 DMA refactor) clean: sub-64 re-ported onto the new dma.rs/chunked.rs byte-primitives; emulator builds; sub-64 + mlen64 + mlen256 cells pass post-merge.
• GQA flash-attention genuinely validates (real K/V, no NaN) across mlen 16/64/256, batch 1/2, seq 4→mlen (MRE ~0.003–0.02; previously all-zero/vacuous).
• Native CLM-60M 1-layer decoder: MLEN=256 PASS (allclose 100%, MSE 8.5e-6; was all-NaN), MLEN=64 PASS 100% (non-regression).
• No regression across linear/ffn/softmax/rms_norm/layer_norm/embedding_add/rope at mlen 16/64/256, incl. batch=2 and seq up to mlen.
• embedding-add @ mlen=64/blen=16 (prior HBM index-OOB crash) now passes (emulator read fix).
• ruff format + ruff check + cargo fmt clean; submodules on merged main (PLENA_Tools#4, PLENA_Compiler#49+Apply repository-wide Ruff formatting and minimal lint cleanup for PR #49 #50).

FFN verification matrix — MLEN < 64 (end-to-end: compile + emulate)

mlen	blen	hidden	inter	up K×out	down K×out	result
8	4	8	16	1×2	2×1	✅ PASS
8	4	16	16	2×2	2×2	✅ PASS
16	8	32	32	2×2	2×2	✅ PASS
16	16	16	64	1×4	4×1	✅ PASS
16	4	64	16	4×1	1×4	✅ PASS
16	4	48	48	3×3	3×3	✅ PASS
32	4	32	64	1×2	2×1	✅ PASS
32	16	64	64	2×2	2×2	✅ PASS
32	32	64	128	2×4	4×2	✅ PASS
32	8	96	96	3×3	3×3	✅ PASS

Companion PRs

• Fix FFN down-projection HBM prefetch advance for multi-K-tile PLENA_Compiler#49 — FFN down-projection K-split (merged)
• Fix native decoder NaN at MLEN>=256 (finite score-mask) + col-block comparison PLENA_Compiler#50 — native MLEN≥256 NaN fix + col-block comparison_params (merged)
• Fix golden comparison for MLEN>=256 column-block-major output PLENA_Tools#4 — native MLEN≥256 column-block comparison + compact golden parser (merged)

Known remaining (tracked, separate work)

• SmolVLM2 full profile / vision branch: vision conv2d im2col (needs VLEN≥64) and K-tiled sub-64 im2col; the SmolVLM2 text decoder needs MLEN≥64 due to head_dim=64. Queued as a separate vision branch (fixes 1–4 + K-tiled im2col) for after this merges. CLM-60M at MLEN=256 already passes.

(Previously-tracked items now resolved: the flash-attention @ MLEN=256 NaN and the embedding-add @ mlen=64/blen=16 crash are both fixed and validated in the matrix.)