Modify the fused attn jax unit test case for head dim qk != head dim v

KshitijLakhani · KshitijLakhani · commit ed071aaaf4df · 2025-06-12T12:04:52.000-07:00
Signed-off-by: Kshitij Janardan Lakhani &lt;klakhani@nvidia.com&gt;
diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
@@ -978,7 +978,7 @@ def check_dqkv(primitive, reference, pad, idx):
             id="2-2048-1024-12-12-64-32-BF16-CROSS",
         ),
         pytest.param(
-            2, 2048, 2048, 12, 6, 64, 32, jnp.float16, id="2-2048-2048-12-6-64-32-FP16-GQA"
+            2, 2048, 2048, 12, 6, 128, 64, jnp.float16, id="2-2048-2048-12-6-128-64-FP16-GQA"
         ),
     ],
 )

Original file line number	Diff line number	Diff line change
`@@ -978,7 +978,7 @@ def check_dqkv(primitive, reference, pad, idx):`
`978`	`978`	`id="2-2048-1024-12-12-64-32-BF16-CROSS",`
`979`	`979`	`),`
`980`	`980`	`pytest.param(`
`981`		`- 2, 2048, 2048, 12, 6, 64, 32, jnp.float16, id="2-2048-2048-12-6-64-32-FP16-GQA"`
	`981`	`+ 2, 2048, 2048, 12, 6, 128, 64, jnp.float16, id="2-2048-2048-12-6-128-64-FP16-GQA"`
`982`	`982`	`),`
`983`	`983`	`],`
`984`	`984`	`)`