[voltjia]

Summary - Add the NineToothed NVIDIA implementation of FusedQkNormRope. - Align the public API and parameter order with vLLM's fused_qk_norm_rope, including forced_token_heads_per_warp = -1 for automatic scheduling. - Apply per-head Q/K RMSNorm and NeoX RoPE in place to packed qkv tensors of shape [tokens, (Hq + Hk + Hv) * D], while preserving the V section unchanged. - Use separate contiguous one-dimensional q_weight and k_weight tensors, preserving the checkpoint format used by model code. - Apply the paired dimension's RMSNorm weight when reading the NeoX paired input, instead of reusing the current dimension's weight. ## Motivation The previous InfiniLM experiment first merged Q/K RMSNorm but still performed RoPE separately. Packing the QKV projection and applying Q/K normalization and RoPE together removes intermediate tensor traffic and keeps the public operator contract aligned with vLLM's fused kernel. This also supersedes #1016. That PR was closed automatically by GitHub when its head branch was renamed; #1019 contains the clean, rebased implementation. Closes N/A - no separate issue exists; supersedes #1016. ### Open-Source Alignment The API and packed-QKV semantics are compared against vLLM commit 16d667348d91ae010cb3cc9d4d6a5b7908a717a9: - Python wrapper: vllm/_custom_ops.py#L353 - C++ declaration: csrc/ops.h#L251 - CUDA implementation: csrc/fused_qknorm_rope_kernel.cu#L773 Aligned with vLLM: - Parameter order: qkv, num_heads_q, num_heads_k, num_heads_v, head_dim, eps, q_weight, k_weight, cos_sin_cache, is_neox, position_ids, forced_token_heads_per_warp. - The Python binding uses forced_token_heads_per_warp = -1 as vLLM does for automatic scheduling. - qkv is a contiguous rank-2 packed projection and is updated in place. - q_weight and k_weight remain independent [head_dim] tensors; no pre-packed two-dimensional norm weight is introduced. - Q and K receive per-head RMSNorm followed by RoPE; V is untouched. Intentional current differences: - This first InfiniOps fast path is guarded to NVIDIA, BF16, head_dim=128, full-dimension RoPE, NeoX layout, contiguous packed QKV, and automatic forced_token_heads_per_warp==-1. - Explicit positive forced-warp values are not accepted; scheduling remains owned by the InfiniOps provider rather than mirroring vLLM's tuning knob. - The kernel is independently implemented with NineToothed; no vLLM source is copied. ## Type of Change - [x] feat — new feature / new operator / new platform - [ ] 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 - [ ] Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer) ## Platforms Affected - [ ] CPU (WITH_CPU) - [x] NVIDIA (WITH_NVIDIA) - [ ] Iluvatar (WITH_ILUVATAR) - [ ] MetaX (WITH_METAX) - [ ] Cambricon (WITH_CAMBRICON) - [ ] Moore (WITH_MOORE) - [ ] Ascend (WITH_ASCEND) - [ ] PyTorch C++ bindings (WITH_TORCH) - [x] Build system / CMake / CI - [x] Python bindings / user-facing API ## Smoke Test Result

N/A - FusedQkNormRope is intentionally excluded from the minimal smoke set. The focused validation below is not a smoke-test result. Run on A100-SXM4-80GB with CUDA 12.6.77 and driver 580.105.08: text pytest -q tests/test_fused_qk_norm_rope.py --devices cuda Running 4 items in this shard .... 4 passed in 0.53s The four cases cover seq_len=1,17, 16 query heads, 8 key heads, 8 value heads, BF16, D128, full NeoX RoPE, and both omission and explicit passing of forced_token_heads_per_warp=-1. ## Test Results on Supported Platforms | Platform | Affected | Build / Smoke Result | Full Result / Notes | | ---------- | :------: | ------------------- | ------------------- | | NVIDIA | Yes | Excluded from smoke by design | Focused suite passed. The operator was included and passed CI at implementation-equivalent e7dc667; the later diff only removes it from smoke selection. | | Iluvatar | No | N/A - not affected | CI unit and shadow unit passed. | | MetaX | No | N/A - not affected | CI unit and shadow unit passed. | | Cambricon | No | N/A - not affected | CI unit and shadow unit passed. | | Moore | No | N/A - not affected | CI unit and shadow unit passed. | | Ascend | No | N/A - not affected | CI unit and shadow unit passed. | Additional validation: - Clean Release CMake build and install succeeded on A100 with CPU, NVIDIA, Torch, and NineToothed enabled at implementation-equivalent e7dc667. - ruff format --check, ruff check, clang-format --dry-run --Werror, and git diff --check passed. - NineToothed code generation was verified with the already-merged CI dependency revision bf8d7c42.

Full pytest output

text pytest -q tests/test_fused_qk_norm_rope.py --devices cuda Running 4 items in this shard .... 4 passed in 0.53s

## Benchmark / Performance Impact

Operator-level CUDA-event benchmark on A100-SXM4-80GB, CUDA 12.6, BF16, 16 Q heads / 8 K heads / 8 V heads, D128, full NeoX RoPE. Each reported value is the mean milliseconds per complete Q/K normalization-plus-RoPE transformation over timed rounds of five calls after warmup. The 1024-token case used 50 rounds; 4096 and 8192 used 30 rounds. Q/K outputs were preallocated for separate baselines, and the fused packed-QKV input was reset outside timed regions.

Baselines:

• Native separate: two RmsNorm calls at implementation index 0 plus two RotaryEmbedding calls at index 0.
• Strong separate: two RmsNorm calls at NineToothed implementation index 9 plus two native RotaryEmbedding calls.
• Fused: one NineToothed FusedQkNormRope call.

Tokens	Native separate	NineToothed-RMS separate	Fused QK norm + RoPE	Fused vs native	Fused vs strongest separate
1024	0.29819 ms	0.26030 ms	0.36692 ms	23.1% slower	40.9% slower
4096	1.05930 ms	0.81250 ms	0.84431 ms	20.3% faster	3.9% slower
8192	2.31373 ms	1.60811 ms	1.67105 ms	27.8% faster	3.9% slower

A post-run correctness check compared fused output with the separate path; V was bit-identical. BF16 differences were confined to a very small number of elements and remained within the benchmark sanity tolerance after accumulated separate-path rounding.

Interpretation: this revision is faster than the non-tuned native baseline at long token counts, but it is not faster than an already NineToothed-RMSNorm separate path. The current in-place packed-QKV implementation also copies the full packed QKV into and out of a workspace to preserve V and avoid source/output hazards. There is therefore no end-to-end performance claim in this PR, and model integration should not assume a win without further kernel optimization.

Notes for Reviewers The current provider is intentionally a guarded fast path and is intentionally absent from the minimal smoke set: NVIDIA, BF16, head_dim=128, full-dimension RoPE, NeoX layout, contiguous packed QKV, contiguous int64 positions, and automatic forced_token_heads_per_warp==-1. Explicit positive forced-warp values remain unsupported because scheduling is delegated to this provider.