@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a28b94e6ef
Fork 0

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 9, 17:06 UTC

[Performance] Split FlashAttn attention and cache update (#25954)

Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Luka Govedič <luka.govedic@gmail.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: mgoin <mgoin64@gmail.com> Co-authored-by: Varun Sundar Rabindranath <varunsundar08@gmail.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <luka.govedic@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com>

ElizaWszola committed Jan 24, 2026 at 02:28 UTC a28b94e6ef60b7f5aa1b97bc8d966a8d12cbc1da
21 files changed +458 -68
tests/v1/attention/test_attention_backends.py
+5
@@ -13,6 +13,7 @@ from tests.v1.attention.utils import (
13 create_common_attn_metadata,
14 create_standard_kv_cache_spec,
15 create_vllm_config,
16 + try_backend_includes_kv_cache_update,
17 try_get_attention_backend,
18 )
19 from vllm.config import ModelConfig
@@ -295,6 +296,10 @@ def run_attention_backend(
296 # Run forward pass
297 # NOTE: The query, key, and value are already shaped correctly
298 # in the calling test function.
299 + if not try_backend_includes_kv_cache_update(actual_backend):
300 + impl.do_kv_cache_update(
301 + mock_layer, key, value, kv_cache, attn_metadata.slot_mapping
302 + )
303 output = impl.forward(
304 mock_layer, query, key, value, kv_cache, attn_metadata, output=output
305 )
tests/v1/attention/utils.py
+12
@@ -130,6 +130,18 @@ def try_get_attention_backend(
130 raise AssertionError("unreachable") from None
131
132
133 +def try_backend_includes_kv_cache_update(
134 + backend: AttentionBackendEnum,
135 +) -> bool:
136 + """Try to get the attention backend class, skipping test if not found."""
137 + try:
138 + backend_class = backend.get_class()
139 + return backend_class.forward_includes_kv_cache_update
140 + except ImportError as e:
141 + pytest.skip(f"{backend.name} not available: {e}")
142 + raise AssertionError("unreachable") from None
143 +
144 +
145 def create_standard_kv_cache_spec(vllm_config: VllmConfig) -> FullAttentionSpec:
146 """Create a FullAttentionSpec from ModelParams only."""
147 return FullAttentionSpec(
tests/v1/kv_connector/unit/test_decode_bench_connector.py
+1 -1
@@ -86,7 +86,7 @@ class DecodeBenchTestRunner:
86 self._block_hasher = get_request_block_hasher(block_size, sha256)
87
88 self._dummy_ctx: ForwardContext = ForwardContext(
89 - no_compile_layers={}, attn_metadata={}, virtual_engine=0
89 + no_compile_layers={}, attn_metadata={}, virtual_engine=0, slot_mapping={}
90 )
91
92 def new_request(self, token_ids: list[int]) -> Request:
tests/v1/kv_connector/unit/test_nixl_connector.py
+8
@@ -548,6 +548,7 @@ class TestNixlHandshake:
548 no_compile_layers={},
549 attn_metadata={},
550 virtual_engine=0,
551 + slot_mapping={},
552 )
553 _before_load = time.perf_counter()
554 connector.start_load_kv(dummy_ctx)
@@ -618,6 +619,7 @@ class TestNixlHandshake:
619 no_compile_layers={},
620 attn_metadata={},
621 virtual_engine=0,
622 + slot_mapping={},
623 )
624 _before_load = time.perf_counter()
625 connector.start_load_kv(dummy_ctx)
@@ -844,6 +846,7 @@ class TestNixlHandshake:
846 no_compile_layers={},
847 attn_metadata={},
848 virtual_engine=0,
849 + slot_mapping={},
850 )
851 _before_load = time.perf_counter()
852 connector.start_load_kv(dummy_ctx)
@@ -1006,6 +1009,7 @@ def test_kv_connector_stats(default_vllm_config, dist_init):
1009 no_compile_layers={},
1010 attn_metadata={},
1011 virtual_engine=0,
1012 + slot_mapping={},
1013 )
1014 connector.start_load_kv(dummy_ctx)
1015
@@ -1767,6 +1771,7 @@ def test_aborted_request_removed_from_worker_in_batch(default_vllm_config, dist_
1771 no_compile_layers={},
1772 attn_metadata={},
1773 virtual_engine=0,
1774 + slot_mapping={},
1775 )
1776 connector.start_load_kv(dummy_ctx)
1777
@@ -1917,6 +1922,7 @@ def test_transfer_failure_logging(
1922 no_compile_layers={},
1923 attn_metadata={},
1924 virtual_engine=0,
1925 + slot_mapping={},
1926 )
1927
1928 # Capture logs from the nixl_connector logger specifically
@@ -2017,6 +2023,7 @@ def test_handshake_failure_returns_finished(default_vllm_config, dist_init):
2023 no_compile_layers={},
2024 attn_metadata={},
2025 virtual_engine=0,
2026 + slot_mapping={},
2027 )
2028 connector.start_load_kv(dummy_ctx)
2029
@@ -2067,6 +2074,7 @@ def test_transfer_setup_failure_returns_finished(default_vllm_config, dist_init)
2074 no_compile_layers={},
2075 attn_metadata={},
2076 virtual_engine=0,
2077 + slot_mapping={},
2078 )
2079 connector.start_load_kv(dummy_ctx)
2080
tests/v1/kv_connector/unit/test_offloading_connector.py
+4 -1
@@ -209,7 +209,10 @@ class RequestRunner:
209 self._block_hasher = get_request_block_hasher(gpu_block_size, sha256)
210
211 self._dummy_ctx: ForwardContext = ForwardContext(
212 - no_compile_layers={}, attn_metadata={}, virtual_engine=0
212 + no_compile_layers={},
213 + attn_metadata={},
214 + virtual_engine=0,
215 + slot_mapping={},
216 )
217
218 def new_request(self, token_ids: list[int]):
tests/v1/spec_decode/test_tree_attention.py
+9
@@ -9,6 +9,7 @@ import torch
9 from tests.v1.attention.utils import (
10 create_standard_kv_cache_spec,
11 create_vllm_config,
12 + try_backend_includes_kv_cache_update,
13 try_get_attention_backend,
14 )
15 from vllm.config import ParallelConfig, SpeculativeConfig
@@ -120,6 +121,14 @@ def forward_attention(
121 key = k.view(-1, num_kv_heads, dim_per_head)
122 value = v.view(-1, num_kv_heads, dim_per_head)
123 output = torch.empty_like(query)
124 + if not try_backend_includes_kv_cache_update(backend):
125 + instance.do_kv_cache_update(
126 + layer=layer,
127 + key=key,
128 + value=value,
129 + kv_cache=kv_cache,
130 + slot_mapping=attn_metadata.slot_mapping,
131 + )
132 return instance.forward(
133 layer=layer,
134 query=query,
vllm/attention/layer.py
+85 -16
@@ -390,29 +390,43 @@ class Attention(nn.Module, AttentionLayerBase):
390 if value is not None:
391 value = value.view(-1, self.num_kv_heads, self.head_size_v)
392 if self.use_direct_call:
393 - forward_context: ForwardContext = get_forward_context()
394 - attn_metadata = forward_context.attn_metadata
395 - if isinstance(attn_metadata, dict):
396 - attn_metadata = attn_metadata[self.layer_name]
397 - self_kv_cache = self.kv_cache[forward_context.virtual_engine]
398 - self.impl.forward(
399 - self, query, key, value, self_kv_cache, attn_metadata, output=output
393 + kv_cache_dummy_dep = None
394 + if not self.attn_backend.forward_includes_kv_cache_update:
395 + kv_cache_dummy_dep = unified_kv_cache_update(
396 + key, value, self.layer_name
397 + )
398 + unified_attention_with_output(
399 + query,
400 + key,
401 + value,
402 + output,
403 + self.layer_name,
404 + kv_cache_dummy_dep=kv_cache_dummy_dep,
405 )
406 else:
407 + kv_cache_dummy_dep = None
408 + if not self.attn_backend.forward_includes_kv_cache_update and (
409 + # torch can only dispatch custom op if a tensor is passed
410 + key is not None or value is not None
411 + ):
412 + kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
413 + key, value, self.layer_name
414 + )
415 torch.ops.vllm.unified_attention_with_output(
403 - query, key, value, output, self.layer_name
416 + query,
417 + key,
418 + value,
419 + output,
420 + self.layer_name,
421 + kv_cache_dummy_dep=kv_cache_dummy_dep,
422 )
423 return output.view(-1, hidden_size)
424 else:
425 + assert self.attn_backend.forward_includes_kv_cache_update, (
426 + "Split KV cache update not supported when output tensor not provided."
427 + )
428 if self.use_direct_call:
408 - forward_context = get_forward_context()
409 - attn_metadata = forward_context.attn_metadata
410 - if isinstance(attn_metadata, dict):
411 - attn_metadata = attn_metadata[self.layer_name]
412 - self_kv_cache = self.kv_cache[forward_context.virtual_engine]
413 - return self.impl.forward(
414 - self, query, key, value, self_kv_cache, attn_metadata
415 - )
429 + return unified_attention(query, key, value, self.layer_name)
430 else:
431 return torch.ops.vllm.unified_attention(
432 query, key, value, self.layer_name
@@ -802,6 +816,55 @@ direct_register_custom_op(
816 )
817
818
819 +def unified_kv_cache_update(
820 + key: torch.Tensor,
821 + value: torch.Tensor,
822 + layer_name: str,
823 +) -> torch.Tensor:
824 + """
825 + Returns a dummy that is passed to unified_attention to signal a side effect and
826 + the data dependency between them to ensure torch.compile preserves ordering.
827 + """
828 + forward_context = get_forward_context()
829 + attn_layer = forward_context.no_compile_layers[layer_name]
830 + kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
831 +
832 + slot_mapping = forward_context.slot_mapping
833 + assert isinstance(slot_mapping, dict), (
834 + f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
835 + )
836 + layer_slot_mapping = slot_mapping.get(layer_name)
837 + if layer_slot_mapping is not None:
838 + assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
839 + f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
840 + )
841 + attn_layer.impl.do_kv_cache_update(
842 + attn_layer,
843 + key,
844 + value,
845 + kv_cache,
846 + layer_slot_mapping,
847 + )
848 +
849 + return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)
850 +
851 +
852 +def unified_kv_cache_update_fake(
853 + key: torch.Tensor,
854 + value: torch.Tensor,
855 + layer_name: str,
856 +) -> torch.Tensor:
857 + return torch.empty(0, device=key.device, dtype=key.dtype)
858 +
859 +
860 +direct_register_custom_op(
861 + op_name="unified_kv_cache_update",
862 + op_func=unified_kv_cache_update,
863 + fake_impl=unified_kv_cache_update_fake,
864 + mutates_args=[],
865 +)
866 +
867 +
868 @maybe_transfer_kv_layer
869 def unified_attention_with_output(
870 query: torch.Tensor,
@@ -811,7 +874,12 @@ def unified_attention_with_output(
874 layer_name: str,
875 output_scale: torch.Tensor | None = None,
876 output_block_scale: torch.Tensor | None = None,
877 + kv_cache_dummy_dep: torch.Tensor | None = None,
878 ) -> None:
879 + # kv_cache_dummy_dep is not used but accepting it creates a data dependency
880 + # that ensures torch.compile preserves ordering between KV cache update and
881 + # attention forward.
882 + del kv_cache_dummy_dep
883 attn_metadata, self, kv_cache = get_attention_context(layer_name)
884
885 self.impl.forward(
@@ -835,6 +903,7 @@ def unified_attention_with_output_fake(
903 layer_name: str,
904 output_scale: torch.Tensor | None = None,
905 output_block_scale: torch.Tensor | None = None,
906 + kv_cache_dummy_dep: torch.Tensor | None = None,
907 ) -> None:
908 return
909
vllm/forward_context.py
+5
@@ -189,6 +189,7 @@ class ForwardContext:
189 # copy from vllm_config.compilation_config.static_forward_context
190 no_compile_layers: dict[str, Any]
191 attn_metadata: dict[str, AttentionMetadata] | list[dict[str, AttentionMetadata]]
192 + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]]
193 """
194 Type Dict[str, AttentionMetadata] for v1, map from layer_name of each
195 attention layer to its attention metadata
@@ -266,6 +267,7 @@ def create_forward_context(
267 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
268 batch_descriptor: BatchDescriptor | None = None,
269 ubatch_slices: UBatchSlices | None = None,
270 + slot_mapping: dict[str, torch.Tensor] | None = None,
271 additional_kwargs: dict[str, Any] | None = None,
272 skip_compiled: bool = False,
273 ):
@@ -282,6 +284,7 @@ def create_forward_context(
284 remaining_moe_layers=remaining_moe_layers,
285 virtual_engine=virtual_engine,
286 attn_metadata=attn_metadata,
287 + slot_mapping=slot_mapping or {},
288 dp_metadata=dp_metadata,
289 cudagraph_runtime_mode=cudagraph_runtime_mode,
290 batch_descriptor=batch_descriptor,
@@ -316,6 +319,7 @@ def set_forward_context(
319 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
320 batch_descriptor: BatchDescriptor | None = None,
321 ubatch_slices: UBatchSlices | None = None,
322 + slot_mapping: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None,
323 skip_compiled: bool = False,
324 ):
325 """A context manager that stores the current forward context,
@@ -374,6 +378,7 @@ def set_forward_context(
378 cudagraph_runtime_mode,
379 batch_descriptor,
380 ubatch_slices,
381 + slot_mapping,
382 additional_kwargs,
383 skip_compiled,
384 )
vllm/model_executor/layers/attention/cross_attention.py
+48 -5
@@ -15,7 +15,7 @@ from vllm.v1.attention.backend import (
15 AttentionMetadata,
16 AttentionType,
17 CommonAttentionMetadata,
18 - subclass_attention_backend,
18 + subclass_attention_backend_with_overrides,
19 )
20 from vllm.v1.attention.selector import get_attn_backend
21 from vllm.v1.kv_cache_interface import CrossAttentionSpec, KVCacheSpec
@@ -72,6 +72,7 @@ def create_cross_attention_backend(
72 ) -> type[AttentionBackend]:
73 prefix = "CrossAttention_"
74 underlying_builder = underlying_attn_backend.get_builder_cls()
75 + underlying_impl = underlying_attn_backend.get_impl_cls()
76
77 class CrossAttentionBuilder(underlying_builder): # type: ignore
78 def build(
@@ -106,18 +107,60 @@ def create_cross_attention_backend(
107 )
108
109 # NOTE (NickLucche) use `new_metadata` instead of `common_*` (initial) here
109 - new_metadata.slot_mapping = _get_cross_slot_mapping(
110 + slot_mapping = _get_cross_slot_mapping(
111 new_metadata.encoder_seq_lens_cpu,
112 new_metadata.block_table_tensor,
113 self.kv_cache_spec,
114 self.device,
115 )
115 - return super().build(common_prefix_len, new_metadata, fast_build)
116 + attn_metadata = super().build(common_prefix_len, new_metadata, fast_build)
117 + attn_metadata.slot_mapping = slot_mapping
118 + return attn_metadata
119 +
120 + # NOTE(Lucas): we need a custom impl so we can use the slot-mapping computed by
121 + # `CrossAttentionBuilder` instead of the one computed by `BlockTable`
122 + # (gpu_model_runner)
123 + class CrossAttentionImpl(underlying_impl): # type: ignore[valid-type,misc]
124 + def forward(
125 + self,
126 + layer: torch.nn.Module,
127 + query: torch.Tensor,
128 + key: torch.Tensor,
129 + value: torch.Tensor,
130 + kv_cache: torch.Tensor,
131 + attn_metadata: AttentionMetadata,
132 + output: torch.Tensor | None = None,
133 + output_scale: torch.Tensor | None = None,
134 + output_block_scale: torch.Tensor | None = None,
135 + ) -> torch.Tensor:
136 + if (
137 + not underlying_attn_backend.forward_includes_kv_cache_update
138 + and attn_metadata is not None
139 + ):
140 + self.do_kv_cache_update(
141 + layer, key, value, kv_cache, attn_metadata.slot_mapping
142 + )
143 +
144 + return super().forward(
145 + layer,
146 + query,
147 + key,
148 + value,
149 + kv_cache,
150 + attn_metadata,
151 + output,
152 + output_scale,
153 + output_block_scale,
154 + )
155
117 - attn_backend = subclass_attention_backend(
156 + attn_backend = subclass_attention_backend_with_overrides(
157 name_prefix=prefix,
158 attention_backend_cls=underlying_attn_backend,
120 - builder_cls=CrossAttentionBuilder,
159 + overrides={
160 + "get_builder_cls": lambda: CrossAttentionBuilder,
161 + "get_impl_cls": lambda: CrossAttentionImpl,
162 + "forward_includes_kv_cache_update": True,
163 + },
164 )
165
166 return attn_backend
vllm/utils/torch_utils.py
+2 -1
@@ -613,8 +613,9 @@ def weak_ref_tensor(tensor: Any) -> Any:
613 Create a weak reference to a tensor.
614 The new tensor will share the same data as the original tensor,
615 but will not keep the original tensor alive.
616 + This ignores 0-size tensors as those don't allocate any memory.
617 """
617 - if isinstance(tensor, torch.Tensor):
618 + if isinstance(tensor, torch.Tensor) and tensor.numel() > 0:
619 return torch.ops._C.weak_ref_tensor(tensor)
620 else:
621 return tensor
vllm/v1/attention/backend.py
+3
@@ -53,6 +53,9 @@ class AttentionBackend(ABC):
53 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
54 supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = ["auto", "bfloat16"]
55
56 + # Does attention's forward() include kv cache update?
57 + forward_includes_kv_cache_update: bool = True
58 +
59 @staticmethod
60 def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
61 return [MultipleOf(1)]
vllm/v1/attention/backends/flash_attn.py
+45 -26
@@ -79,6 +79,8 @@ class FlashAttentionBackend(AttentionBackend):
79 return [16, 32, 64]
80 return [MultipleOf(16)]
81
82 + forward_includes_kv_cache_update: bool = False
83 +
84 @staticmethod
85 def get_name() -> str:
86 return "FLASH_ATTN"
@@ -652,32 +654,6 @@ class FlashAttentionImpl(AttentionImpl):
654 # For decoder and cross-attention, use KV cache as before
655 key_cache, value_cache = kv_cache.unbind(0)
656
655 - # key and value may be None in the case of cross attention. They are
656 - # calculated once based on the output from the encoder and then cached
657 - # in KV cache.
658 - if (
659 - self.kv_sharing_target_layer_name is None
660 - and key is not None
661 - and value is not None
662 - ):
663 - # Reshape the input keys and values and store them in the cache.
664 - # Skip this if sharing KV cache with an earlier attention layer.
665 - # NOTE(woosuk): Here, key and value are padded while slot_mapping is
666 - # not padded. However, we don't need to do key[:num_actual_tokens]
667 - # and value[:num_actual_tokens] because the reshape_and_cache_flash
668 - # op uses the slot_mapping's shape to determine the number of
669 - # actual tokens.
670 - reshape_and_cache_flash(
671 - key,
672 - value,
673 - key_cache,
674 - value_cache,
675 - attn_metadata.slot_mapping,
676 - self.kv_cache_dtype,
677 - layer._k_scale,
678 - layer._v_scale,
679 - )
680 -
657 if self.kv_cache_dtype.startswith("fp8"):
658 # queries are quantized in the attention layer
659 dtype = FlashAttentionBackend.get_fp8_dtype_for_flashattn(
@@ -774,6 +750,49 @@ class FlashAttentionImpl(AttentionImpl):
750 )
751 return output
752
753 + def do_kv_cache_update(
754 + self,
755 + layer: torch.nn.Module,
756 + key: torch.Tensor,
757 + value: torch.Tensor,
758 + kv_cache: torch.Tensor,
759 + slot_mapping: torch.Tensor,
760 + ) -> None:
761 + if self.attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
762 + # For encoder attention,
763 + # we use direct Q, K, V tensors without caching
764 + return
765 +
766 + # key and value may be None in the case of cross attention. They are
767 + # calculated once based on the output from the encoder and then cached
768 + # in KV cache.
769 + if (
770 + self.kv_sharing_target_layer_name is not None
771 + or key is None
772 + or value is None
773 + ):
774 + return
775 +
776 + key_cache, value_cache = kv_cache.unbind(0)
777 +
778 + # Reshape the input keys and values and store them in the cache.
779 + # Skip this if sharing KV cache with an earlier attention layer.
780 + # NOTE(woosuk): Here, key and value are padded while slot_mapping is
781 + # not padded. However, we don't need to do key[:num_actual_tokens]
782 + # and value[:num_actual_tokens] because the reshape_and_cache_flash
783 + # op uses the slot_mapping's shape to determine the number of
784 + # actual tokens.
785 + reshape_and_cache_flash(
786 + key,
787 + value,
788 + key_cache,
789 + value_cache,
790 + slot_mapping,
791 + self.kv_cache_dtype,
792 + layer._k_scale,
793 + layer._v_scale,
794 + )
795 +
796 def _forward_with_dcp(
797 self,
798 query: torch.Tensor,
vllm/v1/spec_decode/eagle.py
+50
@@ -159,6 +159,10 @@ class SpecDecodeBaseProposer:
159 with_numpy=True,
160 )
161
162 + self._slot_mapping_buffer = torch.zeros(
163 + self.max_num_tokens, dtype=torch.int64, device=device
164 + )
165 +
166 # Determine allowed attention backends once during initialization.
167 self.allowed_attn_types: tuple | None = None
168 if current_platform.is_rocm():
@@ -230,6 +234,24 @@ class SpecDecodeBaseProposer:
234 positions = positions[0]
235 self.positions[:num_tokens] = positions
236
237 + def _get_slot_mapping(
238 + self,
239 + num_tokens: int,
240 + slot_mapping: torch.Tensor | None = None,
241 + ) -> dict[str, torch.Tensor]:
242 + """Return slot_mapping dict for EAGLE layers.
243 +
244 + If slot_mapping is provided, copies it into the buffer first.
245 + """
246 + if slot_mapping is not None:
247 + num_actual = slot_mapping.shape[0]
248 + self._slot_mapping_buffer[:num_actual].copy_(slot_mapping)
249 + if num_tokens > num_actual:
250 + self._slot_mapping_buffer[num_actual:num_tokens].fill_(PADDING_SLOT_ID)
251 +
252 + view = self._slot_mapping_buffer[:num_tokens]
253 + return {name: view for name in self.attn_layer_names + self.indexer_layer_names}
254 +
255 def initialize_cudagraph_keys(self, cudagraph_mode: CUDAGraphMode) -> None:
256 """Initialize cudagraph dispatcher keys for eagle.
257
@@ -262,6 +284,9 @@ class SpecDecodeBaseProposer:
284 sampling_metadata: SamplingMetadata,
285 mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
286 num_rejected_tokens_gpu: torch.Tensor | None = None,
287 + slot_mappings: dict[str, torch.Tensor]
288 + | list[dict[str, torch.Tensor]]
289 + | None = None,
290 ) -> torch.Tensor:
291 batch_size = common_attn_metadata.batch_size()
292
@@ -358,6 +383,9 @@ class SpecDecodeBaseProposer:
383 num_tokens=num_input_tokens,
384 num_tokens_across_dp=num_tokens_across_dp,
385 cudagraph_runtime_mode=cudagraph_runtime_mode,
386 + slot_mapping=self._get_slot_mapping(
387 + num_input_tokens, common_attn_metadata.slot_mapping
388 + ),
389 ):
390 ret_hidden_states = self.model(**model_kwargs)
391 if not self.model_returns_tuple():
@@ -396,6 +424,7 @@ class SpecDecodeBaseProposer:
424 positions=positions,
425 hidden_states=hidden_states,
426 common_attn_metadata=common_attn_metadata,
427 + slot_mappings=slot_mappings,
428 )
429 # [batch_size, num_tree_tokens]
430 return torch.cat(draft_token_ids_list, dim=1)
@@ -553,6 +582,9 @@ class SpecDecodeBaseProposer:
582 num_tokens=input_batch_size,
583 num_tokens_across_dp=batch_size_across_dp,
584 cudagraph_runtime_mode=cudagraph_runtime_mode,
585 + slot_mapping=self._get_slot_mapping(
586 + input_batch_size, common_attn_metadata.slot_mapping
587 + ),
588 ):
589 ret_hidden_states = self.model(**model_kwargs)
590 if not self.model_returns_tuple():
@@ -760,6 +792,9 @@ class SpecDecodeBaseProposer:
792 # [num_tokens, hidden_size]
793 hidden_states: torch.Tensor,
794 common_attn_metadata: CommonAttentionMetadata,
795 + slot_mappings: dict[str, torch.Tensor]
796 + | list[dict[str, torch.Tensor]]
797 + | None = None,
798 ) -> list[torch.Tensor]:
799 tree_attn_metadata_builder = self.runner.attn_groups[0][
800 0
@@ -881,6 +916,9 @@ class SpecDecodeBaseProposer:
916 self.vllm_config,
917 num_tokens=num_input_tokens,
918 cudagraph_runtime_mode=cudagraph_runtime_mode,
919 + slot_mapping=self._get_slot_mapping(
920 + num_input_tokens, attn_metadata.slot_mapping
921 + ),
922 ):
923 last_hidden_states, hidden_states = self.model(
924 input_ids=self.input_ids[:num_input_tokens],
@@ -1212,6 +1250,7 @@ class SpecDecodeBaseProposer:
1250 num_tokens: int,
1251 use_cudagraphs: bool = True,
1252 is_graph_capturing: bool = False,
1253 + slot_mappings: dict[str, torch.Tensor] | None = None,
1254 ) -> None:
1255 # FIXME: when using tree-based specdec, adjust number of forward-passes
1256 # according to the depth of the tree.
@@ -1233,12 +1272,23 @@ class SpecDecodeBaseProposer:
1272 if num_tokens_across_dp is not None:
1273 num_tokens_across_dp[self.dp_rank] = num_input_tokens
1274
1275 + # Make sure to use EAGLE's own buffer during cudagraph capture.
1276 + if (
1277 + self.attn_layer_names
1278 + and slot_mappings is not None
1279 + and self.attn_layer_names[0] in slot_mappings
1280 + ):
1281 + slot_mapping_dict = self._get_slot_mapping(num_input_tokens)
1282 + else:
1283 + slot_mapping_dict = slot_mappings or {}
1284 +
1285 with set_forward_context(
1286 None,
1287 self.vllm_config,
1288 num_tokens=num_input_tokens,
1289 num_tokens_across_dp=num_tokens_across_dp,
1290 cudagraph_runtime_mode=cudagraph_runtime_mode,
1291 + slot_mapping=slot_mapping_dict,
1292 ):
1293 if self.supports_mm_inputs:
1294 input_ids = None
vllm/v1/spec_decode/medusa.py
+3
@@ -38,6 +38,9 @@ class MedusaProposer:
38 self,
39 target_hidden_states: torch.Tensor,
40 sampling_metadata: SamplingMetadata,
41 + slot_mappings: dict[str, torch.Tensor]
42 + | list[dict[str, torch.Tensor]]
43 + | None = None, # unused
44 ) -> torch.Tensor:
45 # Generate blocks and compute logits
46 blocks = self.model(target_hidden_states)
vllm/v1/spec_decode/ngram_proposer.py
+4
@@ -3,6 +3,7 @@
3 import os
4
5 import numpy as np
6 +import torch
7 from numba import get_num_threads, jit, njit, prange, set_num_threads
8
9 from vllm.config import VllmConfig
@@ -132,6 +133,9 @@ class NgramProposer:
133 sampled_token_ids: list[list[int]],
134 num_tokens_no_spec: np.ndarray,
135 token_ids_cpu: np.ndarray,
136 + slot_mappings: dict[str, torch.Tensor]
137 + | list[dict[str, torch.Tensor]]
138 + | None = None, # unused
139 ) -> list[list[int]]:
140 # find which requests need ngram proposals
141 valid_ngram_requests = []
vllm/v1/spec_decode/suffix_decoding.py
+5
@@ -1,5 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 from vllm.config import VllmConfig
6 from vllm.v1.worker.gpu_input_batch import InputBatch
7
@@ -33,6 +35,9 @@ class SuffixDecodingProposer:
35 self,
36 input_batch: InputBatch,
37 sampled_token_ids: list[list[int]],
38 + slot_mappings: dict[str, torch.Tensor]
39 + | list[dict[str, torch.Tensor]]
40 + | None = None, # unused
41 ) -> list[list[int]]:
42 """
43 Propose speculative tokens for each request in the input batch. Suffix Decoding
vllm/v1/worker/gpu/attn_utils.py
+12
@@ -140,6 +140,18 @@ def init_kv_cache(
140 return kv_caches
141
142
143 +def build_slot_mappings_by_layer(
144 + slot_mappings: torch.Tensor,
145 + kv_cache_config: KVCacheConfig,
146 +) -> dict[str, torch.Tensor]:
147 + slot_mappings_by_layer: dict[str, torch.Tensor] = {}
148 + for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
149 + slot_mapping = slot_mappings[i]
150 + for layer_name in kv_cache_group.layer_names:
151 + slot_mappings_by_layer[layer_name] = slot_mapping
152 + return slot_mappings_by_layer
153 +
154 +
155 def build_attn_metadata(
156 attn_metadata_builders: list[AttentionMetadataBuilder],
157 num_reqs: int,
vllm/v1/worker/gpu/cudagraph_utils.py
+12 -4
@@ -14,7 +14,10 @@ from vllm.distributed.parallel_state import graph_capture, is_global_first_rank
14 from vllm.forward_context import set_forward_context
15 from vllm.v1.attention.backend import AttentionMetadataBuilder
16 from vllm.v1.kv_cache_interface import KVCacheConfig
17 -from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
17 +from vllm.v1.worker.gpu.attn_utils import (
18 + build_attn_metadata,
19 + build_slot_mappings_by_layer,
20 +)
21 from vllm.v1.worker.gpu.block_table import BlockTables
22 from vllm.v1.worker.gpu.dp_utils import make_num_tokens_across_dp
23 from vllm.v1.worker.gpu.input_batch import InputBuffers
@@ -88,7 +91,7 @@ class CudaGraphManager:
91 positions = mrope_positions[:, :num_tokens]
92 if inputs_embeds is not None:
93 inputs_embeds = inputs_embeds[:num_tokens]
91 - attn_metadata = prepare_inputs_to_capture(
94 + attn_metadata, slot_mappings = prepare_inputs_to_capture(
95 num_reqs,
96 num_tokens,
97 input_buffers,
@@ -98,6 +101,9 @@ class CudaGraphManager:
101 kv_cache_config,
102 )
103 num_tokens_across_dp = make_num_tokens_across_dp(self.dp_size, num_tokens)
104 + slot_mappings_by_layer = build_slot_mappings_by_layer(
105 + slot_mappings, kv_cache_config
106 + )
107
108 # Warm up.
109 with set_forward_context(
@@ -106,6 +112,7 @@ class CudaGraphManager:
112 num_tokens=num_tokens,
113 cudagraph_runtime_mode=CUDAGraphMode.NONE,
114 num_tokens_across_dp=num_tokens_across_dp,
115 + slot_mapping=slot_mappings_by_layer,
116 ):
117 hidden_states = model(
118 input_ids=input_ids,
@@ -125,6 +132,7 @@ class CudaGraphManager:
132 num_tokens=num_tokens,
133 cudagraph_runtime_mode=CUDAGraphMode.NONE,
134 num_tokens_across_dp=num_tokens_across_dp,
135 + slot_mapping=slot_mappings_by_layer,
136 ),
137 torch.cuda.graph(graph, self.pool),
138 ):
@@ -244,7 +252,7 @@ def prepare_inputs_to_capture(
252 attn_metadata_builders: list[AttentionMetadataBuilder],
253 max_model_len: int,
254 kv_cache_config: KVCacheConfig,
247 -) -> dict[str, Any]:
255 +) -> tuple[dict[str, Any], torch.Tensor]:
256 num_tokens_per_req = num_tokens // num_reqs
257
258 query_start_loc_np = np.arange(num_reqs + 1, dtype=np.int32) * num_tokens_per_req
@@ -274,4 +282,4 @@ def prepare_inputs_to_capture(
282 slot_mappings=slot_mappings,
283 kv_cache_config=kv_cache_config,
284 )
277 - return attn_metadata
285 + return attn_metadata, slot_mappings
vllm/v1/worker/gpu/model_runner.py
+10
@@ -24,6 +24,7 @@ from vllm.v1.outputs import ModelRunnerOutput
24 from vllm.v1.worker.gpu.async_utils import AsyncOutput
25 from vllm.v1.worker.gpu.attn_utils import (
26 build_attn_metadata,
27 + build_slot_mappings_by_layer,
28 get_kv_cache_spec,
29 init_attn_backend,
30 init_kv_cache,
@@ -881,6 +882,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
882 if self.uses_mrope:
883 assert input_batch.mrope_positions is not None
884 positions = input_batch.mrope_positions
885 + slot_mappings = self.block_tables.compute_slot_mappings(
886 + input_batch.idx_mapping,
887 + input_batch.query_start_loc,
888 + input_batch.positions[: input_batch.num_tokens],
889 + )
890 + slot_mappings_by_layer = build_slot_mappings_by_layer(
891 + slot_mappings, self.kv_cache_config
892 + )
893 with set_forward_context(
894 input_batch.attn_metadata,
895 self.vllm_config,
@@ -888,6 +897,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
897 # TODO(woosuk): Support piecewise CUDA graph.
898 cudagraph_runtime_mode=CUDAGraphMode.NONE,
899 num_tokens_across_dp=num_tokens_across_dp,
900 + slot_mapping=slot_mappings_by_layer,
901 ):
902 self.kv_connector.pre_forward(scheduler_output)
903 hidden_states = self.model(
vllm/v1/worker/gpu_model_runner.py
+127 -14
@@ -314,6 +314,7 @@ class ExecuteModelState(NamedTuple):
314 aux_hidden_states: list[torch.Tensor] | None
315 ec_connector_output: ECConnectorOutput | None
316 cudagraph_stats: CUDAGraphStat | None
317 + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None
318
319
320 class GPUModelRunner(
@@ -1595,6 +1596,7 @@ class GPUModelRunner(
1596 for_cudagraph_capture: bool = False,
1597 num_scheduled_tokens: dict[str, int] | None = None,
1598 cascade_attn_prefix_lens: list[list[int]] | None = None,
1599 + slot_mappings: dict[int, torch.Tensor] | None = None,
1600 ) -> tuple[PerLayerAttnMetadata, CommonAttentionMetadata | None]:
1601 """
1602 :return: tuple[attn_metadata, spec_decode_common_attn_metadata]
@@ -1628,7 +1630,7 @@ class GPUModelRunner(
1630
1631 kv_cache_groups = self.kv_cache_config.kv_cache_groups
1632
1631 - def _get_block_table_and_slot_mapping(kv_cache_gid: int):
1633 + def _get_block_table(kv_cache_gid: int):
1634 assert num_reqs_padded is not None and num_tokens_padded is not None
1635 kv_cache_spec = kv_cache_groups[kv_cache_gid].kv_cache_spec
1636 if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
@@ -1637,24 +1639,19 @@ class GPUModelRunner(
1639 dtype=torch.int32,
1640 device=self.device,
1641 )
1640 - slot_mapping = torch.zeros(
1641 - (num_tokens_padded,),
1642 - dtype=torch.int64,
1643 - device=self.device,
1644 - )
1642 else:
1643 blk_table = self.input_batch.block_table[kv_cache_gid]
1644 blk_table_tensor = blk_table.get_device_tensor(num_reqs_padded)
1648 - slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded]
1645
1646 # Fill unused with -1. Needed for reshape_and_cache in full cuda
1647 # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID
1652 - slot_mapping[num_tokens:num_tokens_padded].fill_(-1)
1648 blk_table_tensor[num_reqs:num_reqs_padded].fill_(-1)
1649 + return blk_table_tensor
1650
1655 - return blk_table_tensor, slot_mapping
1651 + assert slot_mappings is not None
1652 + block_table_gid_0 = _get_block_table(0)
1653 + slot_mapping_gid_0 = slot_mappings[0]
1654
1657 - block_table_gid_0, slot_mapping_gid_0 = _get_block_table_and_slot_mapping(0)
1655 if self.model_config.enable_return_routed_experts:
1656 self.slot_mapping = slot_mapping_gid_0[:num_tokens].cpu().numpy()
1657 cm_base = CommonAttentionMetadata(
@@ -1779,9 +1776,8 @@ class GPUModelRunner(
1776 num_reqs_padded,
1777 )
1778 if kv_cache_gid > 0:
1782 - cm.block_table_tensor, cm.slot_mapping = (
1783 - _get_block_table_and_slot_mapping(kv_cache_gid)
1784 - )
1779 + cm.block_table_tensor = _get_block_table(kv_cache_gid)
1780 + cm.slot_mapping = slot_mappings[kv_cache_gid]
1781
1782 if self.speculative_config and spec_decode_common_attn_metadata is None:
1783 if isinstance(self.drafter, EagleProposer):
@@ -3119,6 +3115,80 @@ class GPUModelRunner(
3115 pyt_hooks.register_hooks(self.model, self.model.__class__.__name__)
3116 self.layerwise_nvtx_hooks_registered = True
3117
3118 + def _get_slot_mappings(
3119 + self,
3120 + num_tokens_padded: int,
3121 + num_reqs_padded: int,
3122 + num_tokens_unpadded: int,
3123 + ubatch_slices: "UBatchSlices | None" = None,
3124 + ) -> tuple[
3125 + dict[int, torch.Tensor] | None,
3126 + dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
3127 + ]:
3128 + """
3129 + Build slot mappings in both formats needed by the system.
3130 +
3131 + Args:
3132 + num_tokens_padded: Total number of tokens (padded)
3133 + num_reqs_padded: Total number of requests (padded)
3134 + num_tokens_unpadded: Actual number of tokens (unpadded)
3135 + ubatch_slices: Optional ubatch slicing info for DBO
3136 +
3137 + Returns:
3138 + A tuple of:
3139 + - slot_mappings_by_gid: dict[int, torch.Tensor] for attention metadata
3140 + - slot_mappings_by_layer: dict[str, torch.Tensor] or list for ForwardContext
3141 + """
3142 + if not (
3143 + hasattr(self, "kv_cache_config")
3144 + and self.kv_cache_config is not None
3145 + and len(self.kv_cache_config.kv_cache_groups) > 0
3146 + ):
3147 + return None, None
3148 +
3149 + def _get_slot_mapping(kv_cache_gid: int):
3150 + assert num_reqs_padded is not None and num_tokens_padded is not None
3151 + kv_cache_spec = self.kv_cache_config.kv_cache_groups[
3152 + kv_cache_gid
3153 + ].kv_cache_spec
3154 + if isinstance(kv_cache_spec, EncoderOnlyAttentionSpec):
3155 + slot_mapping = torch.zeros(
3156 + (num_tokens_padded,),
3157 + dtype=torch.int64,
3158 + device=self.device,
3159 + )
3160 + else:
3161 + blk_table = self.input_batch.block_table[kv_cache_gid]
3162 + slot_mapping = blk_table.slot_mapping.gpu[:num_tokens_padded]
3163 +
3164 + # Fill unused with -1. Needed for reshape_and_cache in full cuda
3165 + # graph mode. `blk_table_tensor` -1 to match mamba PAD_SLOT_ID
3166 + slot_mapping[num_tokens_unpadded:num_tokens_padded].fill_(-1)
3167 +
3168 + return slot_mapping
3169 +
3170 + slot_mappings_by_gid = {
3171 + gid: _get_slot_mapping(gid)
3172 + for gid, _ in enumerate(self.kv_cache_config.kv_cache_groups)
3173 + }
3174 +
3175 + slot_mappings_by_layer: dict[str, torch.Tensor] = {}
3176 + for gid, kv_cache_group in enumerate(self.kv_cache_config.kv_cache_groups):
3177 + slot_mapping = slot_mappings_by_gid[gid]
3178 + for layer_name in kv_cache_group.layer_names:
3179 + slot_mappings_by_layer[layer_name] = slot_mapping
3180 +
3181 + if ubatch_slices is not None:
3182 + result: list[dict[str, torch.Tensor]] = []
3183 + for ubatch in ubatch_slices:
3184 + sliced_mappings: dict[str, torch.Tensor] = {}
3185 + for layer_name, slot_mapping in slot_mappings_by_layer.items():
3186 + sliced_mappings[layer_name] = slot_mapping[ubatch.token_slice]
3187 + result.append(sliced_mappings)
3188 + return slot_mappings_by_gid, result
3189 +
3190 + return slot_mappings_by_gid, slot_mappings_by_layer
3191 +
3192 @torch.inference_mode()
3193 def execute_model(
3194 self,
@@ -3248,6 +3318,17 @@ class GPUModelRunner(
3318 ubatch_slices_padded,
3319 )
3320
3321 + # True if any attention backend handles KV cache update separately
3322 + # from forward() (i.e., forward_includes_kv_cache_update=False). When true,
3323 + # slot_mappings must use padded dimensions to match the key/value tensors.
3324 + has_separate_kv_update = not all(
3325 + all(
3326 + g.backend.forward_includes_kv_cache_update
3327 + for g in self.attn_groups[id]
3328 + )
3329 + for id, spec in enumerate(self.kv_cache_config.kv_cache_groups)
3330 + if not isinstance(spec.kv_cache_spec, EncoderOnlyAttentionSpec)
3331 + )
3332 pad_attn = cudagraph_mode == CUDAGraphMode.FULL
3333
3334 if self.cache_config.mamba_cache_mode == "align":
@@ -3265,6 +3346,17 @@ class GPUModelRunner(
3346 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
3347 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
3348
3349 + slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
3350 + num_tokens_padded=num_tokens_padded
3351 + if pad_attn or has_separate_kv_update
3352 + else num_tokens_unpadded,
3353 + num_reqs_padded=(
3354 + num_reqs_padded if pad_attn or has_separate_kv_update else num_reqs
3355 + ),
3356 + num_tokens_unpadded=num_tokens_unpadded,
3357 + ubatch_slices=ubatch_slices_padded,
3358 + )
3359 +
3360 attn_metadata, spec_decode_common_attn_metadata = (
3361 self._build_attention_metadata(
3362 num_tokens=num_tokens_unpadded,
@@ -3277,6 +3369,7 @@ class GPUModelRunner(
3369 use_spec_decode=use_spec_decode,
3370 num_scheduled_tokens=scheduler_output.num_scheduled_tokens,
3371 cascade_attn_prefix_lens=cascade_attn_prefix_lens,
3372 + slot_mappings=slot_mappings_by_group,
3373 )
3374 )
3375
@@ -3317,6 +3410,7 @@ class GPUModelRunner(
3410 cudagraph_runtime_mode=cudagraph_mode,
3411 batch_descriptor=batch_desc,
3412 ubatch_slices=ubatch_slices_padded,
3413 + slot_mapping=slot_mappings,
3414 skip_compiled=has_encoder_input,
3415 ),
3416 record_function_or_nullcontext("gpu_model_runner: forward"),
@@ -3399,6 +3493,7 @@ class GPUModelRunner(
3493 aux_hidden_states,
3494 ec_connector_output,
3495 cudagraph_stats,
3496 + slot_mappings,
3497 )
3498 self.kv_connector_output = kv_connector_output
3499 return None
@@ -3435,6 +3530,7 @@ class GPUModelRunner(
3530 aux_hidden_states,
3531 ec_connector_output,
3532 cudagraph_stats,
3533 + slot_mappings,
3534 ) = self.execute_model_state
3535 # Clear ephemeral state.
3536 self.execute_model_state = None
@@ -3468,6 +3564,7 @@ class GPUModelRunner(
3564 aux_hidden_states,
3565 spec_decode_metadata,
3566 spec_decode_common_attn_metadata,
3567 + slot_mappings,
3568 )
3569 self._copy_draft_token_ids_to_cpu(scheduler_output)
3570
@@ -3676,6 +3773,7 @@ class GPUModelRunner(
3773 aux_hidden_states: list[torch.Tensor] | None,
3774 spec_decode_metadata: SpecDecodeMetadata | None,
3775 common_attn_metadata: CommonAttentionMetadata,
3776 + slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
3777 ) -> list[list[int]] | torch.Tensor:
3778 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
3779 spec_config = self.speculative_config
@@ -3687,11 +3785,14 @@ class GPUModelRunner(
3785 sampled_token_ids,
3786 self.input_batch.num_tokens_no_spec,
3787 self.input_batch.token_ids_cpu,
3788 + slot_mappings=slot_mappings,
3789 )
3790 elif spec_config.method == "suffix":
3791 assert isinstance(sampled_token_ids, list)
3792 assert isinstance(self.drafter, SuffixDecodingProposer)
3694 - draft_token_ids = self.drafter.propose(self.input_batch, sampled_token_ids)
3793 + draft_token_ids = self.drafter.propose(
3794 + self.input_batch, sampled_token_ids, slot_mappings=slot_mappings
3795 + )
3796 elif spec_config.method == "medusa":
3797 assert isinstance(sampled_token_ids, list)
3798 assert isinstance(self.drafter, MedusaProposer)
@@ -3716,6 +3817,7 @@ class GPUModelRunner(
3817 draft_token_ids = self.drafter.propose(
3818 target_hidden_states=hidden_states,
3819 sampling_metadata=sampling_metadata,
3820 + slot_mappings=slot_mappings,
3821 )
3822 elif spec_config.use_eagle() or spec_config.uses_draft_model():
3823 assert isinstance(self.drafter, EagleProposer | DraftModelProposer)
@@ -3826,6 +3928,7 @@ class GPUModelRunner(
3928 common_attn_metadata=common_attn_metadata,
3929 mm_embed_inputs=mm_embed_inputs,
3930 num_rejected_tokens_gpu=num_rejected_tokens_gpu,
3931 + slot_mappings=slot_mappings,
3932 )
3933
3934 return draft_token_ids
@@ -4406,6 +4509,13 @@ class GPUModelRunner(
4509
4510 attn_metadata: PerLayerAttnMetadata | None = None
4511
4512 + slot_mappings_by_group, slot_mappings = self._get_slot_mappings(
4513 + num_tokens_padded=num_tokens,
4514 + num_reqs_padded=num_reqs_padded,
4515 + num_tokens_unpadded=num_tokens_unpadded,
4516 + ubatch_slices=ubatch_slices_padded,
4517 + )
4518 +
4519 # If force_attention is True, we always capture attention. Otherwise,
4520 # it only happens for cudagraph_runtime_mode=FULL.
4521 if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL:
@@ -4431,6 +4541,7 @@ class GPUModelRunner(
4541 max_query_len=max_query_len,
4542 ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices,
4543 for_cudagraph_capture=is_graph_capturing,
4544 + slot_mappings=slot_mappings_by_group,
4545 )
4546
4547 with self.maybe_dummy_run_with_lora(
@@ -4499,6 +4610,7 @@ class GPUModelRunner(
4610 cudagraph_runtime_mode=cudagraph_runtime_mode,
4611 batch_descriptor=batch_desc,
4612 ubatch_slices=ubatch_slices_padded,
4613 + slot_mapping=slot_mappings,
4614 ),
4615 ):
4616 outputs = self.model(
@@ -4545,6 +4657,7 @@ class GPUModelRunner(
4657 num_tokens,
4658 use_cudagraphs=use_cudagraphs,
4659 is_graph_capturing=is_graph_capturing,
4660 + slot_mappings=slot_mappings,
4661 )
4662
4663 # We register layerwise NVTX hooks here after the first dynamo tracing is
vllm/v1/worker/gpu_ubatch_wrapper.py
+8
@@ -295,6 +295,7 @@ class UBatchWrapper:
295 self,
296 ubatch_slices,
297 attn_metadata,
298 + slot_mapping,
299 input_ids,
300 positions,
301 inputs_embeds,
@@ -306,6 +307,9 @@ class UBatchWrapper:
307 ) -> list[UbatchMetadata]:
308 # Create one forward context per ubatch
309 forward_contexts = []
310 + # slot_mapping can be None, an empty dict (from create_forward_context
311 + # converting None to {}), or a list of dicts (one per ubatch)
312 + has_slot_mapping = slot_mapping and isinstance(slot_mapping, list)
313 for i, ubatch_slice in enumerate(ubatch_slices):
314 forward_contexts.append(
315 create_forward_context(
@@ -314,6 +318,7 @@ class UBatchWrapper:
318 dp_metadata=dp_metadata[i],
319 batch_descriptor=batch_descriptor,
320 cudagraph_runtime_mode=cudagraph_runtime_mode,
321 + slot_mapping=slot_mapping[i] if has_slot_mapping else None,
322 )
323 )
324
@@ -406,6 +411,7 @@ class UBatchWrapper:
411 return self.cudagraph_wrapper(*args, **kwargs)
412
413 attn_metadata = forward_context.attn_metadata
414 + slot_mapping = forward_context.slot_mapping
415 num_tokens = (
416 ubatch_slices[0].token_slice.stop - ubatch_slices[0].token_slice.start
417 ) * 2
@@ -440,6 +446,7 @@ class UBatchWrapper:
446 ubatch_metadata = self._make_ubatch_metadata(
447 ubatch_slices=ubatch_slices,
448 attn_metadata=attn_metadata,
449 + slot_mapping=slot_mapping,
450 input_ids=input_ids,
451 positions=positions,
452 intermediate_tensors=intermediate_tensors,
@@ -462,6 +469,7 @@ class UBatchWrapper:
469 ubatch_metadata = self._make_ubatch_metadata(
470 ubatch_slices=ubatch_slices,
471 attn_metadata=attn_metadata,
472 + slot_mapping=slot_mapping,
473 input_ids=input_ids,
474 positions=positions,
475 intermediate_tensors=intermediate_tensors,