From 8726f85b9dbed338d9134dfb365859b7acb83617 Mon Sep 17 00:00:00 2001 From: Taimuraz Kaitmazov Date: Sun, 30 Aug 2026 15:44:26 +0300 Subject: [PATCH] mha: size the K/V arg specs by num_KV_heads, not num_heads design.py declares Q/O as (heads, S_q_pad, d) but K/V as (num_KV_heads, S_kv_pad * d). get_arg_spec used num_heads for all four, so under GQA K and V are over by num_heads/num_KV_heads: 8388608 against 2097152 on the GQA param already in test.py, a factor of 4. Two things hid it. num_kv_heads=0 normalizes to heads, which makes the old formula accidentally right for plain MHA, and run_test reads only spec.direction, never spec.shape. It surfaces in an OperatorSequence, where compilation asserts the MLIR arg count equals the computed one. --- iron/operators/mha/op.py | 14 +++++++++----- iron/operators/mha/test.py | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 46 insertions(+), 5 deletions(-) diff --git a/iron/operators/mha/op.py b/iron/operators/mha/op.py index 41c1bcd27..342a961d1 100644 --- a/iron/operators/mha/op.py +++ b/iron/operators/mha/op.py @@ -108,12 +108,16 @@ def get_kernel_artifacts(self): def get_arg_spec(self): seq_padding = self._calculate_seq_padding(self.seq_len, self.num_of_pipelines) - buffer_size = self.num_heads * self.d * seq_padding + # design.py declares Q and O as (heads, S_q_pad, d) but K and V as + # (num_KV_heads, S_kv_pad * d), and treats num_KV_heads == 0 as plain MHA. + kv_heads = self.num_KV_heads if self.num_KV_heads else self.num_heads + q_size = self.num_heads * self.d * seq_padding + kv_size = kv_heads * self.d * seq_padding return [ - AIERuntimeArgSpec("in", (buffer_size,)), # Q - AIERuntimeArgSpec("in", (buffer_size,)), # K - AIERuntimeArgSpec("in", (buffer_size,)), # V - AIERuntimeArgSpec("out", (buffer_size,)), # O + AIERuntimeArgSpec("in", (q_size,)), # Q + AIERuntimeArgSpec("in", (kv_size,)), # K + AIERuntimeArgSpec("in", (kv_size,)), # V + AIERuntimeArgSpec("out", (q_size,)), # O ] def _calculate_seq_padding(self, seq_len, num_pipeline=1): diff --git a/iron/operators/mha/test.py b/iron/operators/mha/test.py index 29c5fd8aa..8173e650f 100755 --- a/iron/operators/mha/test.py +++ b/iron/operators/mha/test.py @@ -80,3 +80,40 @@ def test_mha(seq_len, dim, num_heads, num_pipelines, num_kv_heads, aie_context): assert ( len(errors["O"]) <= max_acceptable_errors ), f"Test failed with {len(errors['O'])} errors (max allowable: {max_acceptable_errors})" + + +@pytest.mark.parametrize( + "seq_len,dim,num_heads,num_pipelines,num_kv_heads", + [ + # GQA: 8 query heads against 2 KV heads. K/V are a quarter of Q/O. + (16384, 64, 8, 8, 2), + # Standard MHA: num_kv_heads == 0 means num_kv_heads == num_heads. + (16384, 64, 1, 8, 0), + ], +) +def test_arg_spec_matches_design_shapes( + seq_len, dim, num_heads, num_pipelines, num_kv_heads +): + """get_arg_spec sizes the runtime buffers; design.py declares the MLIR arg types. + + Under GQA the two disagree on K and V by num_heads/num_KV_heads. Nothing catches + it today: run_test sizes input buffers from the supplied tensor rather than from + the spec, so only a fused OperatorSequence -- which asserts the MLIR arg count + equals the computed one -- would notice. + """ + op = MHA( + num_heads=num_heads, + seq_len=seq_len, + d=dim, + num_KV_heads=num_kv_heads, + num_of_pipelines=num_pipelines, + ) + q, k, v, o = (spec.shape[0] for spec in op.get_arg_spec()) + + # design.py: Q_ty is (heads, S_q_pad, d), KV_ty is (num_KV_heads, S_kv_pad * d). + pad = op._calculate_seq_padding(seq_len, num_pipelines) + kv_heads = num_kv_heads if num_kv_heads else num_heads + assert q == num_heads * pad * dim + assert o == num_heads * pad * dim + assert k == kv_heads * pad * dim + assert v == kv_heads * pad * dim