A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 09:45 UTC
[Model Runner V2] Fix flex attention kv blocks calculation issue (#39353)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 9, 2026 at 13:07 UTC
56e19d7ee20635f87e04089bfaa2f54d52db65e9
1 file changed
+6
-9
vllm/v1/attention/backends/flex_attention.py
+6
-9
@@ -750,11 +750,11 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
750
self.max_model_len = self.model_config.max_model_len
751
max_num_seqs = vllm_config.scheduler_config.max_num_seqs
752
max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens
753
- self.max_num_q_block = (
754
- self.max_model_len + self.q_block_size - 1
755
- ) // self.q_block_size
753
+ self.max_num_query_groups = cdiv(max_num_batched_tokens, self.q_block_size)
754
+ max_num_pages_per_seq = cdiv(self.max_model_len, self.block_size)
755
+ self.max_num_kv_indices = self.q_block_size * max_num_pages_per_seq
756
self.persistent_kv_num_blocks = torch.empty(
757
- self.max_num_q_block, dtype=torch.int32, device=device
757
+ self.max_num_query_groups, dtype=torch.int32, device=device
758
)
759
self.persistent_offset_tensor = torch.empty(
760
max_num_seqs, dtype=torch.int32, device=device
@@ -828,12 +828,9 @@ class FlexAttentionMetadataBuilder(AttentionMetadataBuilder[FlexAttentionMetadat
828
)
829
830
if self.persistent_kv_indices is None:
831
- max_num_kv_block = (
832
- self.max_model_len + self.kv_block_size - 1
833
- ) // self.kv_block_size
831
self.persistent_kv_indices = torch.empty(
835
- self.max_model_len,
836
- max_num_kv_block,
832
+ self.max_num_query_groups,
833
+ self.max_num_kv_indices,
834
dtype=torch.int32,
835
device=self.device,
836
)