gemv: add optional fused RMSNorm/LayerNorm prologue - #167
Open
atassis wants to merge 1 commit into
Open
Conversation
Adds an opt-in `prologue` parameter to GEMV. `prologue="none"` (default) is byte-identical to before; `prologue="rms"`/`"ln"` normalizes the shared B vector once per acquire, before it feeds every matvec call in that batch, so a decode-time norm+GEMV becomes one dispatch instead of a separate elementwise pass. Affine-free (gamma/beta fold into the weight matrix host-side), matching the existing rms_norm.cc/layer_norm.cc kernels' own convention. Both kernels take separate restrict-qualified in/out pointers, so normalizing b in place would violate their aliasing contract (the fix amd#135 used for gelu, an in-place wrapper, isn't available here without touching two kernel files for one call site). Instead each core gets a second L1 Buffer for the normalized vector; matvec reads from that instead of the raw acquire. Reuses rms_norm.cc/layer_norm.cc as-is, no kernel changes. rms_norm.cc and layer_norm.cc exist under both aie2 and aie2p with the same extern-C signature, so unlike the gelu epilogue this prologue is not NPU2-only. Test: test_gemv_norm_prologue compares against an A @ norm(B) golden across three shapes for each of rms/ln, with the standard latency/bandwidth metrics. rms_norm_ref/layer_norm_ref reproduce the kernels' f32 reduction math; checked against a hand-derived formula at rel-L2 < 1e-7 (pure numpy, no device). Not device-verified: I don't have hardware access in this session. Tolerances (rel_tol=0.05, abs_tol=1e-2) follow the gelu epilogue's precedent, not a measurement.
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
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.
Problem
GEMV has no way to fuse a norm ahead of the matvec. A decode-time
RMSNorm/LayerNorm on the input vector currently has to run as its own pass
before the GEMV dispatch, even though the vector is already resident on the
core that is about to consume it.
Fix
Adds an opt-in
prologueparameter toGEMV.prologue="none"(default) isbyte-identical to before.
prologue="rms"/"ln"normalizes the shared Bvector once per acquire in core_body, before it feeds every matvec call in
that batch, so the cost is O(K) per batch rather than O(M) if it were reapplied
per output tile.
Both
rms_norm_bf16_vectorandlayer_normtake separate restrict-qualifiedinput/output pointers, so calling them in place on the acquired B tile would
violate their aliasing contract (the same issue #135 hit for gelu, fixed there
with a new in-place kernel entry point). I did not want to touch two kernel
files for one call site, so instead each core gets a second L1
Bufferforthe normalized vector, and matvec reads from that. No changes to
rms_norm.cc/layer_norm.cc.
Both kernels exist under aie2 and aie2p with the same extern-C signature, so
unlike the gelu epilogue this prologue is not NPU2-only.
Affine-free: gamma/beta fold into the weight matrix host-side, the same
convention rms_norm.cc/layer_norm.cc already use (fixed gamma=1, beta=0).
Test/Evidence
test_gemv_norm_prologuecompares the fused output against anA @ norm(B)golden across three shapes for each of rms/ln, with the standard
latency/bandwidth metrics test_gemv_gelu also reports.
rms_norm_ref/layer_norm_refreproduce the kernels' f32 reduction math.Checked standalone against a hand-derived formula, pure numpy:
rel-L2 4.4e-08 (rms), 8.6e-08 (ln), on a 1024-element vector with nonzero mean.
Known limitations
Not device-verified. The IRON test suite opens the accelerator through the
aie_contextfixture, so this stops at the CPU-only checks above plus apy_compileof the four changed files.rel_tol=0.05, abs_tol=1e-2in thetest follow the gelu epilogue's precedent tolerances, not a measurement on
this shape.