diff --git a/iron/operators/mha/op.py b/iron/operators/mha/op.py
index 41c1bcd2..12f42994 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/O as (heads, S_q_pad, d) and K/V as
+        # (num_KV_heads, S_kv_pad * d); num_KV_heads == 0 means 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 29c5fd8a..c66521d9 100755
--- a/iron/operators/mha/test.py
+++ b/iron/operators/mha/test.py
@@ -80,3 +80,34 @@ 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",
+    [
+        (16384, 64, 8, 8, 2),
+        # num_kv_heads == 0 means plain MHA.
+        (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. The two must agree.
+    """
+    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())
+
+    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