Enable Llama 4 + TRTLLM MHA

[b8zhong]

Motivation

I don't know why it was blocked originally from being used (TBD maybe it was either not in yet, or was untested). A lot better than triton (currently)

Modifications

Unblock it.

python -m sglang.launch_server \
                                                                                                                                                         --model-path /root/.cache/huggingface/hub/models--nvidia--Llama-4-Maverick-17B-128E-Instruct-FP8/snapshots/e91306aa6edc0a8b551b1815027163646638d252 \
                                                                                                                             --tp 8 \                                              
                                                                                                 --attention-backend trtllm_mha \                                                  
                                                                     --trust-remote-code          

python3 -m sglang.bench_serving --backend sglang --num-prompts 64 --dataset-name random --random-input-len 1024 --random-output-len 1024 --random-range-ratio 1 --max-concurrency=8 --flush-cache
Before:

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 8         
Successful requests:                     64        
Benchmark duration (s):                  53.09     
Total input tokens:                      65536     
Total input text tokens:                 65536     
Total input vision tokens:               0         
Total generated tokens:                  65536     
Total generated tokens (retokenized):    65673     
Request throughput (req/s):              1.21      
Input token throughput (tok/s):          1234.33   
Output token throughput (tok/s):         1234.33   
Total token throughput (tok/s):          2468.67   
Concurrency:                             8.00      
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   6632.70   
Median E2E Latency (ms):                 6634.04   
---------------Time to First Token----------------
Mean TTFT (ms):                          144.41    
Median TTFT (ms):                        143.07    
P99 TTFT (ms):                           162.07    
---------------Inter-Token Latency----------------
Mean ITL (ms):                           6.34      
Median ITL (ms):                         6.34      
P95 ITL (ms):                            6.74      
P99 ITL (ms):                            7.00      
Max ITL (ms):                            16.29     
==================================================

After:

============ Serving Benchmark Result ============
Backend:                                 sglang    
Traffic request rate:                    inf       
Max request concurrency:                 8         
Successful requests:                     64        
Benchmark duration (s):                  48.20     
Total input tokens:                      65536     
Total input text tokens:                 65536     
Total input vision tokens:               0         
Total generated tokens:                  65536     
Total generated tokens (retokenized):    65785     
Request throughput (req/s):              1.33      
Input token throughput (tok/s):          1359.59   
Output token throughput (tok/s):         1359.59   
Total token throughput (tok/s):          2719.18   
Concurrency:                             8.00      
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   6022.57   
Median E2E Latency (ms):                 6015.51   
---------------Time to First Token----------------
Mean TTFT (ms):                          139.21    
Median TTFT (ms):                        137.55    
P99 TTFT (ms):                           155.55    
---------------Inter-Token Latency----------------
Mean ITL (ms):                           5.75      
Median ITL (ms):                         5.74      
P95 ITL (ms):                            5.80      
P99 ITL (ms):                            5.91      
Max ITL (ms):                            61.09     
==================================================

Metric	Before (--attention-backend triton)	After (--attention-backend trtllm_mha)	Δ	% Change
Request Throughput (req/s)	1.21	1.33	+0.12	+9.9%
Input Token Throughput (tok/s)	1234.33	1359.59	+125.26	+10.2%
Total Token Throughput (tok/s)	2468.67	2719.18	+250.51	+10.2%
Mean E2E Latency (ms)	6632.70	6022.57	−610.13	−9.2%
Mean ITL (ms)	6.34	5.75	−0.59	−9.3%

Accuracy Tests

Acc is good
After:

python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 50
/usr/local/lib/python3.12/dist-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
  import pynvml  # type: ignore[import]
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1319/1319 [01:19<00:00, 16.56it/s]
Accuracy: 0.933
Invalid: 0.000
Latency: 79.984 s
Output throughput: 1709.423 token/s

[gemini-code-assist]

Summary of Changes

Hello @b8zhong, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request integrates and activates the trtllm_mha attention backend for Llama 4 models, addressing a previous blockage. The primary motivation is to leverage the performance benefits offered by trtllm_mha over the existing triton backend. Comprehensive benchmarks demonstrate significant improvements in throughput and reduced latencies, while accuracy tests confirm that the model's performance remains robust. This change aims to enhance the efficiency and responsiveness of Llama 4 model serving.

Highlights

• Enabled TRTLLM MHA for Llama 4: The trtllm_mha attention backend has been explicitly enabled and added as a valid option for Llama 4 models, allowing for its utilization in server configurations.
• Significant Performance Improvements: Benchmarking against the triton backend shows a notable increase in request throughput (9.9%), input token throughput (10.2%), and total token throughput (10.2%). Mean end-to-end latency and inter-token latency also saw reductions of 9.2% and 9.3% respectively.
• Accuracy Maintained: Accuracy tests on the GSM8K dataset confirm that enabling trtllm_mha does not negatively impact the model's performance, achieving an accuracy of 0.933.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request enables the trtllm_mha attention backend for Llama 4 models, which shows performance improvements according to the benchmarks. The change is simple and correct. My review includes one point of feedback to update an outdated assertion message related to this change to ensure consistency.

[Fridge003]

Is trtllm_mha+llama4 runnable on blackwell or hopper or both of them?

[b8zhong]

Both (since there is fa3 on hopper that is probably faster so i did not change default)

[Fridge003]

Can you change the default backend to trtllm_mha for Llama-4, and add a CI test for it?

Or CI test can be added after B200 CI is reenabled

[b8zhong]

@Fridge003 Done