Motivation

Follow #18073

The old kernel handled both q and k inside one CTA, which kept too much
state live at the same time:

• q
• k
• q_weight
• k_weight
• separate output vectors
• dual reduction buffers in shared memory

The new kernel still performs the work in a single launch, but splits the work
with grid.y = 2:

• blockIdx.y == 0: normalize q
• blockIdx.y == 1: normalize k

This reduces per-thread live state and shrinks the shared reduction buffer from
two lanes to one lane.

On H200 with shape (batch_size=2048, hidden_dim=8192):

• registers/thread: 48 -> 26
• static shared memory/block: 256 B -> 128 B
• theoretical occupancy: 50% -> 100%
• achieved occupancy: 45.25% -> 88.17%
• achieved active warps/SM: 28.96 -> 56.43

Microbench results

H200, bf16:

Modifications

Accuracy Tests

Benchmarking and Profiling

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 Process

1. Ping Merge Oncalls to start the PR flow. 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.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.