Skip to content

fix(gdn): keep prefill beta in fp32 - #1579

Open
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:fix-gdn-prefill-beta-fp32
Open

sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:fix-gdn-prefill-beta-fp32

Conversation

@sufubao

@sufubao sufubao commented Sep 18, 2026

Copy link
Copy Markdown
Collaborator

问题

GDN prefill 的 fused_gdn_gating_kernel 在计算 fp32 sigmoid(b) 后,先转回输入 dtype,再写入 fp32 的 beta_output。输入为 bf16 时,这一步会丢失 beta 的低位精度;后续 chunk GDN 使用的是已舍入的值。相关背景:SGLang issue #38975

修改

  • 直接将 fp32 sigmoid 结果写入 fp32 beta_output,去掉多余的 bf16 转换。
  • 增加 GPU 回归测试,以 bf16 输入验证输出与 fp32 sigmoid 参考值一致。测试在修复前失败,修复后通过。

验证

  • python -m pytest unit_tests/common/basemodel/attention/linear/test_gdn.py unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_gdn_gating.py -q:17 passed(NVIDIA H200)。
  • pre-commit run --files lightllm/common/basemodel/triton_kernel/linear_att/fused_gdn_gating.py unit_tests/common/basemodel/triton_kernel/linear_att/test_fused_gdn_gating.py:Black、flake8 均通过。

@shihaobai
shihaobai self-requested a review September 18, 2026 15:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants