Only allocate encoder metadata for encoder-decoder models (#16527)

This commit is contained in:
Minglei Zhu
2026-01-06 22:17:04 -08:00
committed by GitHub
parent 7d757d6f17
commit 4c85f9d039
@@ -325,6 +325,7 @@ class FlashAttentionBackend(AttentionBackend):
and model_runner.model_config.is_encoder_decoder and model_runner.model_config.is_encoder_decoder
), "Sliding window and cross attention are not supported together" ), "Sliding window and cross attention are not supported together"
self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
self.forward_metadata: FlashAttentionMetadata = None self.forward_metadata: FlashAttentionMetadata = None
# extra metadata for handling speculative decoding topk > 1, extended draft decode and verify # extra metadata for handling speculative decoding topk > 1, extended draft decode and verify
self.forward_metadata_spec_decode_expand: FlashAttentionMetadata = None self.forward_metadata_spec_decode_expand: FlashAttentionMetadata = None
@@ -1574,6 +1575,8 @@ class FlashAttentionBackend(AttentionBackend):
), ),
} }
# Only allocate encoder metadata for encoder-decoder models
if self.is_encoder_decoder:
self.encoder_metadata = { self.encoder_metadata = {
"encoder_page_table": torch.zeros( "encoder_page_table": torch.zeros(
max_bs, max_bs,
@@ -1588,6 +1591,9 @@ class FlashAttentionBackend(AttentionBackend):
max_bs + 1, dtype=torch.int32, device=self.device max_bs + 1, dtype=torch.int32, device=self.device
), ),
} }
else:
# For decoder-only models, skip encoder_metadata allocation
self.encoder_metadata = {}
def init_forward_metadata_capture_cuda_graph( def init_forward_metadata_capture_cuda_graph(
self, self,