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

Use SageAttention fallback for direct VAE

Browse source
This commit is contained in:
Daniel Maddern 2026-08-14 01:00:16 +07:00
parent 9c9e0857c2
commit f85627a9fd
1 changed files with 6 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
src/h3_blackwell_runtime/vae_decoder.py
View file

@ -102,7 +102,12 @@ class Attention(nn.Module):
output = optimized_attention(query, key, value, self.heads, skip_reshape=True) output = optimized_attention(query, key, value, self.heads, skip_reshape=True)
except Exception: except Exception:
output = F.scaled_dot_product_attention(query, key, value).transpose(1, 2) try:
from sageattention import sageattn
output = sageattn(query, key, value, is_causal=False, tensor_layout="HND", smooth_k=False).transpose(1, 2)
except Exception:
output = F.scaled_dot_product_attention(query, key, value).transpose(1, 2)
return self.to_out(output.reshape(batch, sequence, -1).nan_to_num_(0.0)) return self.to_out(output.reshape(batch, sequence, -1).nan_to_num_(0.0))