[Kernel] Phase 3+4: move JIT infra + operator groups into sglang.kernels (RFC #29630) - #31666
Merged
Merged
Conversation
BBuf
requested review from
DarkSharpness,
Edwardf0t1,
Fridge003,
HaiShaw,
HydraQYH,
Qiaolin-Yu,
Ying1123,
celve,
ch-wan,
hebiao064,
hnyls2002,
ispobock,
kpham-sgl,
merrymercy,
mickqian,
yingluosanqian and
yuan-luo
as code owners
July 18, 2026 11:30
BBuf
added
run-ci
run-ci-extra
bypass-maintenance
bypass-fastfail
enable-retry
github-actions
Bot
added
quant
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request refactors the codebase by migrating JIT kernel utility imports from sglang.jit_kernel.utils to sglang.kernels._jit. This change updates imports across various modules, including JIT kernels, operators, layers, models, and tests, ensuring consistent paths for shared JIT compilation and architecture detection utilities. I have no feedback to provide as there are no review comments to assess.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
…it (RFC sgl-project#29630) Relocates the shared JIT build/runtime infrastructure (compile pipeline, arch detection, deps, common utils) from sglang.jit_kernel.utils to an internal sglang.kernels.jit package. Behavior unchanged: - KERNEL_PATH resolves via find_spec('sglang.jit_kernel'), so csrc/ and include/ (per-op C++ sources + shared headers) and every JIT operator stay in sglang.jit_kernel; the build root is unaffected. - ~90 'sglang.jit_kernel.utils' import sites rewritten to 'sglang.kernels.jit'. - __main__ build CLI stays at sglang.jit_kernel (invocation unchanged). Operators migrate group-by-group in Phase 4. Namespace import stays metadata-only (invariant CLEAN); 68 kernels tests pass; lint clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
BBuf
force-pushed
the
bbuf/kernels-phase3-jit-infra
branch
from
July 21, 2026 13:04
b559762 to
f236843
Compare
BBuf
changed the title
BBuf
added a commit
to BBuf/sglang
that referenced
this pull request
Jul 21, 2026
…rnorm (RFC sgl-project#29630) Moves the JIT rmsnorm/qknorm operator wrappers from sglang.jit_kernel.norm to sglang.kernels.ops.layernorm._jit_norm and repoints the group's forward_jit to the local module. A compat shim at sglang.jit_kernel.norm re-exports everything (RFC Phase 4 step 8) so existing in-tree/external import paths keep working; removed in Phase 5. csrc stays in sglang.jit_kernel (KERNEL_PATH unchanged). Stacked on sgl-project#31666. Invariant CLEAN; 68 kernels tests pass; lint clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
BBuf
added a commit
to BBuf/sglang
that referenced
this pull request
Jul 21, 2026
…ps.activation (RFC sgl-project#29630) sglang.jit_kernel.activation -> sglang.kernels.ops.activation._jit_activation; group forward_jit (gated acts + relu2) now imports the local module. Compat shim at the old path re-exports everything (removed in Phase 5). csrc stays in sglang.jit_kernel. Stacked on sgl-project#31666. Invariant CLEAN; 68 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
BBuf
added a commit
to BBuf/sglang
that referenced
this pull request
Jul 21, 2026
… (RFC sgl-project#29630) dsv3_fused_a_gemm / dsv3_router_gemm -> kernels.ops.gemm._jit_*; KernelSpec JIT targets repointed; compat shims at old paths (removed Phase 5). csrc stays in jit_kernel. Stacked on sgl-project#31666. Invariant CLEAN; 68 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
BBuf
added a commit
to BBuf/sglang
that referenced
this pull request
Jul 21, 2026
…ps.kvcache (RFC sgl-project#29630) set_mla_kv_buffer -> kernels.ops.kvcache._jit_set_mla_kv_buffer; mla_buffer.py repointed; compat shim at old path (removed Phase 5). csrc stays in jit_kernel. Stacked on sgl-project#31666. Invariant CLEAN; 68 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
BBuf
added a commit
to BBuf/sglang
that referenced
this pull request
Jul 21, 2026
…nels.ops.quantization (RFC sgl-project#29630) per_tensor_quant_fp8 / per_token_group_quant_8bit(_v2) -> kernels.ops.quantization._jit_*; fp8_kernel.py imports + KernelSpec target repointed; compat shims at old paths (removed Phase 5). csrc stays in jit_kernel. Stacked on sgl-project#31666. Invariant CLEAN; 68 tests pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…it-infra # Conflicts: # python/sglang/jit_kernel/per_token_group_quant_8bit.py # python/sglang/kernels/ops/quantization/fp8_kernel.py # test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py # test/registered/jit/test_per_token_group_quant_8bit.py
…lidated, RFC sgl-project#29630) Folds the five Phase-4 group migrations into the Phase-3 infra PR so the whole jit-namespace move lands as one change: - layernorm (jit_kernel.norm), activation (jit_kernel.activation) - gemm (dsv3_fused_a_gemm / dsv3_router_gemm) - kvcache (set_mla_kv_buffer) - quantization (per_tensor_quant_fp8 / per_token_group_quant / _8bit_v2) Each operator wrapper moves to kernels.ops.<group>._jit_<module>; group forward_jit / KernelSpec JIT targets repoint to the local module; a compat shim at the old sglang.jit_kernel.<module> path re-exports everything (removed in Phase 5). csrc stays in sglang.jit_kernel (KERNEL_PATH via find_spec). Invariant CLEAN; 68 kernels tests pass; lint clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
BBuf
changed the title
…t#29630) kernels/jit/ is the future home of the whole sglang.jit_kernel tree (csrc / include / operators migrate here in later phases), so the shared infra lives in the utils/ subpackage: sglang.kernels.jit -> sglang.kernels.jit.utils. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This was referenced Jul 22, 2026
BBuf
added a commit
to BBuf/sglang
that referenced
this pull request
Jul 23, 2026
…yout - Add the jit/ subpackage (shared JIT build/runtime infra) to the Layout. - All 19 operator groups are populated; drop the stale "reserved placeholders / implementations not moved yet" wording (migration complete via sgl-project#31666/sgl-project#32015/ sgl-project#32045/sgl-project#32072). - Note sglang.jit_kernel removed; ops live at kernels.ops.<group>._jit_<op>; tests/benchmarks under test/registered/kernels/{ops,benchmark}/<group>/; shared test helpers in sglang.test.kernels. - Fix the fused-op parity test path (now under ops/layernorm/). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Zhylkaaa
pushed a commit
to Zhylkaaa/sglang
that referenced
this pull request
Jul 29, 2026
…els (RFC sgl-project#29630) (sgl-project#31666) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
jinzhenfan
pushed a commit
to jinzhenfan/sglang
that referenced
this pull request
Jul 29, 2026
…els (RFC sgl-project#29630) (sgl-project#31666) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This was referenced Aug 10, 2026
Chronostasys
pushed a commit
to MindLab-Research/sglang
that referenced
this pull request
Sep 1, 2026
…ernel Rebased onto the sglang.kernels layout per RFC sgl-project#29630 (sgl-project#31666/sgl-project#32015/sgl-project#32072) now that sglang.jit_kernel is retired. No logic changes: the patch now lives in kernels/jit/csrc/kvcacheio/hicache.cuh and kernels/ops/kvcache/hicache.py.
jakki-amd
pushed a commit
to jakki-amd/sglang
that referenced
this pull request
Sep 9, 2026
…els (RFC sgl-project#29630) (sgl-project#31666) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Atituiset
pushed a commit
to Atituiset/sglang
that referenced
this pull request
Sep 10, 2026
…els (RFC sgl-project#29630) (sgl-project#31666) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
RFC #29630 Phase 3 + Phase 4 (consolidated) — the whole JIT-namespace move lands as one PR. (Per request, the five per-group PRs are folded in here; #31693–#31697 are closed in favor of this.)
Phase 3 — shared infra
sglang/jit_kernel/utils/→sglang/kernels/jit/(compile pipeline, arch, deps, common). ~90 import sites rewritten.KERNEL_PATHresolves viafind_spec('sglang.jit_kernel'), socsrc/,include/, and the build CLI stay insglang.jit_kernel; build root unaffected.Phase 4 — operator groups →
kernels.ops.<group>normactivationdsv3_fused_a_gemm,dsv3_router_gemmset_mla_kv_bufferper_tensor_quant_fp8,per_token_group_quant,per_token_group_quant_8bit_v2Each wrapper →
kernels.ops.<group>._jit_<module>;forward_jit/ KernelSpec JIT targets repoint to the local module; a compat shim at the oldsglang.jit_kernel.<module>path re-exports everything (removed in Phase 5).csrcstays insglang.jit_kernel.Namespace import invariant CLEAN; 68 CPU kernels tests pass; lint clean. Merged with latest main.
Remaining JIT operators (diffusion / moe / attention — tangled subsystems) migrate later.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ⏳ Run #29888251333
Latest PR Test (Extra): ⏳ Run #29888251150