[Bugfix] Preserve tensor causal metadata for grouped attention (#48135)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Codex <codex@openai.com>
Lucas Wilkinson committed
Jul 9, 2026 at 13:57 UTC
e08a9151468190575114de1c996275b993ec940a
1 file changed
+4
-2
vllm/v1/worker/gpu/attn_utils.py
+4
-2
@@ -573,7 +573,7 @@ def build_attn_metadata(
573
mm_req_doc_ranges: dict[int, list[tuple[int, int]]] | None = None,
574
model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
575
for_cudagraph_capture: bool = False,
576
- causal: bool | Mapping[int, bool] = True,
576
+ causal: bool | torch.Tensor | Mapping[int, bool] = True,
577
rswa_prefix_lens: torch.Tensor | None = None,
578
) -> dict[str, Any]:
579
seq_lens = seq_lens[:num_reqs]
@@ -588,7 +588,9 @@ def build_attn_metadata(
588
block_table = block_tables[i]
589
slot_mapping = slot_mappings[i]
590
# Per-group causal for hybrid drafters (mixed SWA/full attention).
591
- group_causal = causal if isinstance(causal, bool) else causal.get(i, True)
591
+ group_causal = (
592
+ causal if isinstance(causal, (bool, torch.Tensor)) else causal.get(i, True)
593
+ )
594
595
common_attn_metadata_extra_kwargs = (
596
model_specific_attn_metadata.get_extra_common_attn_kwargs(i, num_reqs)