daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use fused VAE RMS RoPE when available

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 00:38:05 +07:00
parent a262d44632
commit e96af86e72
1 changed files with 10 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -94,8 +94,16 @@ class Attention(nn.Module):
batch, sequence, _ = x.shape batch, sequence, _ = x.shape
qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head) qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head)
query, key, value = qkv.chunk(3, dim=-1) query, key, value = qkv.chunk(3, dim=-1)
query, key = self.norm_q(query), self.norm_k(key) try:
query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb) 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)) output = F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2))
return self.to_out(output.transpose(1, 2).reshape(batch, sequence, -1).nan_to_num_(0.0)) return self.to_out(output.transpose(1, 2).reshape(batch, sequence, -1).nan_to_num_(0.0))