[DFlash2] Keep the greedy selector walk in range for all-NaN score rows by rodamani · Pull Request #41737 · sgl-project/sglang

rodamani · 2026-09-29T19:34:40Z

Motivation

The greedy branch of _selector_walk_kernel (kernels/ops/speculative/dflash.py) takes best = tl.max(scores) and then the first offset where scores == best. If a whole score row is NaN, best is NaN, no element compares equal, and index falls through to top_k: one past the row, so the walk resolves to a neighbouring row's candidates.

Modifications

Greedy branch of _selector_walk_kernel now makes the same pick as torch.argmax (and _selector_walk_torch): the first NaN if the row has one, else the first maximum. index < top_k for every row, including all-NaN rows, and mixed-NaN rows match the reference.
test/registered/kernels/ops/speculative/test_lilicorr_cuda.py: test_greedy_path_matches_the_reference_on_nan_scores (all-NaN and mixed-NaN rows), comparing tokens and q rows from lilicorr_sample_path against _selector_walk_torch.

DSpark

DFlash-only: DSparkWorkerV2 does not use candidate_selector / selector_walk_triton, so DSpark is unaffected.

Accuracy Tests

CUDA test (test_lilicorr_cuda.py, registered on 1-gpu-large); results from CI. The kernel was also checked against _selector_walk_torch under TRITON_INTERPRET=1 for all-NaN, mixed-NaN, finite, tied and -inf rows (all match).

Speed Tests and Profiling

No hot-path change beyond the fix itself; not separately benchmarked.

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 and Merge Process

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

CI States

Latest PR Test (Base): ❌ Run #37372296257
Latest PR Test (Extra): ❌ Run #37372295607
Latest PR Test (AMD ROCm 10): ❌ Run #37372296288

…n all-NaN score row stays in range (sgl-project#105) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

rodamani · 2026-09-30T00:09:03Z

/rerun-test -c test_basic_sanity_dflash.py

github-actions · 2026-09-30T00:09:30Z

Results for /rerun-test -c test_basic_sanity_dflash.py:

🚀 1-gpu-5090 (1 test): ✅ View workflow run

cd test/ && python3 registered/core/test_basic_sanity_dflash.py

⛔ test/registered/unit/kernels/test_dflash_selector_walk_nan_row.py: File not found: test/registered/unit/kernels/test_dflash_selector_walk_nan_row.py

rodamani · 2026-09-30T00:20:32Z

/tag-and-rerun-ci

…riton interpreter Setting TRITON_INTERPRET inside the test module fails in CI's CPU suite once triton is already imported (tl.sum raises 'Cannot call @triton.jit'd outside of the scope of a kernel'). Move it next to the other speculative kernel tests and register it for base-b-kernel-unit on 1-gpu-large. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

…anager.devin.ai/proxy/github.com/modal-projects/sglang into rohan/up/dflash2-all-nan-selector-row

kpham-sgl

Fix is real; comments inline. Nit: the description's test path is wrong, and the CPU pass/fail numbers can't be right (test is CUDA-only).

Pick the first NaN if the row has one, else the first maximum, so mixed-NaN rows agree with _selector_walk_torch. Move NaN coverage into test_lilicorr_cuda.py as reference comparisons. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

rodamani · 2026-10-05T18:50:18Z

@kpham-sgl comments addressed here!

kpham-sgl · 2026-10-05T19:47:01Z

/rerun-test test/registered/kernels/ops/speculative/test_lilicorr_cuda.py test/registered/kernels/ops/speculative/test_dflash_domino.py test/registered/e2e/models/test_glm53_flash_b200.py

github-actions · 2026-10-05T19:52:35Z

⚠️ Rebase Required Before Re-run

A major update has landed on main. Your PR is diverged relative to required base commit fea3c088a141.

Re-run was not dispatched. What to do:

Rebase your branch onto the latest main and push again
Follow issue #21065 for context
CI-fix PRs may request the bypass-maintenance label to skip this check

rodamani · 2026-10-05T20:51:49Z

/rerun-test test/registered/kernels/ops/speculative/test_lilicorr_cuda.py test/registered/kernels/ops/speculative/test_dflash_domino.py test/registered/e2e/models/test_glm53_flash_b200.py

github-actions · 2026-10-05T21:02:30Z

🚀 1-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/kernels/ops/speculative/test_lilicorr_cuda.py
cd test/ && python3 registered/kernels/ops/speculative/test_dflash_domino.py

🚀 4-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_b200.py

dflash2: greedy selector walk picks the first unbeaten candidate so a…

15e7bf6

…n all-NaN score row stays in range (sgl-project#105) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

github-actions Bot added the jit-kernel label Sep 29, 2026

rodamani marked this pull request as ready for review September 29, 2026 21:15

rodamani requested review from BBuf, DarkSharpness, HaiShaw, HydraQYH, celve and yuan-luo as code owners September 29, 2026 21:15

Merge branch 'main' into rohan/up/dflash2-all-nan-selector-row

adc74c0

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 30, 2026

rodamani and others added 4 commits September 30, 2026 01:41

Merge branch 'rohan/up/dflash2-all-nan-selector-row' of https://git-m…

1ddc24e

…anager.devin.ai/proxy/github.com/modal-projects/sglang into rohan/up/dflash2-all-nan-selector-row

Merge branch 'main' into rohan/up/dflash2-all-nan-selector-row

b8c2bb4

Merge branch 'main' into rohan/up/dflash2-all-nan-selector-row

3da3eab

devin-ai-integration Bot mentioned this pull request Sep 30, 2026

Additional k3 sync fixes modal-projects/sglang#67

Merged

3 of 5 tasks

kpham-sgl reviewed Oct 1, 2026

View reviewed changes

Comment thread python/sglang/kernels/ops/speculative/dflash.py Outdated

Comment thread test/registered/kernels/ops/speculative/test_dflash_selector_walk_nan_row.py Outdated

Merge branch 'main' into rohan/up/dflash2-all-nan-selector-row

9a4e290

kpham-sgl approved these changes Oct 5, 2026

View reviewed changes

kpham-sgl merged commit c765f88 into sgl-project:main Oct 5, 2026
106 of 123 checks passed

devin-ai-integration Bot mentioned this pull request Oct 6, 2026

[Kimi K3] Stream tool-call arguments; DFlash selector + penalizer review follow-ups modal-projects/sglang#72

Open

3 of 5 tasks

Navigation Menu

Repository navigation

[DFlash2] Keep the greedy selector walk in range for all-NaN score rows - #41737

[DFlash2] Keep the greedy selector walk in range for all-NaN score rows#41737
kpham-sgl merged 8 commits into
sgl-project:mainfrom
modal-projects:rohan/up/dflash2-all-nan-selector-row

rodamani commented Sep 29, 2026 •

edited by github-actions Bot

Loading

Uh oh!

rodamani commented Sep 30, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026 •

edited

Loading

Uh oh!

rodamani commented Sep 30, 2026

Uh oh!

kpham-sgl left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

rodamani commented Oct 5, 2026

Uh oh!

kpham-sgl commented Oct 5, 2026

Uh oh!

github-actions Bot commented Oct 5, 2026

Uh oh!

rodamani commented Oct 5, 2026

Uh oh!

github-actions Bot commented Oct 5, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

rodamani commented Sep 29, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

DSpark

Accuracy Tests

Speed Tests and Profiling

Checklist

Review and Merge Process

CI States

Uh oh!

rodamani commented Sep 30, 2026

Uh oh!

github-actions Bot commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

rodamani commented Sep 30, 2026

Uh oh!

kpham-sgl left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

rodamani commented Oct 5, 2026

Uh oh!

kpham-sgl commented Oct 5, 2026

Uh oh!

github-actions Bot commented Oct 5, 2026

⚠️ Rebase Required Before Re-run

Uh oh!

rodamani commented Oct 5, 2026

Uh oh!

github-actions Bot commented Oct 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

rodamani commented Sep 29, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Sep 30, 2026 •

edited

Loading

kpham-sgl left a comment •

edited

Loading

github-actions Bot commented Oct 5, 2026 •

edited

Loading