[Kernel] Avoid batch-size specialization in masked KV writes

[zzzzzzzxh]

Motivation

masked_set_kv_buffer_kernel receives N as a tl.constexpr, although N
only guards pid >= N and the launch grid is exactly (N,). The guard cannot
fire, while every distinct write batch size creates another Triton kernel
specialization.

This PR addresses the CUDA-side finding in #31568. The independent XPU
NT_BUCKET cleanup is intentionally left out so that it can be validated on
the appropriate hardware.

Modifications

• Remove N from the Triton kernel signature and its call site.
• Remove the unreachable pid >= N branch.
• Add a registered one-GPU regression test that launches grids of 17 and 33,
  verifies masked KV writes exactly, and asserts that the second launch reuses
  the first in-memory Triton specialization.
• Register the test in the existing
  base-b-test-1-gpu-large suite and provide its direct pytest
  entry point used by the CI runner.

Accuracy Tests

Reproduction on an NVIDIA RTX 3090 (SM86):

baseline: cache entries after N=17: 1
baseline: cache entries after N=33: 2
patched:  cache entries after N=17: 1
patched:  cache entries after N=33: 1

The committed test was also run through the normal SGLang package import path
in an environment aligned with current main:

Python: 3.11.15
PyTorch: 2.11.0+cu130
Triton: 3.6.0
GPU: NVIDIA GeForce RTX 3090

$ CUDA_VISIBLE_DEVICES=0 PYTHONPATH=python python \
    test/registered/kernels/test_masked_set_kv_buffer.py -f
collected 1 item
test/registered/kernels/test_masked_set_kv_buffer.py::test_batch_size_does_not_create_extra_specializations PASSED
1 passed

Full registry validation also succeeds and resolves the test to:

base-b-test-1-gpu-large (est_time=15s)

The test checks that selected key/value rows are copied exactly, masked-out rows
remain untouched, and changing only the grid size does not add a specialization.

Speed Tests and Profiling

This removes one redundant compilation/cache entry for each distinct N after
the first launch. No model-level throughput change is expected because the
generated write kernel and launch grid are otherwise unchanged.

Checklist

• Format your code according to the Format code with pre-commit. The GitHub lint workflow passes.
• Add unit tests according to the Run and add unit tests. The test is registered in base-b-test-1-gpu-large.
• Update documentation according to Write documentations. No documentation change is needed because this is an internal kernel cleanup with no API or user-facing behavior change.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed. Accuracy and specialization-cache results are included above; model-level speed benchmarking is not applicable because the generated kernel and launch grid are unchanged.
• 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. 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.
4. 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 #32210758233
Latest PR Test (Extra): ❌ Run #32210758198

[gemini-code-assist]

Code Review

This pull request removes the N parameter from set_kv_buffer and the Triton kernel masked_set_kv_buffer_kernel, eliminating the boundary check if pid >= N: return inside the kernel. This optimization prevents unnecessary kernel recompilations and specializations for different batch sizes. Additionally, a new unit test file test_masked_set_kv_buffer.py is added to verify the kernel's correctness and confirm that varying batch sizes do not create extra specializations. No review comments were provided for this pull request.

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.

[zzzzzzzxh]

/tag-and-rerun-ci

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 16f59727c3

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Register the kernel test in a dispatchable suite

This registration resolves to base-b-kernel-unit-test-1-gpu-small, but the CUDA suite allowlist in test/run_suite.py only includes the kernel-unit suites for 1-gpu-large, 4-gpu-b200, and 8-gpu-h200; validate_all_suites() rejects any registered test whose effective suite is not in that list before running a shard. In CI contexts that execute test/run_suite.py, adding this file makes collection fail with an invalid-suite error instead of running the tests, so this should use an existing dispatchable suite such as runner_config="1-gpu-large" or add the corresponding workflow/run_suite support.

Useful? React with 👍 / 👎.

[zzzzzzzxh]

@xiezhq-hermann @ispobock This addresses the CUDA-side specialization issue in #31568. The committed regression test now passes through the normal SGLang import path on an RTX 3090 with the current main kernel stack (torch 2.11.0+cu130, Triton 3.6.0): 1 passed in 7.36s.

Could you please review it and add the run-ci label when appropriate? I tried /tag-and-rerun-ci; as the PR author I can rerun failed jobs, but I do not have permission to add the initial label.

[zzzzzzzxh]

CI update after fixing the test registration:

• The new regression test was executed in base-b-test-1-gpu-large (3) and passed (1 passed, about 12 seconds): https://github.com/sgl-project/sglang/actions/runs/29889309928/job/88829696463
• All eight base-b-test-1-gpu-large shards passed.
• The remaining standard CI failures appear unrelated to this PR: Base C failed in test_deepep_small.py during DeepGEMM CUDA graph capture and then fast-failed dependent shards; the two AMD failures hit Hugging Face API rate limits; and the XPU failure was a DeepSeek-OCR server startup timeout.
• The Extra workflow failures are only the expected gate for the optional run-ci-extra label, which I am not requesting for this narrow CUDA change.

Could a maintainer please trigger a failed-only rerun for the standard CI workflows? @xiezhq-hermann @ispobock Thank you.