Only allocate encoder metadata for encoder-decoder models (#16527)
This commit is contained in:
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user