Repository navigation
fix: restore dict return type in create_masks_for_generate (#46962) - #47267
Closed
davidnichols-ops wants to merge 3 commits into
Closed
davidnichols-ops wants to merge 3 commits into
davidnichols-ops wants to merge 3 commits into
Conversation
…types (#46962) PR #46446 added an early return `return attention_mask` when any layer_type is not in LAYER_PATTERN_TO_MASK_FUNCTION_MAPPING. This broke the dict return type contract that models like Qwen3 rely on via `isinstance(attention_mask, dict)`, causing unnecessary mask recomputation and performance regression. Replace the early return with inline handling: unmapped types get the raw attention_mask as their dict value, so the model still receives a dict and can use pre-computed masks for mapped types. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
The test expected the old bare-tensor return when any layer_type was unmapped, but PR #47267 changed create_masks_for_generate to always return a dict when layer_types exists (preserving the contract that hybrid models rely on via isinstance(attention_mask, dict)). Update the assertions to verify the dict has the correct keys, unmapped types defer the raw mask, and mapped types get pre-computed masks. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
The previous CI run had 0 test failures out of 152,657 tests. The only failing job was examples_torch, which exited with code 137 (OOM kill) during test_generation — a runner memory issue unrelated to this PR. Generated with [Devin](https://devin.ai) Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Contributor
CI recapDashboard: View test results in Grafana |
Author
|
Thanks for taking a look. Could you share the reason for closing this? If there’s a technical issue with the approach or if it doesn’t align with project policy, I’d appreciate the feedback. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
cache_implementation=static#46446 added an early returnreturn attention_maskincreate_masks_for_generatewhen anylayer_typeis not inLAYER_PATTERN_TO_MASK_FUNCTION_MAPPING. This broke the dict return type contract that models like Qwen3, Cohere2, Bamba, DeepseekV32, etc. rely on viaisinstance(attention_mask, dict), causing unnecessary mask recomputation and performance regression (issue [Qwen3]create_masks_for_generatereturn type change after #41251 #46962).attention_maskas their dict value, so the model still receives a dict and can use pre-computed masks for mapped types while deferring to the model for unmapped types.Deleted: the early return
if any(layer_type not in ...): return attention_mask(3 lines) — replaced with a.get()+ conditional in the existing loop, preserving the dict return type contract that all hybrid models depend on.Test plan
full_attention) → returnsdictwithfull_attentionkey (was bare value before)full_attention+sliding_attention) → returnsdictwith both keyslayer_typesattribute → returns bare value (unchanged behavior)Generated with Devin