From c8d7f9020c450758bf8150f8054f7c1ab306c952 Mon Sep 17 00:00:00 2001 From: DefTruth Date: Tue, 7 Jul 2026 12:18:12 +0000 Subject: [PATCH] feat: allow attn backend dispatch for boogu --- src/cache_dit/distributed/transformers/boogu_image.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/cache_dit/distributed/transformers/boogu_image.py b/src/cache_dit/distributed/transformers/boogu_image.py index 87e46fb6..e4233091 100644 --- a/src/cache_dit/distributed/transformers/boogu_image.py +++ b/src/cache_dit/distributed/transformers/boogu_image.py @@ -228,6 +228,7 @@ def _patched_sdpa_call( scale=softmax_scale, enable_gqa=enable_gqa, cp_gqa_strategy=cp_gqa_strategy, + backend=getattr(self, '_attention_backend', None), cp_config=cp_config, ) hidden_states = hidden_states.flatten(-2) @@ -293,6 +294,7 @@ def _patched_varlen_call( scale=softmax_scale, enable_gqa=enable_gqa, cp_gqa_strategy=cp_gqa_strategy, + backend=getattr(self, '_attention_backend', None), cp_config=cp_config, ) hidden_states = hidden_states.flatten(-2)