Skip to content

Commit 6b216e2

Browse files
HJSangclaude
andcommitted
test: give the Qwen3 attention stand-in a non-flash name
transformers resolves any *flash* attn_implementation to a FlashAttention kernel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1 parent b2f4e85 commit 6b216e2

1 file changed

Lines changed: 4 additions & 3 deletions

File tree

‎tests/fast/backends/test_fsdp_hf_compat.py‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -371,12 +371,13 @@ def test_qwen3_attention_receives_the_packed_boundaries(monkeypatch):
371371

372372
seen = []
373373

374-
def flash_stand_in(module, query, key, value, attention_mask, scaling=None, **kwargs):
374+
def attention_stand_in(module, query, key, value, attention_mask, scaling=None, **kwargs):
375375
seen.append(kwargs)
376376
out = torch.nn.functional.scaled_dot_product_attention(query, key, value, scale=scaling, enable_gqa=True)
377377
return out.transpose(1, 2), None
378378

379-
monkeypatch.setitem(AttentionInterface._global_mapping, "flash_stand_in", flash_stand_in)
379+
# not a "flash" name, which transformers would resolve to a real FlashAttention kernel
380+
monkeypatch.setitem(AttentionInterface._global_mapping, "attention_stand_in", attention_stand_in)
380381
cfg = Qwen3Config(
381382
vocab_size=64,
382383
hidden_size=16,
@@ -386,7 +387,7 @@ def flash_stand_in(module, query, key, value, attention_mask, scaling=None, **kw
386387
num_key_value_heads=1,
387388
head_dim=8,
388389
)
389-
cfg._attn_implementation = "flash_stand_in"
390+
cfg._attn_implementation = "attention_stand_in"
390391
model = Qwen3ForCausalLM(cfg).eval()
391392

392393
with torch.no_grad():

0 commit comments

Comments
 (0)