Add optional kitchen VAE fast ops

src/h3_blackwell_runtime/vae_decoder.py

@@ -11,6 +11,11 @@ from safetensors import safe_open
 from torch import nn
 from torch.nn import functional as F
 
+try:
+    import comfy_kitchen  # noqa: F401
+except Exception:
+    comfy_kitchen = None
+
 
 IMAGENET_MEAN = (0.485, 0.456, 0.406)
 IMAGENET_STD = (0.229, 0.224, 0.225)
@@ -19,6 +24,10 @@ LATENTS_STD = (1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.75
 
 
 def _rms_norm(x: torch.Tensor, weight: torch.Tensor | None, eps: float) -> torch.Tensor:
+    if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}:
+        if weight is None:
+            return F.rms_norm(x, (x.shape[-1],), eps=eps)
+        return F.rms_norm(x, weight.shape, weight=weight.to(device=x.device, dtype=x.dtype), eps=eps)
     result = x * torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps).to(x.dtype)
     return result if weight is None else result * weight.to(dtype=x.dtype)
 
@@ -72,6 +81,11 @@ class FeedForward(nn.Module):
 
 def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor:
     """Apply the reference split-half RoPE layout to leading rotary channels."""
+    if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}:
+        try:
+            return torch.ops.comfy_kitchen.apply_rope_split_half1(x, table)
+        except Exception:
+            pass
     pairs = table.shape[-3]
     rot = pairs * 2
     first, second = x[..., :pairs], x[..., pairs:rot]