Skip to content

gemv: add optional fused RMSNorm/LayerNorm prologue - #167

Open
atassis wants to merge 1 commit into
amd:develfrom
atassis:nr-brick-norm-gemv-prologue
Open

gemv: add optional fused RMSNorm/LayerNorm prologue#167
atassis wants to merge 1 commit into
amd:develfrom
atassis:nr-brick-norm-gemv-prologue

Conversation

@atassis

@atassis atassis commented Aug 29, 2026

Copy link
Copy Markdown
Contributor

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 prologue parameter to GEMV. prologue="none" (default) is
byte-identical to before. prologue="rms"/"ln" normalizes the shared B
vector 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_vector and layer_norm take separate restrict-qualified
input/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 Buffer for
the 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_prologue compares the fused output against an A @ 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_ref reproduce 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_context fixture, so this stops at the CPU-only checks above plus a
py_compile of the four changed files. rel_tol=0.05, abs_tol=1e-2 in the
test follow the gelu epilogue's precedent tolerances, not a measurement on
this shape.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant