[Kernel] Fuse Q/K RMSNorm and RoPE (#5990)

Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com>
Signed-off-by: lishunyang12 <lishunyang12@163.com>
Signed-off-by: Mengjie Zhao <zmj0129@gmail.com>
Co-authored-by: lishunyang12 <lishunyang12@163.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

vllm_omni/diffusion/cache/teacache/extractors.py

@@ -1263,6 +1263,7 @@ def extract_minimax_h3_context(
         _BF16_DTYPE,
         _FORWARD_SUPPORTED_KWARGS,
         MINIMAX_H3_ADALN_MODALITY_NUM,
+        _build_rope_table,
         _modulate_scale_shift,
         _required_kwarg,
     )
@@ -1318,7 +1319,7 @@ def extract_minimax_h3_context(
         raise ValueError(f"inverse_indices must be [{seq_len}], got {list(inverse_indices.shape)}")
     device = x.device
 
-    rope_freqs = module.rope(img_position_ids).to(device)
+    rope_table = _build_rope_table(module.rope(img_position_ids).to(device))
 
     decoder_input, t_emb = module._embed(
         x=x,
@@ -1351,7 +1352,7 @@ def extract_minimax_h3_context(
     def run_transformer_blocks() -> tuple[torch.Tensor, ...]:
         hidden, block_rope, block_combined = module.sp_prepare(
             decoder_input,
-            rope_freqs,
+            rope_table,
             combined_indices,
         )
         for block in module.blocks:
@@ -1359,7 +1360,7 @@ def extract_minimax_h3_context(
                 hidden,
                 t_emb=t_emb,
                 combined_indices=block_combined,
-                rope_freqs=block_rope,
+                rope_table=block_rope,
                 cu_seqlens=cu_seqlens,
                 max_seqlen=max_seqlen,
                 packed_total=seq_len,