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 4, 11:15 UTC
[Bugfix][MLA] Size arange_buffer to max_num_batched_tokens to prevent CUDA IMA (#39277)
Signed-off-by: UranusSeven <109661872+UranusSeven@users.noreply.github.com> Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai>
Uranus committed
Apr 30, 2026 at 07:14 UTC
a966aaed30b9132191e8ba88d4f61e76657b690d
1 file changed
+4
-1
vllm/v1/attention/backends/mla/indexer.py
+4
-1
@@ -303,7 +303,10 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):
303
device=self.device,
304
)
305
self.arange_buffer = torch.arange(
306
- scheduler_config.max_num_seqs * next_n,
306
+ max(
307
+ scheduler_config.max_num_seqs * next_n,
308
+ scheduler_config.max_num_batched_tokens,
309
+ ),
310
dtype=torch.int32,
311
device=self.device,
312
)