[shoumikhin]

What is wrong today

In multi-query attention, all query heads share one key head and one value head (Gemma 2B, PaliGemma, pi0, pi0.5). Hugging Face code first copies the shared head once per query head (repeat_kv), then runs two batched matrix multiplications over the copies. TensorRT copies the repeated keys and values and runs its fused attention kernel, which is slow at head size 256 on some GPUs, such as A100 and H100.

What this change does

A new setting, fold_shared_kv_attention (default False), turns on a lowering pass. It folds the query heads into the rows and multiplies against the single head:

before: q [b, h, s, d] @ repeat(k [b, 1, n, d]).transpose  ->  [b, h, s, n]
after:  (q.reshape(b, 1, h*s, d) @ k.transpose).reshape(b, h, s, n)

trt = torch_tensorrt.dynamo.compile(program, arg_inputs=inputs, fold_shared_kv_attention=True)

Things to know:

• The math is the same, but plain matrix multiplications can be less accurate in 16-bit floating point (FP16). At head size 256 the largest error against 32-bit floating point (FP32) was 1.2 to 7 times that of the unfolded attention. use_fp32_acc=True brought it to 0.9 to 2.3 times.
• It needs exactly one key/value head and static shapes. Scaled dot product attention (SDPA) needs decompose_attention=True.
• It does not apply under torch.compile with use_distributed_mode_trace.
• If nothing matched, the pass says so in the log, once per graph.

It is off by default because a Jetson Orin runs the fused kernel faster.

The first commit is a small fix it builds on: replace_fused_rms_norm now sets the fake value (meta["val"]) on the nodes it creates, so a later pass can read shapes after an RMSNorm (root mean square normalization) layer.

Performance

8 query heads, 1 key/value head, head size 256, kernel time:

GPU	queries x keys	before	after
A100	968 x 968	235 us	109 us
A100	50 x 1018	140 us	44 us
H100	968 x 968	86 us	44 us
H100	50 x 1018	75 us	24 us

pi0.5 on an A100: 59.5 to 47.0 ms per action chunk. On a Jetson Orin Nano a denoising step got slower, 14.4 to 15.4 ms.

What was tested

The tests run the lowering pipeline and check the folded shapes, also after an RMSNorm. They cover decomposed SDPA (plain, use_fp32_acc, enable_gqa) and attention compiled with torch_tensorrt.dynamo.compile and with torch.compile, compared against PyTorch. They also check graphs that must stay unchanged: several key/value heads, dynamic lengths, and ten graphs that look like repeat_kv but must not be folded. Removing any one check of the pass makes at least one test fail.

They pass on a Jetson AGX Thor and on x86-64 with an H100 (43 passed), both with TensorRT 11.4 and a PyTorch 2.16 nightly.

[zewenli98]

It's a great addition. LGTM