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

Use paired kitchen RoPE in VAE fast path

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 14:58:11 +07:00
parent 8595dd875e
commit ea7712b4b2
1 changed files with 10 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -109,7 +109,16 @@ class Attention(nn.Module):
qkv = self.to_qkv(x).view(batch, sequence, self.heads, 3 * self.dim_head)
query, key, value = qkv.chunk(3, dim=-1)
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)
if os.getenv("H3_VAE_FAST_OPS", "").lower() in {"1", "true", "yes", "on"}:
try:
rot = rotary_pos_emb.shape[-3] * 2
query_rot, key_rot = torch.ops.comfy_kitchen.apply_rope_split_half(query[..., :rot], key[..., :rot], rotary_pos_emb)
query = torch.cat((query_rot, query[..., rot:]), dim=-1)
key = torch.cat((key_rot, key[..., rot:]), dim=-1)
except Exception:
query, key = _apply_rope_split_half(query, rotary_pos_emb), _apply_rope_split_half(key, rotary_pos_emb)
else:
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