[voltjia]

Summary

• Migrate the Qwen3 fast path from the internal MergedQkRmsNorm experiment to the vLLM-aligned FusedQkNorm operator from InfiniOps #1023.
• Pass the packed rank-2 QKV projection directly to FusedQkNorm; Q and K are normalized in place, while V remains unchanged.
• Keep the checkpoint format unchanged and use the original q_norm.weight and k_norm.weight tensors directly, removing the post-load merged-weight allocation and copies.
• Enable the fused path only for NVIDIA with BF16, head_dim=128, and an active FusedQkNorm implementation. All other configurations retain the existing separate Q/K RMSNorm path.
• Pin InfiniCore #1580 final main merge commit 82095baef0369797132898a8d132ca13ae60807b, which points to the final InfiniOps #1023 master merge commit 1a68b261d19f0ed6f76bbabeab7a25f479f05e07.

Motivation

Qwen3 applies separate RMSNorm to query and key head groups after the packed QKV projection. The final public operator shape follows the vLLM-style packed-QKV FusedQkNorm API, while retaining norm-only behavior so RoPE and the surrounding model path remain unchanged.

Closes N/A - no separate issue exists.

Type of Change

• feat - new feature / new model
• fix - bug fix
• perf - performance improvement (no behavioral change)
• refactor - code restructuring without behavior change
• test - adding or fixing tests only
• docs - documentation only
• build / ci - build system or CI configuration
• chore - tooling, formatting, or other non-code changes
• Breaking change

Test Results of Involved Models on Supported Platforms (Please attach screenshots)

Platform	Result
NVIDIA A100-SXM4-80GB / CUDA 12.6.77 / driver 580.105.08	Passed

Validation performed on the A100 container using the full InfiniOps prefix with CPU, NVIDIA, Torch, NineToothed, Flash-Inference, and W4 Marlin enabled:

InfiniOps #1023 master merge 1a68b261d19f0ed6f76bbabeab7a25f479f05e07
InfiniCore #1580 final main merge 82095baef0369797132898a8d132ca13ae60807b

xmake configure: release, C++11 ABI=1
xmake build infinilm_runtime _infinilm: passed
xmake install infinilm_runtime _infinilm: passed

The basic instruction smoke used a one-layer Qwen3 model with BF16, four query heads, two KV heads, and head_dim=128:

examples/test_infer.py
device=nvidia, dtype=bfloat16, max_new_tokens=8
prefill + decode completed and generated output

Call-path verification used INFINI_OPS_TRACE_CALLS=1:

{"operator_name": "FusedQkNorm", "device_type": "nvidia", "implementation": 9}

The trace contained one FusedQkNorm call for the one-layer prefill and did not fall back to separate Q/K RmsNorm calls.

The A100 smoke and benchmark measurements used the reviewed #1023 operator patch. The operator patch remained unchanged through rebase, and InfiniOps merged #1023 as master commit 1a68b261d19f0ed6f76bbabeab7a25f479f05e07. InfiniLM #609 has also merged into main as commit 12a4492f70dd33235e3d142e333ad4a9629d3b7b. The current #607 head c0e6359 is a clean single commit directly based on InfiniLM main commit 12a4492 and pins the #1023 master commit through InfiniCore #1580 final main merge commit 82095bae.

Benchmark / Performance Impact

Operator-level A100 medians on InfiniOps #1023 show that the aligned packed-QKV kernel itself remains close to the prior internal merged-QK kernel:

Sequence length	Prior MergedQkRmsNorm median	FusedQkNorm median
1024	0.11013 ms	0.11562 ms
8192	0.66273 ms	0.66826 ms

The small kernel cost is offset in InfiniLM by removing the post-load Q/K weight packing allocation and copies and by normalizing Q/K without allocating a separate normalized QK tensor.

End-to-end prefill was measured with examples/bench.py on the same A100, same random one-layer Qwen3 model, BF16, batch size 1, output length 1, prefix caching disabled, synchronous weight loading, and one warmup. Median prefill TTFT is reported. The 1024/4096 rows use three runs; 8192 uses ten runs because that length had materially higher run-to-run variance.

Input length	Prior merged-QK stack	FusedQkNorm stack	Median change
1024	4.16 ms	3.61 ms	-13.2%
4096	13.24 ms	12.99 ms	-1.9%
8192	47.22 ms	46.47 ms	-1.6%

The 8192 samples ranged from roughly 38.6 to 66.9 ms on both stacks, so the long-length result should be read as parity within noise rather than a strong standalone gain.

Notes for Reviewers

• This is the norm-only path selected after measuring the fused norm+RoPE experiment in InfiniOps #1019; RoPE remains intentionally outside this operator.
• The call site deliberately uses the public infini::ops::FusedQkNorm factory rather than probing the private NineToothed provider header. The installed public infini/ops.h exposes the operator and active-implementation query.
• This PR is based directly on current main after build: pin NineToothed fused add RMSNorm #609. build: pin NineToothed fused add RMSNorm #609 carried the reviewed build: pin NineToothed fused_add_rms_norm #606 dependency pin into main after build: pin NineToothed fused_add_rms_norm #606 was accidentally merged into the old stack branch. InfiniOps #1023 has merged; InfiniCore #1580 has merged; its final main merge commit pins the #1023 master commit.

CI / ChatOps

Host behavior, Ruff, and format checks passed on clean head c0e6359 using the historical PR head branch perf/ninetoothed-merged-qk-rms-norm. Native CI run 41998576713 was dispatched on this head and remains queued while the configured InfiniLM runner pool is offline to GitHub Actions.

---

Checklist

Title, Branch, and Commits

• PR title follows Conventional Commits.
• Branch name follows <type>/xxx-yyyy-zzzz.
• Each commit message follows Conventional Commits.
• Every stack commit is meaningful and independently reviewable.
• No stray merge commits from main.
• No fixup! / squash! / wip commits remain.
• N/A - Existing PR/branch/commit does not use the legacy issue format.

Scope and Design

• Changes are minimal and limited to the Qwen3 attention fast path and dependency pin.
• No dead code, debug output, or TODOs remain in the diff.
• No unrelated formatting churn.
• N/A - No public InfiniLM API changes; the changed call targets the pinned InfiniOps API.

General Code Hygiene

• The code is self-explanatory; the fallback condition is expressed directly at the call site.
• Every modified or added file ends with a single trailing newline.
• No trailing whitespace, tab/space mixing, or stray BOMs.
• Identifiers in comments and messages use Markdown backticks.
• All comments and error messages are in English.
• Comments and messages follow language conventions.

C++ Specific

• Code follows the Google C++ Style Guide.
• Error and warning wording follows LLVM conventions.
• Constructor initializer list order matches member declaration order.
• No raw new/delete were added.
• Changed files were formatted with clang-format; the repository formatter uses the same configuration.
• No changes or references were added under csrc/models/llama_legacy/.

Python Specific

• N/A - No Python source files are changed by this PR.

Testing

• N/A - Non-NVIDIA platforms were not run because the new path is NVIDIA-only and other configurations retain the existing implementation.
• Passed a focused one-layer Qwen3 single-request prefill/decode smoke on A100.
• Passed offline prefill benchmarks with examples/bench.py on A100.
• N/A - test/bench/test_benchmark.py was not run because the validation scope was the Qwen3 attention fast path and this repository has no committed Qwen3 random-weight fixture required for that harness.
• N/A - Service behavior is unchanged and the change is below the service layer.

Build, CI, and Tooling

• The project builds cleanly from the dependent full InfiniOps prefix on NVIDIA.
• CI was triggered manually on the updated branch.

Documentation

• N/A - No user-facing configuration, build flag, or developer workflow is introduced by this PR.
• N/A - No breaking change.

Security and Safety

• No secrets, tokens, internal URLs, customer data, or personal hardware identifiers were committed.
• No third-party code was added.
• No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.