[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], with separate one-dimensional q_weight and k_weight tensors. The V section is preserved unchanged.
• Use the paired RMSNorm weight for the paired NeoX dimension, not the current dimension's weight.
• Bump the CI NineToothed dependency to dfce2a5f0ec2a2a9428cb4cc897a344eb7a8f725, which includes the dynamic source-index lowering fix required by the paired-weight expression.

Current Fast Path

The NineToothed provider currently supports 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 are intentionally reported unsupported because scheduling remains delegated to this provider.

Validation

• Clean Release build and install from final commit 16b6639 on latest master 3fac025.
• A100-SXM4-80GB, driver 580.105.08, CUDA 12.6.77.
• pytest -q tests/test_fused_qk_norm_rope.py --devices cuda: 4 passed.
• Covered seq_len=1,17, 16 Q heads, 8 K heads, 8 V heads, BF16, D128, full NeoX RoPE, and both omitted and explicit forced_token_heads_per_warp=-1.
• ruff format --check, ruff check, clang-format --dry-run --Werror, and git diff --check passed.

[voltjia]

Superseded by #1019 after the head branch was renamed to eat/ninetoothed-fused-qk-norm-rope and GitHub closed this PR automatically. #1019 contains the same vLLM-aligned rework at 16b6639.