rope: CUDA-graph-safe Triton ViT RoPE, unblocking Qwen3.5 in encode_server

[ZhengWG]

Motivation

The disaggregated encoder process (python/sglang/srt/disaggregation/encode_server.py) runs the ViT path of multimodal models such as Qwen3.5 / Qwen3.5-MoE / Qwen3-VL via model.get_image_feature / get_video_feature. Those route into Qwen3VLMoeVisionModel.forward(...), which already has a CUDA-graph runner (ViTCudaGraphRunner) gated by SGLANG_VIT_ENABLE_CUDA_GRAPH=1.

In practice CUDA Graph capture for Qwen3.5 failed before this change. Inside VisionAttention.forward the rotary embedding is applied with:

# python/sglang/srt/layers/attention/vision.py
q, k = apply_rotary_pos_emb(q, k, cos, sin)

which on CUDA resolves to:

# python/sglang/srt/layers/rotary_embedding/utils.py
@torch.compile(dynamic=True, backend=get_compiler_backend())
def apply_rotary_pos_emb_native(q, k, cos, sin, unsqueeze_dim=1):
    ...

When this is invoked inside with torch.cuda.graph(graph):, the dynamic-shape guards and RNG state injected by torch.compile violate CUDA Graph capture rules and surface as the random_rng capture failure. The vision graph is never built, so end-to-end ViT latency stays in the unfused, eager-kernel regime.

Modifications

1. New CUDA-graph-safe Triton kernel for NEOX-style vision RoPE:

   • _triton_vision_rope_qk_inplace_kernel and wrapper triton_vision_rope_qk_inplace in python/sglang/srt/layers/rotary_embedding/triton_kernels.py.
   • Operates on q [N, H, D], k [N, KH, D], cos / sin [N, D], in-place; accumulates in fp32 for numerical parity with apply_rotary_pos_emb_native.
   • Triton kernel launches are stateless and contain no Python-side allocations, so they are fully torch.cuda.graph capturable.

2. CUDA dispatch for apply_rotary_pos_emb:

   • New apply_rotary_pos_emb_cuda_triton in python/sglang/srt/layers/rotary_embedding/utils.py.
   • Selected when is_cuda() is true. Silently falls back to the existing apply_rotary_pos_emb_native for non-ViT shapes (e.g. 4D LM-side q/k as in gemma3_causal.py, non-2D cos/sin, mismatched dtypes, non-contiguous strides).
   • CPU and NPU paths are unchanged.

3. Docs updated to advertise Qwen3.5 / Qwen3.5-MoE as supported by ViT CUDA Graph and to document the encoder-server invocation:

   • docs/advanced_features/cuda_graph_for_multi_modal_encoder.md
   • docs_new/docs/advanced_features/cuda_graph_for_multi_modal_encoder.mdx

4. Unit test test/registered/kernels/test_vision_rope_triton.py:

   • Numerical parity vs. apply_rotary_pos_emb_native across (num_tokens, num_heads, num_kv_heads, head_size, dtype) configs covering bf16 / fp16 / fp32 and GQA layouts.
   • In-place semantics check.
   • 4D-layout fallback check.
   • A torch.cuda.CUDAGraph capture + replay correctness test, asserting the exact contract that the old torch.compile path violated.

Accuracy Tests

The kernel is bit-for-bit equivalent (up to fp32 accumulator order) to the existing apply_rotary_pos_emb_native reference. The new test compares against the reference at atol=rtol=1e-2 for bf16/fp16 and 1e-5 for fp32, across multiple shapes and dtypes.

Speed Tests and Profiling

The Triton kernel replaces a torch.compile(dynamic=True)-wrapped sequence of unfused PyTorch ops (q.float(), cos.unsqueeze, rotate_half, q*cos, rotate_half(q)*sin, q_embed.to(orig_q_dtype), ...) with a single fused launch that does the rotation in fp32 and writes back q/k in-place. On top of that, end-to-end the ViT can finally run inside CUDA Graph for Qwen3.5, which collapses the per-layer kernel-launch overhead that dominates small-batch ViT latency on the encoder server.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.