fix: restore dict return type in create_masks_for_generate (#46962)

src/transformers/masking_utils.py

@@ -1524,12 +1524,10 @@ def create_masks_for_generate(
     # If the attribute exists, we need several masks keyed by layer type.
     if hasattr(effective_config, "layer_types"):
         layer_patterns = set(effective_config.layer_types)
-        # Without a registered attention-mask function, defer to the model by returning the raw attention mask
-        if any(layer_type not in LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING for layer_type in layer_patterns):
-            return attention_mask
         causal_masks = {}
         for layer_pattern in layer_patterns:
-            causal_masks[layer_pattern] = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING[layer_pattern](**mask_kwargs)
+            mask_fn = LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING.get(layer_pattern)
+            causal_masks[layer_pattern] = mask_fn(**mask_kwargs) if mask_fn is not None else attention_mask
         return causal_masks
     # In this case, all layers are sliding
     elif getattr(effective_config, "sliding_window", None) is not None:

tests/utils/test_masking_utils.py

@@ -380,19 +380,26 @@ def test_create_masks_for_generate_defers_for_unmapped_layer_types(self):
         `create_masks_for_generate` pre-builds attention masks for compilable caches by mapping each
         `config.layer_types` entry through `LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING`. Hybrid models with a
         non-attention layer type (e.g. ``moe`` / ``mlp`` in Nemotron-H) have no mask function — the helper must
-        not raise `KeyError`; instead the raw attention mask is returned so the model can build its own masks.
+        not raise `KeyError`; instead the unmapped layer type defers to the raw attention mask while mapped
+        types still get their pre-computed masks, preserving the dict return type contract.
         """
         config = LlamaConfig(num_hidden_layers=2)
         config.layer_types = ["full_attention", "moe"]
         attention_mask = torch.ones((1, 5), dtype=torch.long, device=torch_device)
 
-        # Must not raise (previously `KeyError: 'moe'`), and defers the raw mask to the model.
+        # Must not raise (previously `KeyError: 'moe'`), and returns a dict keyed by layer type.
         out = create_masks_for_generate(
             config, inputs_embeds=None, attention_mask=attention_mask, past_key_values=None
         )
-        self.assertIs(out, attention_mask)
-
-        # `None` (no padding) is deferred too.
-        self.assertIsNone(
-            create_masks_for_generate(config, inputs_embeds=None, attention_mask=None, past_key_values=None)
-        )
+        self.assertIsInstance(out, dict)
+        self.assertEqual(set(out.keys()), {"full_attention", "moe"})
+        # Unmapped layer type defers the raw attention mask to the model.
+        self.assertIs(out["moe"], attention_mask)
+        # Mapped layer type gets its pre-computed mask (None for a full mask with no padding).
+        self.assertIsNone(out["full_attention"])
+
+        # `None` (no padding) is deferred for unmapped types, and mapped types compute None too.
+        out_none = create_masks_for_generate(config, inputs_embeds=None, attention_mask=None, past_key_values=None)
+        self.assertIsInstance(out_none, dict)
+        self.assertIsNone(out_none["moe"])
+        self.assertIsNone(out_none["full_attention"])