Fix FFN down-projection HBM prefetch advance for multi-K-tile#49
Merged
Conversation
The down-projection weight prefetch in `_ffn_asm_with_loops` and
`_ffn_asm_fused_up_gate` used `_load_large_int(a_actual, mlen*hidden_size)`
to advance the HBM offset register between K-tile prefetches. This
overwrote the register instead of adding to it, so:
* At K_tiles >= 3 the prefetch for tiles 2..N-1 all read from
HBM[base + mlen*hidden_size] instead of advancing by that stride
each iteration.
* At K_tiles == 2 combined with multi-outer-iter (out_size > mlen),
K-tile 1 of outer iters > 0 read from the wrong absolute address
because the per-outer-iter base offset got dropped.
The up-projection in the same function already used `_addi_large_int`
correctly; the down path now matches. Fix applied at all three
down-projection prefetch sites (lines ~1333, 1699, 1777).
Also: `_emit_ffn_projection_unrolled` switched the K-split V_ADD_VV
accumulator loop count from integer division to `math.ceil` so the
tail vector when `output_elements` is not a multiple of `vlen` is no
longer truncated; `import math` added.
Verified with ATen FFN testbench at MLEN=64 and MLEN=128 across
hidden/inter combinations exercising single-tile, multi-K-tile, and
multi-outer-iter codegen paths (all pass).
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
booth-algo
added a commit
to AICrossSim/PLENA_Simulator
that referenced
this pull request
May 28, 2026
…+ BEHAVIOR.CONFIG After PR #60's config-chain fix exposed the emulator's true MLEN, three ATen tests still failed because their golden references and per-build test plumbing weren't matching the emulator's actual precision and HBM layout. This commit lands the remaining fixes plus the down-side of the K-split bug (in the linked PLENA_Compiler PR). - `golden.py`: `golden_rope`, `golden_rms_norm`, `golden_layer_norm` now pre-quantise their inputs through the HBM MXFP8 path via `_load_to_vector_fp(...)` before BF16 compute, matching the emulator's HBM (MXFP8) → VRAM (BF16) staging. Previously they only quantised at the vector precision and the deltas showed up as 0% match. - `rope_test.py`: pass an explicit compact `tensor_layouts` and compiler-derived `hbm_addrs` (read from `prog._compiler.get_hbm_layout(name).hbm_base_addr`). Without this, a stale `tensor_layouts.json` (or default padding) shifted QROT/COS/SIN off the addresses the prefetch ISA reads. - `norm_test.py`: install `fp_preload[eps_offset]=eps`, `fp_preload[reci_hid_offset]=1/hidden` at the offsets the chosen norm op actually reads (rms uses 0/1, layer uses 1/2). Without preload the emulator divided by zero and emitted `inf`. Default `hidden = mlen` to avoid mixing the norm test with the FFN K-split path. - `flash_attention_gqa_test.py`: set `args.hlen = h_qkv` before `setup_hw` so the per-build HLEN matches the packed-attention constraint `broadcast_amount * hlen == mlen`. - `flash_attention_mha_test.py`: default `seq_len` / `head_dim` to MLEN so the comparison-row layout matches at MLEN=128 (previously the 64×64 defaults at MLEN=128 produced a 4096-vs-16384 length mismatch in `slice_rows`). - `configurable.py`: `HardwareConfig.write_toml` now also writes the tile dimensions under a `BEHAVIOR.CONFIG` section. `PlenaCompiler`'s `_behavior_config_value` reads from `mode="BEHAVIOR"` (see `PLENA_Compiler/utils/load_config.py`); without the section, HLEN silently defaulted to MLEN and broke packed attention codegen. - `PLENA_Compiler` submodule: bumped to the fix branch with the `_load_large_int` → `_addi_large_int` correction in the down- projection HBM prefetch (see AICrossSim/PLENA_Compiler#49). Sweep result at MLEN=64 and MLEN=128 (linear, rms_norm, layer_norm, fpvar_softmax, embedding_add, rope, flash_attn_gqa, flash_attn_mha, ffn, norm rms, norm layer): all pass. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
7 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The down-projection weight prefetch in
_ffn_asm_with_loopsand_ffn_asm_fused_up_gateused_load_large_int(a_actual, mlen*hidden_size)to advance the HBM offset register between K-tile prefetches. This overwrote the register instead of adding to it, causing several silent miscomputations:K_tiles >= 3the prefetch for K-tiles 2..N-1 all read fromHBM[base + mlen*hidden_size]instead of advancing by that stride each iteration.K_tiles == 2combined with multi-outer-iter (out_size > mlen), K-tile 1 of outer iters > 0 read from the wrong absolute address because the per-outer-iter base offset got dropped.The up-projection in the same function already used
_addi_large_intcorrectly; the down path now matches at all three call sites.Also:
_emit_ffn_projection_unrolled's K-split V_ADD_VV accumulator loop count is nowmath.ceil(output_elements / vlen)instead of integer division so the tail vector whenoutput_elementsis not a multiple ofvlenis no longer dropped.Test plan
Verified with the ATen FFN testbench in PLENA_Simulator (PR AICrossSim/PLENA_Simulator#60). Matrix spans single-tile, multi-K-tile, multi-outer-iter, and BLEN 4–64:
¹ MLEN=256 additionally required the HBM tile-align fix on the simulator side (PR AICrossSim/PLENA_Simulator#60); with both, FFN passes at MLEN=256.