[mani1106]

Motivation

Dense Qwen3 prefill launches QK normalization and cached RoPE separately, then FA3 or FlashInfer materializes contiguous Q. This change combines normalization, the same cached RoPE, and contiguous Q output in one CUDA launch. K is normalized and rotated in its original packed QKV storage; Q input and V remain unchanged.

It uses the existing --enable-fused-qk-norm-rope opt-in and remains disabled by default on CUDA. Attention kernels and KV-cache handling are unchanged.

Modifications

• Kernel: Add lazily registered attention.qk_norm_rope with FP16/BF16 support, tensor/device/dtype/stride/alignment checks, and explicit K/output mutation metadata. D128 processes two rows per warp; D64/D256 use one. Share the existing standalone half-warp normalization arithmetic, round normalized values before FP32-table rotation, and preserve shuffle participation for an unmatched final row.
• Dense dispatch: Fuse only ordinary dense-Qwen3 EXTEND on Hopper, head dimension 128, TP1, unquantized weights, FA3/FlashInfer, full-head unscaled NeoX RoPE, and matching activation/weight dtypes. Keep speculative, mixed, verification, inherited/non-Qwen3, deterministic, compiled, and other unvalidated configurations on the separate path. Preserve existing BF16 table decode guards.
• Coverage/documentation: Extend existing kernel and model fixtures for native byte parity, Q/K/V ownership, padding, empty input, fallback dispatch, and changing-input CUDA graphs. Add a cold-L2 preparation benchmark including the baseline consumer's Q conversion, and document the opt-in in the Qwen3 cookbook.

Disabling the existing flag or reverting this contribution rolls the change back; there is no checkpoint, schema, or persisted-data migration. On upstream base-versus-PR comparisons, keep the existing flag enabled in both trees so BF16 decode treatment is identical.

Accuracy Tests

Current contribution-source checks, run on H200 and the checkout's Python path:

Command	Result
python -m pytest test/registered/kernels/ops/attention/test_table_qk_norm_rope.py -q	35 passed; includes FP16/BF16 native-byte parity, ownership, padding and empty input
python -m pytest test/registered/unit/models/test_qwen3_qk_norm_rope.py -q	21 passed; decode preservation, exclusions, BF16/FP16 changing-input graphs
python -m pytest test/registered/kernels/ops/layernorm/test_qknorm.py -k two_rows -xq	3 existing standalone half-warp equivalence regressions passed
compute-sanitizer --tool memcheck --error-exitcode 99 python -m pytest test/registered/kernels/ops/attention/test_table_qk_norm_rope.py -k hybrid -xq	7 passed; 0 errors
Same selected cases under compute-sanitizer --tool synccheck --error-exitcode 99	7 passed; 0 errors
python -m pytest test/registered/unit/kernels/test_kernels_namespace.py -q	104 tests and 2 subtests passed

Fresh compatible-serving validation: one H200, TP1, SGLang v0.5.20 serving paths, Torch 2.13.0.2+cu130, Triton 3.8.0, FlashInfer 0.6.18, and kernel 0.4.7.0+cu130. The QKV producer, cached NeoX RoPE and attention consumer paths used by these cases were checked; the exact PR CUDA/shared-normalization files from 01b927b were loaded through a compatibility adapter. The CUDA header SHA256 is eb5c6486fdfd265bf8aa9a54ea52be72ce126516f50802b747ea951e4bf430e0. The comparison uses optimized #45214 standalone normalization in the baseline and preserves identical original decode treatment in both arms.

Paired real-model smoke outputs are exact at 128, 4096 and 8192 input tokens for Qwen3-4B BF16/FA3 and Qwen3-Embedding-0.6B FP16/FlashInfer. A direct compatible-runtime FP16/BF16 contract check also matches Q/K exactly, with Q input/V unchanged, K updated and contiguous Q.

Paired Qwen3-4B GSM8K-200 with sgl-eval 0.1.3:

Arm	Correct	Errors	Truncations
Optimized split baseline	188/200 (94%)	0	0
PR CUDA hybrid	188/200 (94%)	0	0

The same 200 questions, expected answers, generation texts and per-question grades match exactly. The existing verified canonical corpus was reused; no prompts, grading or tolerances were changed.

Checkpoints:

• Qwen/Qwen3-4B revision 1cfa9a7208912126459214e8b04321603b3df60c
• Qwen/Qwen3-Embedding-0.6B revision 97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3

Accuracy command/settings (MODEL is the verified local snapshot of the pinned checkpoint):

sgl-eval run gsm8k --base-url http://127.0.0.1:31122/v1 \
  --model "$MODEL" --num-examples 200 --num-threads 1 \
  --max-tokens 2048 --temperature 0 --top-p 1 --seed 42 \
  --no-thinking --chat-template-kwarg enable_thinking=false \
  --out-dir "$OUTPUT_DIR"

Speed Tests and Profiling

Fresh compatible-serving comparison: baseline A1 → candidate B → baseline A2, with fresh servers, identical checkpoint/runtime/seed/cache/chunking/graph settings, and 200 measured requests per window at concurrency 4 after warmup. This is six server phases and nine serving windows, not a repeat of the broader candidate matrix. All 1800 measured requests completed without errors.

For each metric, the baseline below is the arithmetic mean of A1 and A2's window metric. Change is 100 * (B / mean(A1, A2) - 1); these are descriptive bracketed comparisons, not confidence intervals or pooled percentiles.

Workload	p50 change	p95 change	p99 change	Completed requests/s change
Qwen3-4B BF16/FA3, 4096 input / 1 output	-2.39%	-2.75%	-2.93%	+2.27%
Qwen3-Embedding-0.6B FP16/FlashInfer, 4096 input	-3.70%	-1.47%	-2.89%	+3.65%
Qwen3-Embedding-0.6B FP16/FlashInfer, 8192 input control	-1.69%	-1.59%	-1.76%	+1.86%

Absolute per-window latency:

Workload	A1 p50 ms	B p50 ms	A2 p50 ms	A1 p95 ms	B p95 ms	A2 p95 ms
Qwen3-4B 4096→1	251.470	245.000	250.533	256.680	249.602	256.639
Embedding 4096	55.062	52.898	54.794	56.167	55.252	55.986
Embedding 8192	160.953	157.954	160.397	163.223	160.947	163.859

Generation latency is the complete non-streaming HTTP request through response-body read/JSON parse with one generated token: a one-token request-latency/TTFT proxy, not first-streamed-token TTFT. Embedding latency is complete embedding response latency. Throughput is completed requests per second over each measured window.

Serving reproduction protocol

For an independent source-pair reproduction, compare base b50d4394945d3dcacea05ee4afd334a9b225cbdc with PR 01b927b0ba6992c18b6dd185f0490ff528eeb0fe on one Hopper GPU. Enable the existing flag in both source trees; use the pinned model revisions above. The recorded compatible-v0.5.20 measurements use the adapter described above, so these source-pair instructions specify the workload rather than package that adapter.

Common launch settings in every arm:

OMP_NUM_THREADS=8 python -m sglang.launch_server \
  --model-path "$MODEL" --dtype "$DTYPE" --attention-backend "$BACKEND" \
  --enable-fused-qk-norm-rope \
  --host 127.0.0.1 --port 31122 --tp-size 1 --random-seed 42 \
  --chunked-prefill-size 4096 --max-prefill-tokens 32768 \
  --max-total-tokens 196608 --max-running-requests 64 \
  --mem-fraction-static 0.7 --log-level warning \
  --cuda-graph-config '{"prefill":{"backend":"breakable","bs":[1,2,4,8,16,32,64,128,256,512,1024,2048,4096,8192,16384,32768],"max_bs":32768},"decode":{"backend":"full","bs":[1,2,4,8,16,32],"max_bs":32}}'

For Qwen3-4B use DTYPE=bfloat16, BACKEND=fa3, and append --return-hidden-states-mode last. For Qwen3-Embedding-0.6B use DTYPE=float16, BACKEND=flashinfer, and append --is-embedding. Use local model snapshots with offline Hugging Face/Transformers access and only one visible GPU.

The HTTP client uses explicit token IDs, not tokenizer-generated random text:

def ids(length, seed):
    return [1000 + (i * 13 + seed * 79) % 3000 for i in range(length)]

• Generation: POST /generate with input_ids=ids(4096, seed), sampling_params={"temperature": 0, "max_new_tokens": 1, "ignore_eos": True}, return_logprob=True, and return_hidden_states="last". Hidden-state/logprob response serialization is included in the measured latency.
• Embedding: POST /v1/embeddings with one token-ID array in input=ids(length, seed) and encoding_format="float"; measure lengths 4096 and 8192 separately.
• Order/cache: fresh servers in A1 -> B -> A2 order for each model/configuration. Before each window issue eight concurrent warmups with seeds 700 + i (i=0..7), wait for all eight, then POST /flush_cache?timeout=60 once. Read that response as text, not JSON. Submit 200 measured requests with seeds 1000 + i (i=0..199) through a semaphore of four, with no flush between requests.
• Timing: after acquiring the semaphore, start the request timer immediately before the HTTP POST; validate status, read the complete response body as text and JSON-decode it before stopping. Client semaphore queue time is excluded. Compute percentiles over the 200 successful request times. A separate wall timer surrounds submission/completion of all 200 requests; throughput is 200 / wall_seconds.

bench_serving --disable-stream reproduces the non-streaming timing class, but ordinary random-ids mode does not guarantee the same IDs, diagnostic fields or JSON-read boundary. Use the explicit HTTP protocol above when matching this workload.

Untimed actual-server profiles confirm qk_norm_rope_kernel engagement in the candidate:

Profiled workload	Baseline GPU kernels	Candidate GPU kernels	Reduction
Qwen3-4B	2081	2009	72
Qwen3-Embedding-0.6B	364	308	56

The activation-dtype Q-copy kernel family disappears; other copy kernels remain. Native decode normalization remains in the generation candidate trace. Separate current-source preparation profiles show three preparation kernels → one, with no aten::contiguous in the hybrid path. D128 BF16/FP16 int32/int64 specializations compile with 56 registers and zero spills on SM90a. Profiles are untimed engagement/copy-removal evidence and are not mixed into serving measurements.

Current-source cold-L2 preparation microbenchmark: python test/registered/kernels/benchmark/attention/bench_qk_norm_rope.py passed all six token/dtype shapes. At 4096 tokens, 32 Q heads / 8 KV heads / D128, split norm + cached RoPE + Q conversion measured 79.456 → 33.455 us (BF16) and 79.331 → 31.787 us (FP16). The marker framework rotates inputs to defeat L2 reuse.

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.
• Local code review completed.

Ruff 0.15.1, isort 7.0.0, clang-format 20.1.7, registered-test admission, pytest entrypoint lint and git diff --check passed. Repository pre-commit checks passed for all 10 contribution files with pre-commit run --show-diff-on-failure --files $(git diff --name-only b50d439..HEAD); no files were changed. Both local Claude and Copilot GPT complete-diff reviewers covered all 10 files and returned no actionable findings; contribution code remains at the reviewed 01b927b.

Final audit reran the 7 hybrid cases, 21 model dispatch/graph cases and 3 shared half-warp regressions: all 31 passed without code changes. Renewed complete-diff Standards and Spec/functionality reviews returned no actionable findings. Base CI currently stops at the draft gate ("PR is draft. Blocking CI."); this is not a GPU-test failure or a CI pass.

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.

Implementation and local validation were assisted by GitHub Copilot; automated authorship is disclosed in the commit trailer.

---

CI States

Latest PR Test (Base): ✅ Run #41798402706
Latest PR Test (Extra): ❌ Run #41798402657
Latest PR Test (AMD ROCm 10): ⏳ Run #41798402707

Tip

Base CI only covers the core cases, and its workload may not exercise your change. Use /rerun-test to run the tests your change touches, or copy this prompt for an agent:

Find the registered tests this PR affects and run them with /rerun-test, following https://github.com/sgl-project/sglang/blob/main/.agents/skills/rerun-test/SKILL.md

[mintlify]

Preview approval required

This pull request comes from a fork, so Mintlify won't build a preview until someone with write access approves commit b1a2b65.

Approve preview

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[mani1106]

Fixed the unrelated CI failure - Inkling checkpoint replay fixture in d6f3d18. After #47173, the draft graph runner reads backend.cache_loc_dtype instead of returning a fixed torch.int64. The fixture still returned a bare Mock(), so real buffer allocation failed with TypeError: ones(): argument 'dtype' must be torch.dtype, not Mock. This is an upstream test-fixture mismatch, unrelated to the Qwen3 fusion changes.

The open fix #47858 addresses this by setting cache_loc_dtype=torch.int64 on the mock backend. I reused that exact one-line test correction here so this PR is not blocked by the same fixture error. Production code, capture/replay assertions, and parameter combinations are unchanged. Since the fixture edit is identical, it is expected to merge cleanly when the upstream fix lands; any intervening changes still need to be checked at merge time.

Python syntax, the isolated before/after mock dtype contract check, and git diff --check passed locally. Full GPU replay validation remains for CI: this local environment has no PyTorch/CUDA, and pre-commit is unavailable.

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/tag-and-rerun-ci

[BBuf]

/tag-and-rerun-ci

[mani1106]

/rerun-failed-ci

[mani1106]

@BBuf Base CI passed. Can you help take a look?? Thank you!

[BBuf]

LGTM.

[mani1106]

/rerun-failed-ci try again

[mani1106]

Applied the exact one-line Arm64 test fix proposed in #48223 in commit 093cf61 to address the unrelated CI regression tracked in #48221.

test/registered/cpu/arm64/test_moe.py now explicitly imports sgl_kernel so its torch.ops.sgl_kernel operators are registered before use. The test no longer depends on the transitive import side effect removed by #47731. This is the only change in this commit; the Qwen3 fusion, eligibility, kernel implementation and numerical behavior are unchanged.

[BBuf]

/tag-and-rerun-ci

[mani1106]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci

[BBuf]

/rerun-failed-ci