Separate q,k,v projection in sliding attention layers

[patrickji2014]

Description

In Gemma-4, existing model implementation use concatenate qkv weights in head_dim. when compute the qkv projections, this requires split the concatenated value back into q, k, v, which introduces some overhead from slice operations.

Perf metric

Batch size 818

With the change in this pr, 2.8% lift in the tok/s is observed.

============ Serving Benchmark Result ============
Successful requests:                     3272      
Failed requests:                         0         
Benchmark duration (s):                  32.38     
Total input tokens:                      3350528   
Total generated tokens:                  1636000   
Request throughput (req/s):              101.06    
Output token throughput (tok/s):         50529.19  
Peak output token throughput (tok/s):    78056.00  
Peak concurrent requests:                3272.00   
Total token throughput (tok/s):          154012.98  
---------------Time to First Token----------------
Mean TTFT (ms):                          5977.31   
Median TTFT (ms):                        5845.37   
P99 TTFT (ms):                           10696.54  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          51.29     
Median TPOT (ms):                        51.62     
P99 TPOT (ms):                           58.51     
---------------Inter-token Latency----------------
Mean ITL (ms):                           51.29     
Median ITL (ms):                         25.36     
P99 ITL (ms):                            201.25    
==================================================

baseline is run with xla_tpu_allow_conv_input_fusion_with_downcast_convert=true to
Fuse Sub-16 bytes Quantization Into Matmul

============ Serving Benchmark Result ============
Successful requests:                     3272      
Failed requests:                         0         
Benchmark duration (s):                  33.31     
Total input tokens:                      3350528   
Total generated tokens:                  1636000   
Request throughput (req/s):              98.23     
Output token throughput (tok/s):         49115.39  
Peak output token throughput (tok/s):    75799.00  
Peak concurrent requests:                3272.00   
Total token throughput (tok/s):          149703.70 
---------------Time to First Token----------------
Mean TTFT (ms):                          6332.15   
Median TTFT (ms):                        6260.21   
P99 TTFT (ms):                           11430.22  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          52.59     
Median TPOT (ms):                        52.97     
P99 TPOT (ms):                           60.38     
---------------Inter-token Latency----------------
Mean ITL (ms):                           52.59     
Median ITL (ms):                         51.39     
P99 ITL (ms):                            209.43    
==================================================

Batch Size 512

Baseline Total token throughput (tok/s): 137,241
With this pr, Total token throughput (tok/s): 139,284

Tests

One shot correctness check

curl http://localhost:8000/v1/chat/completions -X POST -H "Content-Type: application/json" -d '{"model": "google/gemma-4-26B-A4B-it", "
messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Who won the world series in 2020?"}], "max_toke
ns": 25, "temperature": 0.0}'


{"id":"chatcmpl-96d956199de39e85","object":"chat.completion","created":1786487110,"model":"google/gemma-4-26B-A4B-it","choices":[{"index":0,"message":{"role":"assistant","content":"The **Los Angeles Dodgers** won the 2020 World Series. They defeated the Tampa Bay Rays 4 games","refusal":null,"annotations":null,"audio":null,"function_call":null,"reasoning":null},"logprobs":null,"finish_reason":"length","stop_reason":null,"token_ids":null,"routed_experts":null}],"service_tier":null,"system_fingerprint":"vllm-0.26.1rc1.dev256+gf5bb701fa-tp2-c9a71676","usage":{"prompt_tokens":37,"total_tokens":62,"completion_tokens":25,"prompt_tokens_details":null},"prompt_logprobs":null,"prompt_token_ids":null,"prompt_text":null,"kv_transfer_params":null,"ec_transfer_params":null,"metrics":null}

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a Github issue, please include a link, e.g.,:
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[lk-chen]

LGTM, could you update PR title and description?