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, 08:30 UTC
[Perf] Reduce H2D pageable memory copies (#38794)
Signed-off-by: jackcfwang <jackcfwang@tencent.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Chaofan Wang committed
Apr 10, 2026 at 15:03 UTC
49d20346e411099a09acd55031b4f1a507bad99c
2 files changed
+62
-25
vllm/v1/attention/backends/triton_attn.py
+20
-18
@@ -85,40 +85,42 @@ class TritonAttentionMetadata:
85
scheduler_metadata: torch.Tensor | None = None
86
prefix_scheduler_metadata: torch.Tensor | None = None
87
mm_prefix_range: dict[int, list[tuple[int, int]]] | None = None
88
+ mm_prefix_range_tensor: torch.Tensor | None = None
89
89
- @property
90
- def mm_prefix_range_tensor(self) -> torch.Tensor | None:
90
+ @staticmethod
91
+ def compute_mm_prefix_range_tensor(
92
+ mm_prefix_range: dict[int, list[tuple[int, int]]] | None,
93
+ num_seqs: int,
94
+ device: torch.device,
95
+ ) -> torch.Tensor | None:
96
"""Convert mm_prefix_range dict to padded tensor for Triton kernel.
97
98
Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges.
99
Empty ranges have start==end==0, which kernel skips via is_valid check.
100
"""
96
- # TODO(Isotr0py): Move to model runner's attention metadata
97
- # preparation to avoid duplicate computation.
98
- if self.mm_prefix_range is None:
101
+ if mm_prefix_range is None:
102
return None
103
101
- num_seqs = self.seq_lens.shape[0]
102
- device = self.seq_lens.device
103
-
104
# Collect ranges, using [(0,0)] for empty sequences to ensure uniform dims
105
range_lists = [
106
- self.mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs)
106
+ mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs)
107
]
108
109
# Return None if all ranges are trivial (only (0,0) placeholders)
110
if all(r == [(0, 0)] for r in range_lists):
111
return None
112
113
- # Create 2D tensors with shape (num_ranges, 2) for each sequence
114
- range_tensors = [
115
- torch.tensor(r, dtype=torch.int32, device=device).view(-1, 2)
116
- for r in range_lists
117
- ]
118
-
119
- return torch.nested.nested_tensor(
120
- range_tensors, layout=torch.jagged
121
- ).to_padded_tensor(0)
113
+ # Build on CPU first then move to GPU in a single H2D transfer
114
+ max_ranges = max(len(r) for r in range_lists)
115
+ # Pad all sequences to the same number of ranges
116
+ padded = []
117
+ for r in range_lists:
118
+ padded_r = list(r) + [(0, 0)] * (max_ranges - len(r))
119
+ padded.append(padded_r)
120
+ # Create tensor with efficient H2D transfer
121
+ return torch.tensor(padded, dtype=torch.int32, device=device).view(
122
+ num_seqs, max_ranges, 2
123
+ )
124
125
126
class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMetadata]):
vllm/v1/worker/gpu_model_runner.py
+42
-7
@@ -2332,13 +2332,8 @@ class GPUModelRunner(
2332
req_idx = self.input_batch.req_id_to_index[req_id]
2333
req_doc_ranges[req_idx] = image_doc_ranges
2334
2335
- if isinstance(attn_metadata, list):
2336
- for ub_metadata in attn_metadata:
2337
- for _metadata in ub_metadata.values():
2338
- _metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined]
2339
- else:
2340
- for _metadata in attn_metadata.values():
2341
- _metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined]
2335
+ # Set mm_prefix_range for all attention metadata
2336
+ self._set_mm_prefix_range_for_metadata(attn_metadata, req_doc_ranges)
2337
2338
if spec_decode_common_attn_metadata is not None and (
2339
num_reqs != num_reqs_padded or num_tokens != num_tokens_padded
@@ -6487,6 +6482,46 @@ class GPUModelRunner(
6482
return
6483
self.reorder_batch_threshold = reduce(min_none_high, reorder_batch_thresholds) # type: ignore[assignment]
6484
6485
+ def _set_mm_prefix_range_for_metadata(
6486
+ self,
6487
+ attn_metadata: Any,
6488
+ req_doc_ranges: dict[int, list[tuple[int, int]]],
6489
+ ) -> None:
6490
+ """Set mm_prefix_range for all attention metadata objects.
6491
+
6492
+ This method handles both list and non-list attention metadata,
6493
+ computing mm_prefix_range_tensor once and sharing it across all
6494
+ metadata objects to avoid redundant host-to-device transfers.
6495
+ """
6496
+ from vllm.v1.attention.backends.triton_attn import (
6497
+ TritonAttentionMetadata,
6498
+ )
6499
+
6500
+ # Get all metadata objects from either list or dict structure
6501
+ metadata_list = []
6502
+ if isinstance(attn_metadata, list):
6503
+ for ub_metadata in attn_metadata:
6504
+ metadata_list.extend(ub_metadata.values())
6505
+ else:
6506
+ metadata_list.extend(attn_metadata.values())
6507
+
6508
+ # Set mm_prefix_range for all metadata and compute tensor once
6509
+ shared_tensor = None
6510
+ for metadata in metadata_list:
6511
+ metadata.mm_prefix_range = req_doc_ranges # type: ignore[attr-defined]
6512
+
6513
+ # Only compute tensor for TritonAttentionMetadata
6514
+ if isinstance(metadata, TritonAttentionMetadata):
6515
+ if shared_tensor is None:
6516
+ shared_tensor = (
6517
+ TritonAttentionMetadata.compute_mm_prefix_range_tensor(
6518
+ req_doc_ranges,
6519
+ metadata.seq_lens.shape[0], # type: ignore[attr-defined]
6520
+ metadata.seq_lens.device, # type: ignore[attr-defined]
6521
+ )
6522
+ )
6523
+ metadata.mm_prefix_range_tensor = shared_tensor
6524
+
6525
def may_reinitialize_input_batch(
6526
self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
6527
) -> None: