[Perf][MLA Sparse] Pin req_id_per_token before non_blocking H2D on XPU and ROCm (#54295)
Signed-off-by: Keyi Li <likey6688@gmail.com> Co-authored-by: Keyi Li <likey6688@gmail.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Keyi Li committed
Aug 28, 2026 at 19:00 UTC
43196f2458bebb822cd825c0370b6f157b66f293
2 files changed
+4
-3
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+2
-1
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.attention.mla_attention import (
17
)
18
from vllm.platforms import current_platform
19
from vllm.triton_utils import tl, triton
20
+from vllm.utils.torch_utils import np_to_pinned_tensor
21
from vllm.v1.attention.backend import (
22
AttentionBackend,
23
AttentionCGSupport,
@@ -515,7 +516,7 @@ class ROCMAiterMLASparseMetadataBuilder(
516
self._prev_req_extent = new_req_extent
517
self._prev_indices_extent = new_indices_extent
518
self.req_id_per_token_buffer[:new_req_extent].copy_(
518
- torch.from_numpy(req_id_per_token), non_blocking=True
519
+ np_to_pinned_tensor(req_id_per_token), non_blocking=True
520
)
521
query_lens = (
522
common_attn_metadata.query_start_loc[1:]
vllm/v1/attention/backends/mla/xpu_mla_sparse.py
+2
-2
@@ -13,7 +13,7 @@ from vllm.logger import init_logger
13
from vllm.model_executor.layers.attention.mla_attention import (
14
get_mla_dims,
15
)
16
-from vllm.utils.torch_utils import is_quantized_kv_cache
16
+from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor
17
from vllm.v1.attention.backend import (
18
AttentionBackend,
19
AttentionCGSupport,
@@ -158,7 +158,7 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]
158
# Zero-fill for cudagraphs
159
self.req_id_per_token_buffer.fill_(0)
160
self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
161
- torch.from_numpy(req_id_per_token), non_blocking=True
161
+ np_to_pinned_tensor(req_id_per_token), non_blocking=True
162
)
163
164
req_id_per_token = self.req_id_per_token_buffer[:num_tokens]