Fp8 prefill attn kernel integration by 1am9trash · Pull Request #18528 · sgl-project/sglang

1am9trash · 2026-02-10T08:39:23Z

Motivation

Replace the current BF16 prefill kernel with an FP8 kernel to improve performance.
The FP8 path uses aiter kernels:

mla_prefill_ps_asm_fwd
mla_reduce_v1 kernel

Modifications

Add the SGLANG_AITER_FP8_PREFILL_ATTN environment variable to control whether FP8 prefill is enabled (enabled by default).
Generate the required metadata at each step in init_forward_metadata().
In forward_extend(), cast BF16 inputs to FP8 (scale = 1) and call FP8 kernels.

Accuracy Tests

gsm8k cmd: python benchmark/gsm8k/bench_sglang.py --num-questions 2000 --parallel 2000 --port 9000
GPQA cmd: python -m sglang.test.run_eval --eval-name gpqa--port 9000 --num-examples 198 --max-tokens 65535 --repeat 8
LiveBenchCode: not supported in sglang test.

Prefill Type	GSM8K	GPQA (mean)	GPQA (max)	LiveBenchCode
BF16 Prefill	0.941	0.679	0.707
FP8 Prefill	0.945	0.669	0.697

Benchmarking and Profiling

server cmd:

SGLANG_AITER_FP8_PREFILL_ATTN=1 python3 -m sglang.launch_server --model-path DeepSeek-R1-MXFP4-Preview/ --tensor-parallel-size 8 --trust-remote-code --chunked-prefill-size 131072 --host 0.0.0.0 --port 9000 --log-requests --disable-radix-cache --mem-fraction-static 0.8 --max-running-requests 64 --kv-cache-dtype fp8_e4m3 --attention-backend aiter --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

gemini-code-assist · 2026-02-10T08:39:41Z

Summary of Changes

Hello @1am9trash, 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 an FP8 prefill attention kernel to enhance performance by replacing the existing BF16 kernel. It introduces an environment variable for controlling this feature, adds new metadata generation logic, and modifies the attention forward pass to utilize FP8 inputs and specialized aiter kernels for prefill operations.

Highlights

FP8 Prefill Attention Kernel Integration: The existing BF16 prefill kernel has been replaced with an FP8 kernel to improve performance, leveraging mla_prefill_ps_asm_fwd and mla_reduce_v1 aiter kernels.
Environment Variable Control: A new environment variable, SGLANG_AITER_FP8_PREFILL_ATTN, has been introduced to enable or disable the FP8 prefill functionality, with FP8 prefill enabled by default.
Metadata Generation for FP8: New methods (make_mla_prefill_ps_meta_data_buffer, make_mla_prefill_ps_meta_data) and updates to init_forward_metadata have been added to generate the necessary metadata for the FP8 prefill process.
Conditional FP8 Kernel Execution: The forward_extend function has been modified to conditionally cast BF16 inputs to FP8 and invoke the new FP8 prefill kernels when the feature is enabled.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

python/sglang/srt/layers/attention/aiter_backend.py
- Imported get_ps_metadata_info_v1, get_ps_metadata_v1, mla_prefill_ps_asm_fwd, and mla_reduce_v1 from aiter.ops.
- Added _use_fp8_prefill_attn boolean flag controlled by the SGLANG_AITER_FP8_PREFILL_ATTN environment variable.
- Implemented make_mla_prefill_ps_meta_data_buffer to allocate buffers for FP8 prefill metadata.
- Implemented make_mla_prefill_ps_meta_data to populate FP8 prefill metadata using get_ps_metadata_v1.
- Updated init_forward_metadata to conditionally prepare FP8 prefill metadata when the feature is enabled.
- Modified forward_extend to conditionally convert inputs to FP8 and call mla_prefill_ps_asm_fwd and mla_reduce_v1 for prefill attention.

Activity

The author has provided a detailed motivation for the change, focusing on performance improvement using FP8 kernels.
Benchmarking results, including a command and an image, have been shared to demonstrate the impact on inference speed.
Accuracy tests are marked as "To be updated," indicating pending work or results.
The checklist for code quality, tests, documentation, and benchmarks remains unchecked, suggesting further actions are required before merging.

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.

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 integrates an FP8 prefill attention kernel to improve performance, controlled by the SGLANG_AITER_FP8_PREFILL_ATTN environment variable. The changes include adding new helper functions to generate metadata for the FP8 kernel and updating the forward_extend method to use this new path.

My review identifies a critical bug in the FP8 casting logic that could lead to a NameError, and another critical issue with hardcoded FP8 range values that could cause incorrect behavior on certain hardware. I've also pointed out some minor issues related to an unused variable and naming consistency to improve code quality.

1am9trash · 2026-02-11T02:18:51Z

Co-authored-by: @kkHuang-amd

Co-authored-by: kkHuang-amd <wunhuang@amd.com>

1am9trash added 2 commits February 10, 2026 09:19

Add support for fp8 prefill kernel in aiter_backend

9467cc9

Change the flag name to SGLANG_AITER_FP8_PREFILL_ATTN (default=True)

50cd748

1am9trash requested review from Fridge003, Qiaolin-Yu, hebiao064, ispobock and merrymercy as code owners February 10, 2026 08:39

gemini-code-assist Bot reviewed Feb 10, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/attention/aiter_backend.py Outdated

Comment thread python/sglang/srt/layers/attention/aiter_backend.py Outdated

Comment thread python/sglang/srt/layers/attention/aiter_backend.py

Comment thread python/sglang/srt/layers/attention/aiter_backend.py Outdated

yctseng0211 added amd run-ci labels Feb 10, 2026

1am9trash added 2 commits February 10, 2026 16:48

Remove unused q_descale

ddd75c4

Fix name error

0f12a8b

HaiShaw requested changes Feb 10, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/attention/aiter_backend.py Outdated

1am9trash added 2 commits February 10, 2026 17:14

Remove q, k, v's clamp()

eaed4ae

Only use fp8 prefill attn kernel on mi35x

d5e9ead

Merge branch 'main' into fp8_prefill_integration

406b614

HaiShaw approved these changes Feb 11, 2026

View reviewed changes

HaiShaw merged commit a8eef53 into sgl-project:main Feb 11, 2026
100 of 132 checks passed

alphabetc1 pushed a commit to alphabetc1/sglang that referenced this pull request Feb 11, 2026

Fp8 prefill attn kernel integration (sgl-project#18528)

4744038

Co-authored-by: kkHuang-amd <wunhuang@amd.com>

Johnsonms pushed a commit to Johnsonms/sglang that referenced this pull request Feb 14, 2026

Fp8 prefill attn kernel integration (sgl-project#18528)

fb436b5

Co-authored-by: kkHuang-amd <wunhuang@amd.com>

magicYang1573 pushed a commit to magicYang1573/sglang that referenced this pull request Mar 9, 2026

Fp8 prefill attn kernel integration (sgl-project#18528)

46e1ab0

Co-authored-by: kkHuang-amd <wunhuang@amd.com>

Wangzheee pushed a commit to Wangzheee/sglang that referenced this pull request Mar 21, 2026

Fp8 prefill attn kernel integration (sgl-project#18528)

fa891a3

Co-authored-by: kkHuang-amd <wunhuang@amd.com>

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fp8 prefill attn kernel integration#18528

Fp8 prefill attn kernel integration#18528
HaiShaw merged 7 commits intosgl-project:mainfrom
1am9trash:fp8_prefill_integration

1am9trash commented Feb 10, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Feb 10, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

1am9trash commented Feb 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

1am9trash commented Feb 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Feb 10, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

1am9trash commented Feb 11, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

1am9trash commented Feb 10, 2026 •

edited

Loading