feat: openai oss attention sink support with trtllm-gen backend

[yyihuang]

Motivation

Modifications

Accuracy Test

python3 -m sglang.launch_server --model-path lmsys/gpt-oss-20b-bf16 --trust-remote-code --attention-backend trtllm_mha --enable-triton-kernel-moe --mem-fraction-static 0.7 --tp-size 8 --disable-cuda-graph --disable-hybrid-swa-memory

throughput
python3 benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1000 --parallel 1000

Accuracy: 0.473
Invalid: 0.047
Latency: 29.988 s
Output throughput: 14635.344 token/s

lm_eval

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
gsm8k	3	flexible-extract	8	exact_match	↑	0.8946	±	0.0085
		strict-match	8	exact_match	↑	0.2919	±	0.0125

Benchmark & Profiling

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Summary of Changes

Hello @yyihuang, 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 introduces support for the 'attention sink' mechanism within the trtllm-gen backend, specifically for OpenAI OSS models. The changes integrate a new attention_sink parameter into the multi-head attention forward passes, allowing for an additional value to be applied per head in the softmax denominator, which can be crucial for managing context and improving model stability or performance in long sequences.

Highlights

• Attention Sink Variable Initialization: Introduced an attention_sink variable in trtllm_mha_backend.py that conditionally holds layer.attention_sinks if layer.enable_attention_sink is true, otherwise None. This variable is intended to represent an additional value per head for the softmax denominator in attention calculations.
• Decode Path Integration: The newly defined attention_sink variable is now passed as a sink argument to the underlying attention computation function within the forward_decode method, enabling attention sink functionality during single token generation (lines 267-268).
• Extend Path Integration: Similarly, the attention_sink variable is also passed as a sink argument to the attention computation function within the forward_extend method, extending attention sink support to the context processing/prefill phase (lines 324-325).

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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 or fill out our survey to provide feedback.

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 adds support for attention sinks in the TRT-LLM MHA backend, which is a great feature.

I've identified a couple of areas for improvement regarding code maintainability:

• Code Duplication: There's an opportunity to refactor duplicated logic for determining the attention_sink value into a helper method. This aligns with the DRY (Don't Repeat Yourself) principle.
• Stale TODOs: The TODO comments related to adding attention sink support should be updated or removed now that the feature is being implemented to avoid confusion.

A significant concern is the lack of tests for this new functionality. The PR checklist for adding unit tests is also unchecked. Adding tests would be crucial to ensure correctness and prevent future regressions.

Please see my detailed comments on the code. Overall, good work on adding this feature.

[gemini-code-assist]

With the addition of the sink parameter, the TODO comment on the preceding line (266) appears to be at least partially outdated. It's good practice to update or remove TODO comments when the work is done to avoid confusion for future developers. If the nvfp4 scale factor part is still pending, the comment could be updated to reflect only that.

[gemini-code-assist]

This line is a duplicate of the logic in forward_decode at line 250. To improve maintainability and follow the DRY (Don't Repeat Yourself) principle, consider refactoring this into a private helper method, for example _get_attention_sink(self, layer).

[gemini-code-assist]

Similar to the forward_decode function, with the addition of this sink parameter, the TODO on line 323 seems outdated. Please consider updating or removing it to keep the code clean and maintainable.

[yyihuang]

add attention sink support on trtllm_mha backend, following #8822 or #8824