# SINGLE_GPU_OUTPUT_START
Native JAX bf16 GQA + SWA:
Mean time: 5.109810829162598 ms

TE DotProductAttention GQA + SWA:
Mean time: 0.09856224060058594 ms
# SINGLE_GPU_OUTPUT_END

# MLA_OUTPUT_START
TE DeepSeek-style MLA head dimensions: q/k head dim=128, v head dim=64
Output shape=(2, 4096, 128, 64), dtype=bfloat16
Grad shapes=[(2, 4096, 128, 128), (2, 4096, 8, 128), (2, 4096, 8, 64)]
# MLA_OUTPUT_END
