Fix: mistake sigmoid in kda by strgrb · Pull Request #17508 · sgl-project/sglang

strgrb · 2026-01-21T14:05:23Z

Motivation

This pr #17108 introduced a bug after resolving conflicts with main, that beta do an extra sigmoid in decode.
But when I test with gsm8k and mmlu, it gives the correct scores. So I use logprob to verify it:

I use an input of 1 2 3 5 8, and get the logprobs:

before pr 17108: [[-0.17581729590892792, 220, None], [-0.1355297565460205, 1718, None], [-0.2613707184791565, 220, None], [-0.02005852572619915, 2714, None], [-0.24853865802288055, 220, None], [-0.012261606752872467, 3587, None], [-0.23973554372787476, 220, None], [-0.007739911787211895, 3807, None], [-0.14101243019104004, 220, None], [-0.004367693793028593, 7173, None], [-0.25840112566947937, 220, None], [-0.005275614093989134, 17242, None], [-0.41415637731552124, 220, None], [-0.004919685889035463, 22851, None], [-0.12799102067947388, 220, None], [-0.003638078924268484, 38685, None], [-0.09586101025342941, 220, None], [-0.0013765154872089624, 30830, None], [-0.08789035677909851, 220, None], [-0.0007250064518302679, 48105, None], [-0.21781061589717865, 220, None], [-0.009240597486495972, 20320, None], [-0.000614216027315706, 22, None], [-0.06670233607292175, 220, None], [-0.0030783668626099825, 30049, None], [-0.0010446811793372035, 19, None], [-0.036415670067071915, 220, None], [-0.0008166077313944697, 39922, None], [-0.0005688241217285395, 16, None], [-0.030416714027523994, 220, None], [-0.00824251864105463, 46983, None], [-0.007040098775178194, 20, None]]
after pr 17108: [[-0.17581729590892792, 220, None], [-0.13807269930839539, 1718, None], [-0.22698462009429932, 220, None], [-0.02354123815894127, 2714, None], [-0.2504671812057495, 220, None], [-0.014055313542485237, 3587, None], [-0.2388245314359665, 220, None], [-0.006573718972504139, 3807, None], [-0.14944948256015778, 220, None], [-0.006166716571897268, 7173, None], [-0.16696849465370178, 220, None], [-0.005782779306173325, 17242, None], [-0.5210497975349426, 220, None], [-0.006434556096792221, 22851, None], [-0.15668046474456787, 220, None], [-0.004042905755341053, 38685, None], [-0.12107308954000473, 220, None], [-0.0016079128254204988, 30830, None], [-0.0959644466638565, 220, None], [-0.0013440155889838934, 48105, None], [-0.183536097407341, 220, None], [-0.0059707374311983585, 20320, None], [-0.0005921755800954998, 22, None], [-0.0742078647017479, 220, None], [-0.003984726965427399, 30049, None], [-0.0022481423802673817, 19, None], [-0.04476695507764816, 220, None], [-0.0012746074935421348, 39922, None], [-0.0008269704412668943, 16, None], [-0.03894393518567085, 220, None], [-0.00423392141237855, 46983, None], [-0.01138688251376152, 20, None]]
after this fix: [[-0.17581729590892792, 220, None], [-0.1355297565460205, 1718, None], [-0.26084673404693604, 220, None], [-0.019883597269654274, 2714, None], [-0.24637433886528015, 220, None], [-0.01390707679092884, 3587, None], [-0.2394082248210907, 220, None], [-0.0077369543723762035, 3807, None], [-0.13882087171077728, 220, None], [-0.004371016751974821, 7173, None], [-0.24931126832962036, 220, None], [-0.00554867135360837, 17242, None], [-0.43724241852760315, 220, None], [-0.005725057329982519, 22851, None], [-0.126996710896492, 220, None], [-0.0034029928501695395, 38685, None], [-0.09386846423149109, 220, None], [-0.0014907924924045801, 30830, None], [-0.08422954380512238, 220, None], [-0.0007373951375484467, 48105, None], [-0.19275186955928802, 220, None], [-0.008889622986316681, 20320, None], [-0.0006095697171986103, 22, None], [-0.06565475463867188, 220, None], [-0.002948106499388814, 30049, None], [-0.000979539705440402, 19, None], [-0.036493945866823196, 220, None], [-0.0008892634068615735, 39922, None], [-0.0006393534131348133, 16, None], [-0.030334722250699997, 220, None], [-0.007627647835761309, 46983, None], [-0.006255688611418009, 20, None]]

Modifications

Accuracy Tests

Benchmarking and Profiling

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-01-21T14:05:41Z

Summary of Changes

Hello @strgrb, 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 resolves a subtle bug within the 'KimiLinearAttentionBackend' where an erroneous 'sigmoid' activation was applied to the 'beta' parameter during the decode operation. The change ensures that the 'sigmoid' function is only invoked when appropriate, thereby restoring the correct logprobability calculations and maintaining the integrity of the model's output.

Highlights

Bug Fix for KimiLinearAttentionBackend: Addresses a bug introduced in a previous PR (17108) where the 'beta' calculation in the 'KimiLinearAttentionBackend' incorrectly applied an extra 'sigmoid' function during the decode phase.
Conditional Sigmoid Application: The 'sigmoid' activation is now conditionally applied to the 'beta' variable. It is only applied when the 'forward_mode' is not 'decode', ensuring the correct mathematical operation for different model states.
Logprob Verification: The fix has been validated by comparing logprobabilities, confirming that the corrected logic produces the expected numerical results, resolving the discrepancy caused by the extra 'sigmoid'.

🧠 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.

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 correctly fixes a bug where an extra sigmoid function was applied to beta during the decode phase in KDA. The change correctly moves the sigmoid application to be conditional on not being in decode mode. I've also suggested a small refactoring to improve code clarity by reducing duplication.

ispobock · 2026-01-21T14:34:04Z

Could you add an unit test for it?

ispobock · 2026-01-21T14:34:21Z

/rerun-stage stage-b-test-large-2-gpu

github-actions · 2026-01-21T14:34:58Z

✅ Triggered stage-b-test-large-2-gpu to run independently (skipping dependencies).

github-actions · 2026-01-21T14:35:04Z

🔗 View workflow run

strgrb · 2026-01-22T07:43:03Z

Could you add an unit test for it?

@ispobock I wrote a unittest and random generated data may always pass, and model dumped data works well. I will find a way to generate reasonable data and then commit the unittest.

strgrb · 2026-01-22T09:02:11Z

@ispobock Now unittest is pushed, and with this ut, we can see this bug:
without extra sigmoid:

abs_diff_out=tensor(0., device='cuda:0', dtype=torch.bfloat16), abs_diff_state=tensor(1.1921e-07, device='cuda:0')
.
----------------------------------------------------------------------
Ran 1 test in 3.620s

OK

with extra sigmoid

abs_diff_out=tensor(0.0125, device='cuda:0', dtype=torch.bfloat16), abs_diff_state=tensor(0.2750, device='cuda:0')
F
======================================================================
FAIL: test_kda_fused_sigmoid_gating_recurrent (test_kda_fused_sigmoid_gating_recurrent.TestKDAFusedSigmoidGatingRecurrent)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/ossfs/workspace/kda/test_kda_fused_sigmoid_gating_recurrent.py", line 115, in test_kda_fused_sigmoid_gating_recurrent
    self.assertTrue(torch.allclose(core_attn_out, core_attn_out_ref))
AssertionError: False is not true

----------------------------------------------------------------------
Ran 1 test in 3.613s

FAILED (failures=1)

and we should use fp32 ssm_states to avoid random input always pass

strgrb · 2026-01-22T09:11:46Z

/tag-run-ci-label

strgrb · 2026-01-22T14:12:11Z

/rerun-failed-ci

strgrb · 2026-01-23T01:34:18Z

/rerun-failed-ci

strgrb · 2026-01-23T02:31:50Z

/rerun-failed-ci

strgrb · 2026-01-23T03:11:27Z

/rerun-failed-ci

strgrb · 2026-01-23T03:13:50Z

@ispobock It seems most ci passed, and an amd ci failed with llada. I try it once more

ispobock · 2026-01-22T09:39:39Z

+
+
+@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
+class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):


Could you add this test to https://github.com/sgl-project/sglang/tree/main/test/registered/attention?
Also register it like

sglang/test/registered/attention/test_triton_attention_backend.py

Line 23 in f33022d

register_cuda_ci(est_time=200, suite="stage-b-test-large-1-gpu")

.

@ispobock Is this correct? I move the test here and register it to stage-b-test-large-1-gpu

strgrb · 2026-01-23T16:45:00Z

/rerun-failed-ci

strgrb · 2026-01-24T01:23:37Z

/rerun-failed-ci

fix kda

72760ac

gemini-code-assist Bot reviewed Jan 21, 2026

View reviewed changes

Comment thread python/sglang/srt/models/kimi_linear.py Outdated

accept suggestion

85ed42d

add ut

599d609

github-actions Bot added the run-ci label Jan 22, 2026

Merge branch 'main' into fix/kda

20c0cc1

ispobock reviewed Jan 23, 2026

View reviewed changes

move ut to registered

b9b4241

ispobock approved these changes Jan 23, 2026

View reviewed changes

ispobock merged commit 176da1b into sgl-project:main Jan 24, 2026
255 of 270 checks passed

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

Fix: mistake sigmoid in kda (sgl-project#17508)

0dc9f76



		@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA")
		class TestKDAFusedSigmoidGatingRecurrent(unittest.TestCase):

Conversation

strgrb commented Jan 21, 2026

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Jan 21, 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!

ispobock commented Jan 21, 2026

Uh oh!

ispobock commented Jan 21, 2026

Uh oh!

github-actions Bot commented Jan 21, 2026

Uh oh!

github-actions Bot commented Jan 21, 2026

Uh oh!

strgrb commented Jan 22, 2026

Uh oh!

strgrb commented Jan 22, 2026

Uh oh!

strgrb commented Jan 22, 2026

Uh oh!

strgrb commented Jan 22, 2026

Uh oh!

strgrb commented Jan 23, 2026

Uh oh!

strgrb commented Jan 23, 2026

Uh oh!

strgrb commented Jan 23, 2026

Uh oh!

strgrb commented Jan 23, 2026

Uh oh!

ispobock Jan 22, 2026

Choose a reason for hiding this comment

Uh oh!

strgrb Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock Jan 23, 2026

Choose a reason for hiding this comment

Uh oh!

strgrb commented Jan 23, 2026

Uh oh!

strgrb commented Jan 24, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants