From cc2956f10acc4933b5161fffd76a6327455474c8 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Fri, 28 Aug 2026 02:38:54 +0300 Subject: [PATCH 1/2] gemv: add optional fused RMSNorm/LayerNorm prologue 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 #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. --- iron/operators/gemv/design.py | 43 +++++++++++++++-- iron/operators/gemv/op.py | 80 +++++++++++++++++++++++--------- iron/operators/gemv/reference.py | 19 ++++++++ iron/operators/gemv/test.py | 59 +++++++++++++++++++++++ 4 files changed, 175 insertions(+), 26 deletions(-) diff --git a/iron/operators/gemv/design.py b/iron/operators/gemv/design.py index 28d8b42c4..fb526990d 100644 --- a/iron/operators/gemv/design.py +++ b/iron/operators/gemv/design.py @@ -8,7 +8,7 @@ from aie.dialects.aie import T from aie.helpers.dialects.scf import _for as range_ from aie.helpers.taplib import TensorAccessPattern -from aie.iron import Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker +from aie.iron import Buffer, Kernel, ObjectFifo, Program, Runtime, TaskGroup, Worker """ Matrix-vector design @@ -37,6 +37,7 @@ def my_matvec( func_prefix="", verbose=False, epilogue="none", + prologue="none", ): if m_output is None: m_output = m_input @@ -101,6 +102,26 @@ def my_matvec( [np.int32, L1_C_ty], ) + # Optional norm applied to the shared B vector once per acquire, before it feeds every + # matvec call in this batch -- so the cost is O(K) per batch, not O(M) as re-normalizing + # per output tile would be. Both kernels take separate in/out pointers (restrict on + # both), so normalizing needs a second L1 buffer rather than an in-place call. + assert prologue in ("none", "rms", "ln") + norm_kernel = None + b_norm_bufs = [None] * cols + PROLOGUE_EPSILON = 1e-5 # matches layer_norm.cc's hardcoded epsilon + if prologue != "none": + norm_fn = "rms_norm_bf16_vector" if prologue == "rms" else "layer_norm" + norm_args = ( + [L1_B_ty, L1_B_ty, np.int32, np.float32] + if prologue == "rms" + else [L1_B_ty, L1_B_ty, np.int32] + ) + norm_kernel = Kernel( + f"{func_prefix}{norm_fn}", f"{func_prefix}{kernel_object}", norm_args + ) + b_norm_bufs = [Buffer(type=L1_B_ty, name=f"B_norm_{i}") for i in range(cols)] + A_L3L1_fifos = [ ObjectFifo(L1_A_ty, name=f"A_L3L1_{i}", depth=2) for i in range(cols) ] @@ -111,10 +132,20 @@ def my_matvec( ObjectFifo(L1_C_ty, name=f"C_L1L3_{i}", depth=2) for i in range(cols) ] - def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, gelu_kernel=None): + def core_body( + A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, norm_kernel, b_norm, gelu_kernel + ): one_idx = index.constant(1) for _ in range_(0xFFFFFFFF): # batch dim handled as part of this loop - b = B_L3L1_fifo.acquire(1) + b_raw = B_L3L1_fifo.acquire(1) + if norm_kernel is not None: + if prologue == "rms": + norm_kernel(b_raw, b_norm, K, PROLOGUE_EPSILON) + else: + norm_kernel(b_raw, b_norm, K) + b = b_norm + else: + b = b_raw # The kernel function computes m output rows; each core is responsible for (M/cols) output rows, so we need to call the kernel (M/cols)/m times. for i_idx in range_(M // m_output // cols): c = C_L1L3_fifo.acquire(1) @@ -138,8 +169,10 @@ def core_body(A_L3L1_fifo, B_L3L1_fifo, C_L1L3_fifo, matvec, gelu_kernel=None): B_L3L1_fifos[i].cons(), C_L1L3_fifos[i].prod(), matvec, - ] - + ([gelu_kernel] if epilogue == "gelu" else []), + norm_kernel, + b_norm_bufs[i], + gelu_kernel, + ], ) for i in range(cols) ] diff --git a/iron/operators/gemv/op.py b/iron/operators/gemv/op.py index 55c87239c..7e2c9defa 100644 --- a/iron/operators/gemv/op.py +++ b/iron/operators/gemv/op.py @@ -32,6 +32,10 @@ class GEMV(MLIROperator): # "none" (default) leaves the output unchanged; "gelu" applies GELU(tanh approx). # repr=False keeps operator/artifact names stable for the default path. epilogue: str = field(default="none", repr=False) + # Optional norm applied to the shared vector once before the matvec, fusing a + # RMSNorm/LayerNorm decode prologue into the same dispatch. Affine-free (gamma/beta + # fold into the weight matrix host-side, same convention the norm kernels already use). + prologue: str = field(default="none", repr=False) context: object = field(default=None, repr=False) _name_aliases: ClassVar[Dict[str, str]] = { @@ -63,6 +67,10 @@ def __post_init__(self): raise ValueError( f"gelu epilogue needs tile_size_output % 16 == 0 (got {self.tile_size_output})" ) + if self.prologue not in ("none", "rms", "ln"): + raise ValueError( + f"unknown prologue {self.prologue!r} (expected 'none', 'rms' or 'ln')" + ) MLIROperator.__init__(self, context=self.context) @@ -73,17 +81,25 @@ def name(self) -> str: # both would emit the same .mlir/.xclbin, and in a shared build dir a cached unfused # build can then satisfy the fused op (running the raw matvec with no activation). base = super().name - if self.epilogue == "none": - return base - return f"{base}_epi{self.epilogue}" + suffix = "" + if self.prologue != "none": + suffix += f"_pro{self.prologue}" + if self.epilogue != "none": + suffix += f"_epi{self.epilogue}" + return base + suffix @property def _kernel_link_file(self): - # With the gelu epilogue the core also links the gelu kernel, so the object becomes an - # archive of (matvec, gelu); the plain matvec stays a single object. - if self.epilogue == "gelu": - return f"gemv_{self.K}k_{self.kernel_vector_size}vs_gelu_kernels.a" - return f"gemv_{self.K}k_{self.kernel_vector_size}vs.o" + # With a prologue and/or the gelu epilogue the core links extra kernels, so the + # object becomes an archive; with neither, the plain matvec stays a single object. + if self.prologue == "none" and self.epilogue == "none": + return f"gemv_{self.K}k_{self.kernel_vector_size}vs.o" + suffix = "" + if self.prologue != "none": + suffix += f"_pro{self.prologue}" + if self.epilogue != "none": + suffix += f"_epi{self.epilogue}" + return f"gemv_{self.K}k_{self.kernel_vector_size}vs{suffix}_kernels.a" def get_mlir_artifact(self): mlir_verbose = getattr(self.context, "mlir_verbose", False) @@ -106,6 +122,7 @@ def get_mlir_artifact(self): "verbose": mlir_verbose, "kernel_object": self._kernel_link_file, "epilogue": self.epilogue, + "prologue": self.prologue, }, ), ) @@ -123,6 +140,24 @@ def get_kernel_artifacts(self): f"-DVEC_SIZE={self.kernel_vector_size}", ], ) + extra_objs = [] + if self.prologue != "none": + # rms_norm.cc / layer_norm.cc exist under both aie2 and aie2p with the same + # extern-C signature, so the prologue works on NPU1 and NPU2 alike. + norm_src = "rms_norm.cc" if self.prologue == "rms" else "layer_norm.cc" + extra_objs.append( + KernelObjectArtifact( + norm_src.replace(".cc", ".o"), + dependencies=[ + SourceArtifact( + self.context.base_dir + / "aie_kernels" + / get_kernel_dir() + / norm_src + ) + ], + ) + ) if self.epilogue == "gelu": # The gelu kernel lives in aie2p/gelu.cc, so the fused epilogue is NPU2-only. if get_kernel_dir() != "aie2p": @@ -130,20 +165,23 @@ def get_kernel_artifacts(self): "gemv gelu epilogue is only available on NPU2 (aie2p); " f"current kernel dir is {get_kernel_dir()!r}" ) - gelu_obj = KernelObjectArtifact( - "gelu.o", - dependencies=[ - SourceArtifact( - self.context.base_dir / "aie_kernels" / "aie2p" / "gelu.cc" - ) - ], - ) - return [ - KernelArchiveArtifact( - self._kernel_link_file, dependencies=[matvec_obj, gelu_obj] + extra_objs.append( + KernelObjectArtifact( + "gelu.o", + dependencies=[ + SourceArtifact( + self.context.base_dir / "aie_kernels" / "aie2p" / "gelu.cc" + ) + ], ) - ] - return [matvec_obj] + ) + if not extra_objs: + return [matvec_obj] + return [ + KernelArchiveArtifact( + self._kernel_link_file, dependencies=[matvec_obj] + extra_objs + ) + ] def get_arg_spec(self): batch_dim = (self.num_batches,) if self.num_batches > 1 else () diff --git a/iron/operators/gemv/reference.py b/iron/operators/gemv/reference.py index 140d8a0de..30951eadd 100644 --- a/iron/operators/gemv/reference.py +++ b/iron/operators/gemv/reference.py @@ -74,3 +74,22 @@ def gelu_tanh_approx(x): xf = np.asarray(x, dtype=np.float32) inner = 0.79788456 * (xf + 0.044715 * xf**3) return 0.5 * xf * (1.0 + np.tanh(inner)) + + +def rms_norm_ref(x, epsilon=1e-5): + """RMSNorm, matching aie_kernels/{aie2,aie2p}/rms_norm.cc's rms_norm_bf16_vector: + f32 sum-of-squares reduction, affine-free (gamma=1), computed in float32. + """ + xf = np.asarray(x, dtype=np.float32) + inv_rms = 1.0 / np.sqrt(np.mean(xf * xf) + epsilon) + return xf * inv_rms + + +def layer_norm_ref(x, epsilon=1e-5): + """LayerNorm, matching aie_kernels/{aie2,aie2p}/layer_norm.cc's layer_norm: + f32 mean/var reduction, affine-free (gamma=1, beta=0), computed in float32. + """ + xf = np.asarray(x, dtype=np.float32) + mean = np.mean(xf) + var = np.mean(xf * xf) - mean * mean + return (xf - mean) / np.sqrt(var + epsilon) diff --git a/iron/operators/gemv/test.py b/iron/operators/gemv/test.py index 052aff9f8..f9af419ac 100755 --- a/iron/operators/gemv/test.py +++ b/iron/operators/gemv/test.py @@ -10,6 +10,8 @@ generate_golden_reference, generate_golden_reference_batched, gelu_tanh_approx, + rms_norm_ref, + layer_norm_ref, ) from iron.common.device_utils import get_kernel_dir import numpy as np @@ -186,3 +188,60 @@ def test_gemv_gelu( print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") assert not errors, f"Test failed with errors: {errors}" + + +@pytest.mark.metrics( + Latency=r"Latency \(us\): (?P[\d\.]+)", + Bandwidth=r"Effective Bandwidth: (?P[\d\.e\+-]+) GB/s", + Throughput=r"Throughput: (?P[\d\.e\+-]+) GFLOP/s", +) +@pytest.mark.parametrize( + "M,K,num_aie_columns,tile_size_input,tile_size_output,prologue,norm_ref", + [ + pytest.param(128, 128, 1, 32, 128, "rms", rms_norm_ref), + pytest.param(2048, 8192, 1, 1, 2048, "rms", rms_norm_ref), + pytest.param(8192, 2048, 1, 4, 1024, "rms", rms_norm_ref), + pytest.param(128, 128, 1, 32, 128, "ln", layer_norm_ref), + pytest.param(2048, 8192, 1, 1, 2048, "ln", layer_norm_ref), + pytest.param(8192, 2048, 1, 4, 1024, "ln", layer_norm_ref), + ], +) +def test_gemv_norm_prologue( + M, + K, + num_aie_columns, + tile_size_input, + tile_size_output, + prologue, + norm_ref, + aie_context, +): + """GEMV with a fused RMSNorm/LayerNorm prologue vs an A @ norm(B) golden.""" + golden_ref = generate_golden_reference(M=M, K=K) + b_ref = golden_ref["B"].to(torch.float32).numpy() + b_norm = torch.from_numpy(norm_ref(b_ref).astype(np.float32)).to(torch.bfloat16) + c_pro = (golden_ref["A"] @ b_norm).to(torch.bfloat16) + + operator = GEMV( + M=M, + K=K, + num_aie_columns=num_aie_columns, + tile_size_input=tile_size_input, + tile_size_output=tile_size_output, + prologue=prologue, + context=aie_context, + ) + + input_buffers = {"matrix": golden_ref["A"].flatten(), "vector": golden_ref["B"]} + output_buffers = {"output": c_pro} + + errors, latency_us, bandwidth_gbps = run_test( + operator, input_buffers, output_buffers, rel_tol=0.05, abs_tol=1e-2 + ) + + print(f"\nLatency: {latency_us:.1f} us") + gflops = (2.0 * M * K) / (latency_us * 1e-6) / 1e9 + print(f"Throughput: {gflops:.6e} GFLOP/s") + print(f"Effective Bandwidth: {bandwidth_gbps:.6e} GB/s\n") + + assert not errors, f"Test failed with errors: {errors}" From 1b7d479122a78e76a20f9cd106437532c9063357 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Andr=C3=A9=20R=C3=B6sti?= Date: Tue, 8 Sep 2026 15:20:00 -0600 Subject: [PATCH 2/2] Update iron/operators/gemv/op.py --- iron/operators/gemv/op.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/iron/operators/gemv/op.py b/iron/operators/gemv/op.py index 7e2c9defa..01ef8f9a2 100644 --- a/iron/operators/gemv/op.py +++ b/iron/operators/gemv/op.py @@ -34,7 +34,7 @@ class GEMV(MLIROperator): epilogue: str = field(default="none", repr=False) # Optional norm applied to the shared vector once before the matvec, fusing a # RMSNorm/LayerNorm decode prologue into the same dispatch. Affine-free (gamma/beta - # fold into the weight matrix host-side, same convention the norm kernels already use). + # fold into the weight matrix host-side). prologue: str = field(default="none", repr=False) context: object = field(default=None, repr=False)