Merge pull request AI-Hypercomputer#3074 from AI-Hypercomputer:xibin/qkv

PiperOrigin-RevId: 865089762

Author: Google-ML-Automation

src/MaxText/configs/base.yml

@@ -455,6 +455,7 @@ logical_axis_rules: [
                       ["kv_lora_up_proj",[]],
                       ['norm', ['tensor', 'tensor_transpose']],
                       ['layers', 'stage'],
+                      ['qkv', []],
                       ['kv', []],
                       ['kv_head_dim', []],
                       ['cache_batch_prefill', []],