[voltjia]

Supersedes #1016, which GitHub closed automatically when its head branch was renamed.

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.

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 the content-identical implementation commit 16b6639; final commit e7dc667 differs only by restoring the existing NineToothed CI pin.
• Verified code generation with the already-merged NineToothed dependency pinned by CI at bf8d7c42.
• 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.