Fix GQA computation when fused_qkv=True

[dinodeep]

Description

This PR fixes a bug with fused_qkv=True and Grouped Query Attention (GQA) configuration. Previously, when performing GQA (num_query_heads > num_kv_heads), fused_qkv=True would ignore the num_kv_heads argument and assume that the user was configuring MHA and use num_query_heads for the number of query, key, and value heads.

This PR fixed that bug by adding support for fused_qkv with GQA so that fused_qkv will appropriately perform GQA when the user configures it.

Previously, when executing a Llama-70B model (which performs GQA with 64 query heads, and 8 kv heads) we would have the following optimized HLO generated when fused_qkv=True

%bitcast.17.0 = bf16[4,8196,16,128]{3,2,1,0} bitcast(%fusion-done.18), ...
%bitcast.18.0 = bf16[4,8196,16,128]{3,2,1,0} bitcast(%fusion-done.19), ...
%bitcast.19.0 = bf16[4,8196,16,128]{3,2,1,0} bitcast(%fusion-done.20), ...
...
%te_fused_attn_forward_ffi.0 = (bf16[4,8196,16,128]{3,2,1,0}, f32[4,16,8196,1]{3,2,1,0}, u32[2,4]{1,0}, u8[32]{0}) custom-call(%bitcast.17.0, %bitcast.18.0, %bitcast.19.0, ...

with the following (incorrectly reported) performance since the FLOP math that maxtext's metrics computes additionally reported the total FLOPs assuming GQA (since that is what the user configured) despite compiling and executing an MHA-based attention layer prior to the fix.

completed step: 9, seconds: 0.271, TFLOP/s/device: 860.509, Tokens/s/device: 30245.551, total_weights: 32784, loss: 5.065, lm_loss: 5.065, perplexity: 158.440
completed step: 10, seconds: 0.266, TFLOP/s/device: 876.609, Tokens/s/device: 30811.451, total_weights: 32784, loss: 4.504, lm_loss: 4.504, perplexity: 90.407
completed step: 11, seconds: 0.268, TFLOP/s/device: 871.081, Tokens/s/device: 30617.162, total_weights: 32784, loss: 4.062, lm_loss: 4.062, perplexity: 58.113
completed step: 12, seconds: 0.272, TFLOP/s/device: 856.683, Tokens/s/device: 30111.098, total_weights: 32784, loss: 3.730, lm_loss: 3.730, perplexity: 41.668
completed step: 13, seconds: 0.269, TFLOP/s/device: 865.395, Tokens/s/device: 30417.291, total_weights: 32784, loss: 3.484, lm_loss: 3.484, perplexity: 32.589
completed step: 14, seconds: 0.270, TFLOP/s/device: 863.312, Tokens/s/device: 30344.092, total_weights: 32784, loss: 3.295, lm_loss: 3.295, perplexity: 26.983

With the changes in this PR, we get the correct HLO that performs GQA (see num key/value heads < num query heads)

%bitcast.17.0 = bf16[4,8196,16,128]{3,2,1,0} bitcast(%fusion-done.19), ...
%bitcast.18.0 = bf16[4,8196,2,128]{3,2,1,0} bitcast(%fusion-done.20), ... 
%bitcast.19.0 = bf16[4,8196,2,128]{3,2,1,0} bitcast(%fusion-done.21), ...
...
%te_fused_attn_forward_ffi.0 = (bf16[4,8196,16,128]{3,2,1,0}, f32[4,16,8196,1]{3,2,1,0}, u32[2,4]{1,0}, u8[32]{0}) custom-call(%bitcast.17.0, %bitcast.18.0, %bitcast.19.0, ...

with correctly reported performance (GQA flops with GQA execution step time)

completed step: 10, seconds: 0.249, TFLOP/s/device: 935.746, Tokens/s/device: 32890.037, total_weights: 32784, loss: 4.630, lm_loss: 4.630, perplexity: 102.515
completed step: 11, seconds: 0.249, TFLOP/s/device: 937.688, Tokens/s/device: 32958.283, total_weights: 32784, loss: 4.227, lm_loss: 4.227, perplexity: 68.486
completed step: 12, seconds: 0.250, TFLOP/s/device: 932.114, Tokens/s/device: 32762.377, total_weights: 32784, loss: 3.923, lm_loss: 3.923, perplexity: 50.567
completed step: 13, seconds: 0.248, TFLOP/s/device: 941.151, Tokens/s/device: 33080.000, total_weights: 32784, loss: 3.698, lm_loss: 3.698, perplexity: 40.378
completed step: 14, seconds: 0.249, TFLOP/s/device: 935.093, Tokens/s/device: 32867.088, total_weights: 32784, loss: 3.525, lm_loss: 3.525, perplexity: 33.942

Tests

I validated this change by running Llama3-70b with fused_qkv prior to the new changes, inspected the optimized HLO profile, and observed MHA-based attention despite configuring GQA. Then, with these changes, I re-ran the same workflow and observed the correct computation graph.

Reproducer command:

XLA_FLAGS="--xla_dump_to=/opt/logs/fused_qkv_dump --xla_dump_hlo_as_text=true --xla_dump_hlo_pass_re=.*" mpirun -np 4 --allow-run-as-root python3 -m maxtext.trainers.pre_train.train     src/maxtext/configs/base.yml     model_name=llama3-70b     per_device_batch_size=1     max_target_length=8196     scan_layers=true     steps=15     remat_policy=minimal     use_iota_embed=true     logits_dot_in_fp32=false     enable_checkpointing=false     base_output_directory=local_train     dataset_path=local     dataset_type=synthetic     attention=cudnn_flash_te     hardware=gpu_multiprocess     run_name=test max_segments_per_seq=32 ici_tensor_sequence_parallelism=4 fused_mlp=true  fused_qkv=true base_num_decoder_layers=4 override_model_config=True

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[codecov]

Codecov Report

❌ Patch coverage is 80.95238% with 4 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/attentions.py	80.95%	2 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[gobbleturk]

can you explain this logic - I would think logical rules are still desired here even with PP

[dinodeep]

From my understanding, with pipeline parallelism, when logical_rules is None, the Flax logical_to_mesh_axes method will use the rules from the current nn.partitioning.axis_rules(...) context as opposed to the fixed rules provided by the config. These context rules are used with pipeline parallelism here for example. In these cases, we wouldn't want to use the config logical axis rules but the modified rules from the context manager instead.

This is a pattern that is done throughout other parts of the codebase such as here. Let me know if you would like me to consolidate them into a single helper