From d556544b30ddf03a6b686a6cdad545dc4b932f16 Mon Sep 17 00:00:00 2001 From: Anai-Guo Date: Sun, 30 Aug 2026 18:20:30 -0700 Subject: [PATCH] fix(lumina): pass to_compute_mask, not use_cache, to LLaDABlock.attention attention() takes (q, k, v, attention_bias, layer_past, to_compute_mask) and has no use_cache parameter. Three of the four call sites still pass use_cache=, which raises TypeError; only LLaDALlamaBlock's non-checkpointed branch -- the path the shipped block_type=llama config takes -- is correct. --- pipelines/lumina_dimmo/lumina_dimoo.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pipelines/lumina_dimmo/lumina_dimoo.py b/pipelines/lumina_dimmo/lumina_dimoo.py index d16ca17ef..9a4b90898 100644 --- a/pipelines/lumina_dimmo/lumina_dimoo.py +++ b/pipelines/lumina_dimmo/lumina_dimoo.py @@ -874,10 +874,10 @@ class LLaDASequentialBlock(LLaDABlock): if self._activation_checkpoint_fn is not None: att, cache = self._activation_checkpoint_fn( # type: ignore - self.attention, q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache + self.attention, q, k, v, attention_bias, layer_past=layer_past ) else: - att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache) + att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past) x = x + self.dropout(att) @@ -967,7 +967,7 @@ class LLaDALlamaBlock(LLaDABlock): if self._activation_checkpoint_fn is not None: att, cache = self._activation_checkpoint_fn( # type: ignore - self.attention, q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache + self.attention, q, k, v, attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask ) else: att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask)