[NPU] Add head_dim=256 to _can_use_tnd whitelist (#28635)

2026-06-18 17:55:50 +08:00
parent 2a9cce5d27
commit bb9d31f22d
1 changed files with 1 additions and 1 deletions
@@ -376,7 +376,7 @@ class AscendAttnBackend(AttentionBackend):
        """Check if TND layout is supported."""
        d = layer.qk_head_dim
        v = layer.v_head_dim
-        return (d == v and d in (128, 192)) or (d == 192 and v == 128)
+        return (d == v and d in (128, 192, 256)) or (d == 192 and v == 128)

    def get_verify_buffers_to_fill_after_draft(self):
        """