From f85627a9fd09c31cf94d928a828beb76e4289d7f Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Fri, 14 Aug 2026 01:00:16 +0700 Subject: [PATCH] Use SageAttention fallback for direct VAE --- src/h3_blackwell_runtime/vae_decoder.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/h3_blackwell_runtime/vae_decoder.py b/src/h3_blackwell_runtime/vae_decoder.py index bed1238..18ebb05 100644 --- a/src/h3_blackwell_runtime/vae_decoder.py +++ b/src/h3_blackwell_runtime/vae_decoder.py @@ -102,7 +102,12 @@ class Attention(nn.Module): output = optimized_attention(query, key, value, self.heads, skip_reshape=True) 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))