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)