[Refactor] spec: Move platform dispatch and Triton kernels into `kernels/ops` by hnyls2002 · Pull Request #45504 · sgl-project/sglang

hnyls2002 · 2026-10-21T03:51:48Z

Cache locs: add fill_front_out_cache_loc_func and assign_draft_cache_locs_contiguous_func in kernels/ops/speculative/cache_locs.py, each dispatching to the Triton kernel or the sgl_kernel CPU op; move_accept_tokens_to_target_kvcache and the draft cache-loc assignment call them, and spec_utils drops its CPU-op imports, next_power_of_2, and the unused assign_extend_cache_locs forward
Move Triton kernels out of srt: compute_dflash_correct_drafts_and_bonus to kernels/ops/speculative/dflash.py (so dspark_accept no longer imports it from srt), table_qk_norm_rope_ to kernels/ops/attention/table_qk_norm_rope.py, and uno_tree.py to kernels/ops/speculative/uno_tree.py unchanged
Add reconstruct_indices_from_tree_mask_func in reconstruct_tree.py (Triton on XPU, the compiled op elsewhere); ngram_worker drops its per-platform imports
Move TreeMaskMode, the tree-build and greedy-verify Triton launchers, and verify_tree_greedy_func into kernels/ops/speculative/spec_tree.py, and split the platform dispatch out of build_tree_kernel_efficient into build_tree_kernel_efficient_func; eagle_utils keeps the mask allocation
Resolve the speculative sampling ops once per platform in kernels/ops/sampling/renorm.py (top_k_renorm_prob, top_p_renorm_prob) and kernels/ops/speculative/verify_sampling.py (tree_speculative_sampling_target_only, chain_speculative_sampling_triton), None where a platform has no kernel; DFlash and EAGLE import from there instead of each keeping its own table
Remove dead code: the unused cache-loc kernels align_evict_mask_to_page_size, get_src_tgt_cache_loc, get_target_cache_loc, and filter_finished_cache_loc_kernel with their spec_utils forwards (and the unused assign_req_to_token_pool forward), create_num_accept_tokens_filter, compute_ragged_extend_lengths, read_draft_checkpoint_gamma, and the decoupled-spec request-id codec, DraftControlInbox, and ReadyDraftControls, which only their tests used
Register the new entry points; test_next_verify_prep.py calls the draft wrapper over a smaller matrix; test_graph_pool_borrow.py patches the sampling op where eagle_utils binds it; drop test_decoupled_spec_io.py, whose schema has no production caller

No behavior change on CUDA. Elsewhere: on NPU, DFlash's target-probability renormalization uses the sgl_kernel_npu ops EAGLE already uses, with the torch_npu path as fallback; on CPU the draft assignment passes window_cols, which equals num_steps there.

CI States

Latest PR Test (Base): 🚫 Run #39556705672
Latest PR Test (Extra): ❌ Run #39556705491
Latest PR Test (AMD ROCm 10): ⏳ Run #39556705596

…kernel test

…upled-spec rid codec and inbox

hnyls2002 · 2026-10-21T06:46:22Z

/rerun-test test_next_verify_prep.py test_graph_pool_borrow.py test_move_accept_tokens_targets.py test_build_eagle_tree.py test_spec_ngram.py registered/spec/uno/test_uno.py test_spec_eagle.py test_spec_eagle_topk.py test_eagle_reject_sampling.py registered/spec/dflash/test_dflash.py test_dspark_kernel_parity.py test_dflash_overlap_hostsync.py test_renorm.py test_dflash_logits.py test_verify_mask.py test_kernels_namespace.py test_spec_utils_reexport_surface.py test_uno_tree_sparse_sampling.py test_eagle_worker_v2_topk1_fastpath.py test_dspark_accept_tail.py

hnyls2002 · 2026-10-21T06:46:28Z

/tag-and-rerun-ci

github-actions · 2026-10-21T06:46:47Z

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

cd test/ && python3 registered/kernels/ops/speculative/test_next_verify_prep.py
cd test/ && python3 registered/spec/test_spec_ngram.py
cd test/ && python3 registered/spec/uno/test_uno.py
cd test/ && python3 registered/kernels/ops/sampling/test_renorm.py

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

cd test/ && python3 registered/unit/model_executor/runner_utils/test_graph_pool_borrow.py
cd test/ && python3 registered/unit/spec/test_move_accept_tokens_targets.py
cd test/ && python3 registered/spec/utils/test_build_eagle_tree.py
cd test/ && python3 registered/spec/eagle/test_spec_eagle.py
cd test/ && python3 registered/spec/eagle/test_spec_eagle_topk.py
cd test/ && python3 registered/spec/dflash/test_dflash.py
cd test/ && python3 registered/spec/dspark/test_dspark_kernel_parity.py
cd test/ && python3 registered/unit/spec/test_dflash_overlap_hostsync.py

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

cd test/ && python3 registered/spec/eagle/test_eagle_reject_sampling.py

🚀 ubuntu-latest (7 tests): ✅ View workflow run

cd test/ && python3 registered/unit/spec/test_dflash_logits.py
cd test/ && python3 registered/unit/layers/attention/test_verify_mask.py
cd test/ && python3 registered/unit/kernels/test_kernels_namespace.py
cd test/ && python3 registered/unit/spec/test_spec_utils_reexport_surface.py
cd test/ && python3 registered/unit/spec/test_uno_tree_sparse_sampling.py
cd test/ && python3 registered/unit/spec/test_eagle_worker_v2_topk1_fastpath.py
cd test/ && python3 registered/unit/spec/test_dspark_accept_tail.py

… importable - hooks.py: torch 2.14's `_CUDAGraph` binding reads `keep_graph` in `__init__`, so overriding only `__new__` let a bare `torch.cuda.CUDAGraph()` drop its raw graph at capture end, and decode graph capture failed in cuGraphInstantiate with CUDA_ERROR_INVALID_VALUE. - lifted.py: defer annotations. `Tuple`, `Optional` and `torch` are bound only in the SGLang module each function is rebound into, so the first verify raised `NameError: Tuple`. - lifted.py: `eagle_sample` drops its inline sampler imports, which referenced `_is_npu` (removed from eagle_utils by sgl-project/sglang#45504), and uses eagle_utils' module names. - lifted.py: `_draft_extend_for_decode` passes `sampling_seed=` and `positions=` again, as current SGLang does. Tested on 8x B300, GLM-5.3-NVFP4, TP8, EAGLE 3/1/4, SGLang f4a2df3dac, all spec_step switches on.

…traverse — entries 102,103,119-supersede,120,194,195,261-264 3 markers consumed (eagle_utils RE-PLANTED; spec_utils re-plant; eagle_info VERIFIED-ABSORBED). B4 tree count now 12. PER-FILE: - eagle_utils.py (entry 103): RE-PLANTED the tensors-after-early-return graft (num_needed_tokens == 0 returns before any torch.tensor -- the 17F were the fail-fast + alloc-path rebuilding CPU tensors every round; topk>1 max_alloc_len back to the list-builtin max) MERGED with upstream's in-range mirrorless block: it rides AFTER the tensors now (it reads cur_kv_lens_cpu) with a getattr-guarded seq_lens_cpu read -- identical semantics for the real ScheduleBatch (attr always present), stale-stub-proof. - spec_utils.py (entries 102,120,195; 261 deduped): * entry 261 sgl-project#45491: DEDUPED with evidence -- git log 12dc623..2a6c204 proves 573b571 IN RANGE; its verify-window semantics SURVIVE upstream's later refactors (sgl-project#45504 kernel move + sgl-project#46990 typed views collapsed the triton assign_extend_cache_locs pair) as fill_front_out_cache_loc_func(num_correct_drafts+1, batch.out_cache_loc, ...) carrying sgl-project#45491's OWN comment text; the fork-hunks window-form is thus superseded-in-place (origin/main has neither old name -- the replacement landed within the same range). * entry 195 sgl-project#44317: RE-PLANTED the traverse_tree host walk (tolist + numpy bitmask view; int words) + the needs_vocab_mask() skip (the backend-side methods ALREADY landed upstream in constrained/ -- consumer half was the dropped one) + the leftover-grammar_mask clear on the grammars-present-but-unconstrained path (the merged tree had the clear on the SUCCESS path only -- the stale-mask class this entry exists for). traverse_tree 3F -> green. * entry 102 (PLE on the ReplaySSM branches) + entry 120 (circular-ring fold before a plain extend): RE-PLANTED, RETARGETED to sgl-project#45504/sgl-project#46980 reality: kernel import path attention.fla.* -> attention.linear.* (upstream's kernels regroup), and the fold reads mamba_pool.get_speculative_mamba2_params_all_layers() (upstream sgl-project#46980 deleted the req-pool delegate; same route the verify-commit uses now). test_replayssm_extend_fold collection 1err -> 2P (the documented 'pair': its sibling test_replayssm_ring_fold.py was PRUNED by upstream's own sgl-project#44690 (file gone; a stale .pyc remains) and the 'gdn cutedsl probe' err belongs to kernels/ops/attention (6 CUDA-import collection errs there -- B1/B5 zone, not this dir): the spec-dir population is now 0err. test_replayssm_verify_ple_commit 1F -> green (stub re-anchored: router part_view seam + PLE pools off-pin restored -- the entry-102 pin the upstream file rewrite dropped, present in the 309-tip test verbatim). - eagle_info.py (entry 194): VERIFIED-ABSORBED -- the call sites pass num_token_blocks=None + kv_lens_sum, and runner_ids.pack_read_stream auto-sizes with the fork's OWN helper (spec_kv_index_token_blocks, now hosted in kernels/ops/attention/kv_indices) -- the 'AUTO-SIZED upgrade' B3-4 recorded; identical tiling math. - Entry-119 supersession chain CLOSED as briefed: sgl-project#45435's upstream form (tails padded -1 + mamba_lazy_undo_prealloc_at_boundary + untracked_decode_batch_idx gate) is the tree's reality (B4-2 audit); entry 119's production half stays dead. - test_ngram_mamba_verify_update.py (B2-marked file, B2 marker INTACT): restored only the two PLE off-pins + the entry-102 comment that the 309-tip fork version carried and upstream's rewrite dropped -- my re-plant re-creates the condition they exist for. ACTUALS (damage -> after, CVD=9, prod venv, per-file): test_eagle_decode_allocation 17F -> ** 18P ** (16-matrix + capacity + 1) test_spec_utils_traverse_tree 3F -> green (file pin incl.) test_replayssm_extend_fold 1err -> ** 2P ** test_replayssm_verify_ple_commit 1F -> green DIR unit/spec 363P/20F/21S/238sub/2err -> ** 388P/0F/21S/238sub/0err ** (two consecutive full-dir re-runs EXACT; DoD(b): zero non-env F). Lint: ruff format + F821/E9/F401 clean on the six files. No GPU; no push. Co-Authored-By: OpenCode <noreply@opencode.ai>

hnyls2002 added 3 commits October 20, 2026 20:13

take post-verify kv move targets from the verify window; add test

5c415e3

reuse assign_extend_cache_locs over the verify window

450eb79

add fill_front_out_cache_loc_func for the move targets

0f6175e

hnyls2002 requested review from BBuf, DarkSharpness, HaiShaw, HydraQYH, Qiaolin-Yu, Ying1123, celve, merrymercy and yuan-luo as code owners October 21, 2026 03:51

github-actions Bot added the jit-kernel label Oct 21, 2026

Base automatically changed from lsyin/fix-ngram-move-read-done to main October 21, 2026 04:12

move draft cache-loc dispatch into a wrapper; drop dead forward; add …

5160b6e

…kernel test

hnyls2002 changed the title ~~[Refactor] spec: Add fill_front_out_cache_loc_func for post-verify move targets~~ [Refactor] spec: Move cache-loc platform dispatch into kernels/ops wrappers Oct 21, 2026

hnyls2002 added 5 commits October 20, 2026 21:21

merge main

eb3fced

move dflash and uno triton kernels into kernels/ops

9206526

wrap reconstruct_indices_from_tree_mask dispatch

fb46a73

move tree build and greedy verify dispatch into spec_tree

08ce18e

resolve spec sampling ops in one per-platform table

736fb6f

hnyls2002 requested a review from kpham-sgl as a code owner October 21, 2026 05:27

hnyls2002 changed the title ~~[Refactor] spec: Move cache-loc platform dispatch into kernels/ops wrappers~~ [Refactor] spec: Move platform dispatch and Triton kernels into kernels/ops Oct 21, 2026

hnyls2002 added 2 commits October 20, 2026 23:17

remove dead spec kernels, helpers, and forwards; drop the unused deco…

f594cd2

…upled-spec rid codec and inbox

trim comments; trim spec tests

dfb93ae

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

hnyls2002 merged commit 6efdd33 into main Oct 21, 2026
264 of 307 checks passed

hnyls2002 deleted the lsyin/spec-front-cache-loc-wrapper branch October 21, 2026 08:22

Eliaaazzz mentioned this pull request Oct 21, 2026

[Perf] spec: Drop the draft_probs buffer from tree_speculative_sampling_target_only #45616

Open

4 of 5 tasks

This was referenced Oct 21, 2026

[CI Monitor] Daily Report - 2026-10-21 bingxche/sglang-ci-bot#237

Open

[CI Monitor] Daily Report - 2026-10-22 bingxche/sglang-ci-bot#241

Open

This was referenced Oct 22, 2026

[NPU] Enable NGRAM speculative decoding on Ascend NPUs #38197

Closed

[Feat] [NPU] speculative: Enable NGRAM speculative decoding on Ascend NPUs #45797

Open

he-yufeng mentioned this pull request Oct 22, 2026

[Fix] ngram: Wire SGLANG_NGRAM_FORCE_GREEDY_VERIFY back into NGRAM verify #44147

Closed

3 of 5 tasks

gilfordting mentioned this pull request Oct 30, 2026

[Fix] sampling: deterministic top-p / top-k renorm when tp_size > 1 (TP rank divergence in DFlash/DSpark) #38565

Merged

5 tasks done

Navigation Menu

Repository navigation

[Refactor] spec: Move platform dispatch and Triton kernels into `kernels/ops` - #45504

[Refactor] spec: Move platform dispatch and Triton kernels into `kernels/ops`#45504
hnyls2002 merged 11 commits into
mainfrom
lsyin/spec-front-cache-loc-wrapper

hnyls2002 commented Oct 21, 2026 •

edited by github-actions Bot

Loading

Uh oh!

hnyls2002 commented Oct 21, 2026

Uh oh!

hnyls2002 commented Oct 21, 2026

Uh oh!

github-actions Bot commented Oct 21, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

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

Uh oh!

CI States

Uh oh!

hnyls2002 commented Oct 21, 2026

Uh oh!

hnyls2002 commented Oct 21, 2026

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

hnyls2002 commented Oct 21, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Oct 21, 2026 •

edited

Loading