Use fused VAE RMS RoPE when available

src/h3_blackwell_runtime/vae_decoder.py

@@ -94,6 +94,14 @@ class Attention(nn.Module):
         batch, sequence, _ = x.shape
         qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head)
         query, key, value = qkv.chunk(3, dim=-1)
+        try:
+            import comfy_kitchen  # Registers the standalone CUDA extension operators.
+
+            del comfy_kitchen
+            query, key = query.contiguous(), key.contiguous()
+            weight = torch.ones(self.dim_head, device=query.device, dtype=query.dtype)
+            torch.ops.comfy_kitchen.rms_rope_split_half_(query, key, rotary_pos_emb, weight, weight, self.norm_q.eps, rotary_pos_emb.shape[-3] * 2)
+        except Exception:
             query, key = self.norm_q(query), self.norm_k(key)
             query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb)
         output = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))