From e4f145e9d51434383bf70bed3e8dacf9b05f90d6 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Fri, 14 Aug 2026 00:50:03 +0700 Subject: [PATCH] Mirror upstream VAE block math --- src/h3_blackwell_runtime/vae_decoder.py | 20 ++++++-------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index 81aa42d..ecaced5 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -67,7 +67,7 @@ class FeedForward(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: gate, value = self.w1(x).chunk(2, dim=-1) - return self.w2(F.silu(gate) * value) + return self.w2(F.silu(gate).mul_(value)) def _apply_rope_split_half(x: torch.Tensor, table: torch.Tensor) -> torch.Tensor: @@ -94,16 +94,8 @@ 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) + 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) query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2) try: from comfy.ldm.modules.attention import optimized_attention @@ -126,8 +118,8 @@ class TransformerBlock(nn.Module): self.scale2 = nn.Parameter(torch.empty(dim, device=device)) def forward(self, x: torch.Tensor, rotary_pos_emb: torch.Tensor) -> torch.Tensor: - x = x + self.attn(self.norm1(x), rotary_pos_emb) * self.scale1.to(x.dtype) - return x + self.ff(self.norm2(x)) * self.scale2.to(x.dtype) + x = x.addcmul_(self.attn(self.norm1(x), rotary_pos_emb), self.scale1.to(x.dtype)) + return x.addcmul_(self.ff(self.norm2(x)), self.scale2.to(x.dtype)) class ViT3DDecoder(nn.Module): @@ -156,7 +148,7 @@ class ViT3DDecoder(nn.Module): h = block(h, rope) output = self.proj_out(self.norm_out(h))[:, :patches] output = output.view(batch, latent_t, latent_h, latent_w, self.out_channels, self.patch_size_t, self.patch_size, self.patch_size) - return output.permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size) + return output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous().reshape(batch, self.out_channels, latent_t * self.patch_size_t, latent_h * self.patch_size, latent_w * self.patch_size) class MiniMaxH3VideoVAE(nn.Module):