[BBuf]

Motivation

Port the standalone kernels from the Kimi K3 Day0 work to main first. Keeping
the kernels and their direct tests separate makes the reusable pieces reviewable
before the more invasive model, scheduler, and serving integration changes.

Modifications

• Add the shared JIT/TMA support and generic prerequisite kernels used by Kimi
  K3.
• Port the standalone KDA kernels, including ReplaySSM, packed decode, and the
  NVIDIA/PTX prefill implementations.
• Port the Kimi K3 compute kernels for SiTU, quantization, MLA output gating,
  attention residuals, MTP, and image preprocessing.
• Port the multi-GPU fused all-reduce, GEMM + all-gather, GEMM + all-reduce, and
  sequence-parallel collective kernels.
• Keep four hardware-grouped registered test entrypoints for H100 prerequisite,
  B200 compute, B200 collective, and Blackwell prefill correctness.
• Remove tuner/benchmark scaffolding and Day0 runtime dispatch/integration
  changes; keep model, scheduler, speculative decoding, multimodal, and serving
  runtime integration out of this PR.

Accuracy Tests

• pre-commit run --all-files
• git diff --check
• python3 -m compileall -q python/sglang/kernels test/registered/kernels test/registered/jit
• After consolidation, all 9 retained H100 tests passed in 12.91s. The suite
  covers MLA BF16/FP8 scatter and concat, generic prerequisite kernels, and the
  representative ReplaySSM ring parity path.
• The 13 Blackwell-only test methods imported successfully on H100 and skipped
  through their architecture guards.
• The SM100/SM103-specific KDA, compute, and collective tests are registered for
  the B200/GB300 CI jobs. They were not run locally because no Blackwell machine
  was reachable; H100 results are not presented as Blackwell validation.

Speed Tests and Profiling

Not applicable to this extraction. It does not enable the Kimi K3 serving path,
and no benchmark files are added. Performance validation belongs with the
follow-up runtime integration after the Blackwell kernel jobs pass.

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 applicable accuracy results; speed benchmarking is deferred because
  this PR does not enable a serving path.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Run the registered B200 and GB300 kernel jobs to complete the
   architecture-specific acceptance matrix.

---

CI States

Latest PR Test (Base): ❌ Run #30679474553
Latest PR Test (Extra): ✅ Run #30679474497

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[DarkSharpness]

How does this compile? Should we move this to AOT kernel first?

[BBuf]

It builds through torch.utils.cpp_extension.load (ninja), not the sglang JIT — the only kernel under kernels/jit/ that does, cached separately in ~/.cache/torch_extensions.

Measured on B300: 109.2 s cold, 5.0 s warm, and the .so carries only sm_103a, so B200/GB200 get no kernel image although the cookbook lists B200. Agreed on AOT.

[DarkSharpness]

Is this already moved to AOT? We can defer the migration of this kernel? I remember torch.utils.cpp_extension.load is not concurrent-safe?

[DarkSharpness]

Is this related to this PR?

[BBuf]

You are right to ask. The file is byte-identical to main so it looks unrelated, but this PR is meant to export all the kimi-k3 kernel changes and is_flash_attention_v4_available() was missing. Restored.

It also prompted a file-by-file audit against kimi-k3, which found more lossy exports — dropped PDL in pack_topk_ids and concat_mla, the radix-select fast path (3.1-3.5x) and the num_warps heuristic in moe_fused_gate, v2p_ptr/PAGE_MULT in kv_indices (this one hard-crashes the K3 runtime), HEAD_DIM masking in verify_splitkv, the stride guard in kda_blackwell/kernel_h, do_not_specialize in gpu_tensor_hash, and the ROCm --use_fast_math strip in jit/utils/compile.py. All restored; the kernel tree is now byte-identical to kimi-k3 apart from the fixes in this review.

[BBuf]

Validation of 22cb347d47

Numbers in that commit message were taken mid-way through the restore, before the last batch of lossy exports was put back. Here is the full set, and it supersedes them.

Setup. 8× B300 (sm_103) TP8 single node, kimi-k3 runtime running against this PR's kernels, flashinfer_mxfp4 via the 0617 SiTU cubin pool, non-spec, no --enable-symm-mem, cookbook B300 low-latency recipe. ~/.cache/sglang and ~/.cache/tvm-ffi cleared before every run so nothing is served from a pre-PR artifact. bench_serving random, range-ratio 1.0, num_prompts = 5 × conc, --flush-cache per point. Baseline is the same box and recipe running kimi-k3's own kernels.

bs=1

	baseline	this PR
8192/1024 output	112.07 tok/s	112.16
8192/1024 TTFT	369.58 ms	369.20
8192/1024 TPOT	8.57 ms	8.56
1024/1024 output	114.61 tok/s	114.77
1024/1024 TPOT	8.49 ms	8.49

The 8k TTFT match matters on its own: that figure is the 8192-token prefill, and the review fixes in this PR touch the MLA KV write path (copy_bytes, elect.sync), so a prefill regression there would show up here. It does not.

Throughput point, 8192/1024 concurrency 24 (the SLO-optimal cell for this recipe). First reading was TTFT 5649 ms, which did not reproduce — it was the first load after a cold JIT cache. Three repeats:

	baseline	rep 1	rep 2	rep 3
output tok/s	843.14	858.38	858.03	857.76
total tok/s	7588.22	7725.38	7722.28	7719.85
mean TTFT	4606 ms	4748	4755	4759
median TTFT	—	4761	4776	4772
P99 TTFT	—	7933	7945	7956
mean TPOT	22.87 ms	23.34	23.34	23.35

Repeats agree within 0.1%. Steady state is +1.7% throughput, TTFT +3.3%, TPOT +2.1%. Worth noting the shape: theoretical output is roughly conc/TPOT, so baseline 24/22.87 = 1049 against 843 measured is 80.4% utilisation, while 24/23.35 = 1028 against 858 is 83.5%. Slightly slower per step, meaningfully less pipeline idle — the direction you would expect from PDL being back.

Accuracy. GSM8K 200q, three runs: 0.995 / 0.995 / 0.995, invalid 0.005 each. Baseline 0.990 — one question out of 200.

One process note. The checklist currently says "Speed Tests and Profiling: Not applicable to this extraction". The v2p_ptr removal argues against that for kernel-export PRs: it passes review by diff and passes a compile gate, because it is a Triton runtime kwargs check, and only surfaces when the ported kernels run one decode step under the real runtime. That single run is what turned up the other twelve regressions too.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[BBuf]

Pushed b78f2aab53 (consolidation, rebased onto 7009efd7ec). Validated end to end on 8×B300; details below, including a correction to my earlier reply about align_single_token — I said we would drop it, and that was wrong.

Correcting the align_single_token reply

I removed it, then put it back. The reasoning that justified removing it is worth writing down because it applies to every file in this PR.

The argument was: its only caller is fused_marlin_moe.py, that file is not in this PR, so the kernel is unreachable here. The premise is true and the conclusion does not follow. This PR exists so kimi-k3 can drop its own kernel copies and use main's, so the question is not "who calls it inside this diff" but "who calls it on the branch this export serves". There, the caller is live:

# kimi-k3: srt/layers/moe/fused_moe_triton/fused_marlin_moe.py:244
if M == 1 and topk <= 32 and expert_map is None:
    from sglang.kernels.ops.moe.moe_align_single_token import moe_align_single_token

Three properties make removal a hard failure rather than a slow path:

• the import is inside the branch, so a missing module raises ImportError at request time instead of falling through to the moe_align_block_size call in the else;
• the guard is M == 1, i.e. single-token decode — the case this kernel exists to speed up (one single-warp sort instead of the align + count_and_sort pair) and the case bs=1 numbers measure;
• Marlin is not exotic: with --moe-runner-backend at auto it is what gets picked whenever the SiTU cubin pool is absent, which is the default on a fresh box.

The file, its Python wrapper and its test case are all restored. That makes this commit three consolidations, not four.

While checking this I also swept the rest of the export. Of the files kimi-k3 changes under python/sglang/kernels/, exactly one is now not carried here: distributed/communicator.cuh. That one is deliberate — the PR supplies the same five system-scope PTX helpers from kimi_k3/comm/ptx_sys.cuh in the same device::distributed namespace, and the two helpers kimi-k3 additionally defines there (load_relaxed_sys, red_add_release_sys) are referenced by nothing in this PR. Its Semaphore refactor is cosmetic: same asm, immediate folded to a register operand.

What was measured, and on what

Running the PR branch by itself proves nothing — that tree is main plus kernels, and main has no K3 model support. So both arms are a composite: kimi-k3 @ 578edb240, with every kernel file it changes reverted to the merge-base, then this PR's kernels laid on top. Reverting matters as much as overlaying: leave kimi-k3's communicator.cuh in place and its helpers collide with ptx_sys.cuh's.

Scope is python/sglang/kernels/ only. The PR's two other source files (environ.py, custom_all_reduce_v2.py) are deliberately not overlaid — both were exported from kimi-k3, so it already holds them and additionally holds things this PR does not carry (SGLANG_ENABLE_CUSTOM_ALL_REDUCE_V2_MULTINODE, read by parallel_state.py:2001). Overlaying them replaces a superset with a subset.

Arms are one commit apart, so the delta isolates the consolidation:

	017c367079 (before)	f79e7b8327 (after)
bs=1 output tok/s	105.42	106.38
bs=1 TPOT	8.47 ms	8.47 ms
bs=1 TTFT p50	261.01 ms	239.31 ms
C=24 output tok/s	1000.94	1001.36
C=24 TPOT	19.66 ms	19.71 ms
GSM8K (200q)	0.995	0.995

8×B300 SXM6, TP8, --mem-fraction-static 0.85, --mamba-full-memory-ratio 7.21, MoE flashinfer_mxfp4 via the 0617 SiTU pool, attention auto-resolved to trtllm_mla. Both arms booted to the identical pool split (max_mamba_cache_size=476, max_total_num_tokens=136128), and random-range-ratio 1.0 fixes in=1024 / out=256 exactly.

TPOT is the indicative number here: it excludes TTFT, so it is the part shared activation / radix / mbarrier code could actually move, and it is flat in both regimes.

One honest note on the first pass: the after-arm initially measured 98.70 bs=1 tok/s, a 6.4% deficit. That was cold JIT — the consolidation changes source hashes, so its modules had to build while the before-arm hit an already-populated ~/.cache/tvm-ffi. The tell was that only mean/p99 TTFT moved (438/1679 ms vs 266/305) while median TTFT was already better and TPOT was unchanged; the warmup pass in the same run showed p99 5436 ms. Re-running the after-arm against a warm cache gives the 106.38 above. Reporting it because "re-ran it and it got better" deserves the mechanism attached.

Covering the path the main run cannot reach

flashinfer_mxfp4 never calls moe_align_single_token, so none of the above would have caught the removal. Separate arm, system interpreter (flashinfer 0.6.16rc5), no pool, --moe-runner-backend marlin:

• single-request generation correct (17 × 23 → 391);
• GSM8K 100q: accuracy 1.000, invalid 0.000;
• and the load is verified rather than inferred — moe_align_single_token…so is mapped in every scheduler process (/proc/<pid>/maps), which only happens via that if M == 1 import.

Verification of the consolidation itself

Each of the three is a move of existing text, not a retype, and two are checked bit-exactly:

• SiTU activation (kimi_k3/situ_activation.cuh) — bit-identical across three shapes covering both kHasLinearBeta branches.
• Radix primitives (moe/radix_select_common.cuh) — bit-identical across 16 tensors (ids and weights × four token counts × radix fast path and fused gate). The first attempt here was a retype and it silently changed every selected expert id while leaving the weights identical; that is why these are lifted verbatim now.
• mbarrier wrappers (sgl_kernel/mbarrier.cuh) — same global namespace ptx both consumers already open, so no call site changes.

All ten touched JIT modules cold-compile (artifacts confirmed newly built in ~/.cache/tvm-ffi, not cache hits).

One caveat on comparability

These runs are TP8 without DCP, so the absolute numbers are not comparable to the cookbook's B300 TP8+DCP8 cell. That is a branch-side version knot, not something this PR causes: cutedsl_mla_backend.py:329 passes enable_dcp=True, a kwarg flashinfer 0.6.15.post1 (the pyproject.toml pin) does not have, while 0.6.16rc5 cannot use the 0617 SiTU pool because the pool's overlay/include shadows the wheel's KernelRunner.h (TrtllmGenBatchedGemmRunnerOptions has no member perTokenSfDtype). Of the two reachable configurations I picked the one that exercises what this commit touches — the SiTU kernels are on the flashinfer_mxfp4 path, not the Marlin one. Both arms ran identically, so the A/B stands.