[Kimi K3][Kernel] Fuse MoonViT Q/K complex RoPE

[gcanlin]

Purpose

MoonViT currently materializes the real and imaginary RoPE components in every
encoder layer and invokes ApplyRotaryEmb separately for the strided Q and K
views of its packed QKV projection. This PR adds an SM90+ Triton path that:

• reads the packed-QKV Q/K views directly;
• rotates Q and K in one kernel and produces the same contiguous output layout;
• keeps the token count as a runtime argument to avoid recompilation for random
  image sizes; and
• falls back to the existing implementation for unsupported devices, dtypes,
  shapes, and empty inputs.

The kernel is adapted from SGLang's fused vision RoPE implementation, while its
arithmetic order follows vLLM's current interleaved RoPE path.

This is not a duplicate of an existing open PR. Searches for MoonViT RoPE,
Kimi vision RoPE, fused QK complex RoPE, and kimi_k25_vit found no PR
implementing this kernel. #53011 addresses eager torch.compile cache behavior,
and #40600 targets full ViT CUDA graphs; neither fuses Q/K complex RoPE.

AI assistance (OpenAI Codex) was used for implementation and testing. The human
submitter must review every changed line and be able to defend the change before
this draft is marked ready for review.

Test Plan

1. Run pre-commit hooks on all changed files.
2. Compare BF16/FP16 fused results against FP32 complex-multiply references for
   Kimi K3's packed-QKV layout and dynamic token counts.
3. Capture and replay the fused path with a CUDA graph after changing inputs.
4. Benchmark the exact current MoonViT boundary (ApplyRotaryEmb twice,
   including real/imag materialization) against the fused path on B200.
5. Before marking ready, run a trained Kimi K3 multimodal end-to-end comparison.

Test Result

Local hooks:

.venv/bin/pre-commit run --files \
  vllm/model_executor/layers/rotary_embedding/vision.py \
  vllm/model_executor/models/kimi_k25_vit.py \
  tests/kernels/core/test_vision_rope.py \
  benchmarks/kernels/benchmark_vision_rope.py

All hooks passed, including ruff, formatting, mypy, SPDX, and forbidden-import
checks.

B200 kernel tests in the vLLM nightly container:

/workspace/.venv/bin/python -m pytest /workspace/test_vision_rope.py -q
8 passed, 14 warnings in 17.87s

The tests cover BF16 and FP16, 1/257/4096 tokens, packed QKV strides,
contiguous outputs, unsupported-input rejection, and CUDA graph replay. In a
direct BF16 comparison with the existing vLLM path, the maximum difference was
one BF16 ULP (0.0078125), affecting 0.00152% of elements.

B200 BF16 microbenchmark, 12 heads and head dimension 128:

Tokens	Existing Q/K RoPE	Fused Q/K RoPE	Speedup
256	150.02 us	11.81 us	12.7x
1,024	150.75 us	12.06 us	12.5x
4,096	151.15 us	20.48 us	7.4x
16,384	302.11 us	57.31 us	5.3x

A 27-layer random-weight MoonViT run using Kimi K3's vision geometry produced
finite outputs with cosine similarity 0.999912 between the existing and fused
paths. This is a numerical stress test, not a trained-model evaluation.

Trained Kimi K3 multimodal end-to-end serving/evaluation has not been run
yet. This PR remains a draft until that result and human line-by-line review are
complete.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR.
• The test plan.
• The available test and benchmark results.
• No documentation update is required for this internal kernel optimization.

[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.

[gau-nernst]

q: this doesn't look very different from RoPE used in LLMs. Can we reuse existing RoPE kernels in vLLM? also, since you are doing this, can we fuse QK-Norm+RoPE as well (I actually don't know if MoonViT has QK-Norm, just wondering)