A high-throughput and memory-efficient inference and serving engine for LLMs
[Perf][Attention] Pin MLA chunked-context metadata tensors so H2D copies are truly non-blocking (#45074)
Signed-off-by: zixi-qi <zixi@inferact.ai>
qizixi committed
Jun 10, 2026 at 15:56 UTC
e2db0222e9b2ff37f7a1c54a1d8043e61f8b0b8d
1 file changed
+5
-3
vllm/model_executor/layers/attention/mla_attention.py
+5
-3
@@ -1664,7 +1664,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1664
.unsqueeze(1)
1665
.expand(-1, num_prefills)
1666
* max_context_chunk
1667
- )
1667
+ ).pin_memory()
1668
chunk_ends = torch.min(
1669
context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk
1670
)
@@ -1680,7 +1680,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1680
1681
max_token_num_over_chunk = chunk_total_token.max().item()
1682
token_to_seq_tensor_cpu = torch.zeros(
1683
- [num_chunks, max_token_num_over_chunk], dtype=torch.int32
1683
+ [num_chunks, max_token_num_over_chunk],
1684
+ dtype=torch.int32,
1685
+ pin_memory=True,
1686
)
1687
range_idx = torch.arange(num_prefills, dtype=torch.int32)
1688
for i in range(num_chunks):
@@ -1724,7 +1726,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1726
.unsqueeze(1)
1727
.expand(-1, num_prefills)
1728
* padded_local_max_context_chunk_across_ranks
1727
- )
1729
+ ).pin_memory()
1730
local_chunk_ends = torch.min(
1731
padded_local_context_lens_cpu.unsqueeze(0),
1732
local_chunk_starts