diff --git a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py index 551811e60e8..e2d54821ce9 100644 --- a/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py +++ b/tests/distributed/test_eplb_fused_moe_layer_dep_nvfp4.py @@ -225,6 +225,12 @@ def _test_eplb_fml(env, world_size: int, test_config: TestConfig): logical_to_physical_map, logical_replica_count, ) + fml.router.eplb_state.should_record_tensor = torch.ones( + (), dtype=torch.bool, device=device + ) + fml.router.eplb_state.num_unpadded_tokens_tensors = [ + torch.tensor(0, dtype=torch.int32, device=device) + ] out_after_shuffle = [] with set_forward_context( diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index d1bcd3241aa..552063988fa 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -1332,6 +1332,9 @@ def _test_body_eplb( eplb_moe_layer.router.eplb_state.should_record_tensor = torch.ones( (), dtype=torch.bool, device=device ) + eplb_moe_layer.router.eplb_state.num_unpadded_tokens_tensors = [ + torch.tensor(0, dtype=torch.int32, device=device) + ] # Get "after" output with rearranged weights and EPLB routing with set_forward_context( diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index 41dea812193..62a4968a0d1 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -61,12 +61,14 @@ def setup_eplb_state( global_num_experts, dtype=torch.int64, device="cuda" ) should_record_tensor = torch.ones((), dtype=torch.bool, device="cuda") + num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32, device="cuda")] return EplbLayerState( expert_load_view=expert_load_view, logical_to_physical_map=logical_to_physical_map, logical_replica_count=logical_replica_count, should_record_tensor=should_record_tensor, + num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) @@ -782,3 +784,67 @@ def test_eplb_map_with_redundancy( torch.testing.assert_close(load, exp_load) else: assert load.sum().item() == 0 + + +@pytest.mark.parametrize( + "l2p_map, replica_count, num_physical, topk_ids, " + "num_unpadded, expected_out, expected_load", + [ + pytest.param( + [[0], [1], [2], [3]], + [1, 1, 1, 1], + 4, + [[0, 1], [2, 3], [0, 2], [1, 3]], + 2, + [[0, 1], [2, 3], [0, 2], [1, 3]], + # only rows 0,1 counted: expert 0→1, 1→1, 2→1, 3→1 + [1, 1, 1, 1], + id="half_padded", + ), + pytest.param( + # record everything (None = no padding info) + [[0], [1], [2], [3]], + [1, 1, 1, 1], + 4, + [[0, 1], [2, 3], [0, 2], [1, 3]], + None, + [[0, 1], [2, 3], [0, 2], [1, 3]], + [2, 2, 2, 2], + id="no_padding_info", + ), + ], +) +def test_eplb_map_num_unpadded_tokens( + l2p_map, + replica_count, + num_physical, + topk_ids, + num_unpadded, + expected_out, + expected_load, +): + l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda") + rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda") + load = torch.zeros(num_physical, dtype=torch.int32, device="cuda") + rec = torch.tensor(True, dtype=torch.bool, device="cuda") + ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda") + num_unpadded_t = ( + torch.tensor(num_unpadded, dtype=torch.int32, device="cuda") + if num_unpadded is not None + else None + ) + + out = eplb_map_to_physical_and_record( + topk_ids=ids, + expert_load_view=load, + logical_to_physical_map=l2p, + logical_replica_count=rc, + record_enabled=rec, + num_unpadded_tokens=num_unpadded_t, + ) + + exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda") + torch.testing.assert_close(out, exp_out) + + exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda") + torch.testing.assert_close(load, exp_load) diff --git a/tests/model_executor/test_routed_experts_capture.py b/tests/model_executor/test_routed_experts_capture.py index d1a542396e6..9efee9eec82 100644 --- a/tests/model_executor/test_routed_experts_capture.py +++ b/tests/model_executor/test_routed_experts_capture.py @@ -91,6 +91,7 @@ def test_base_router_capture_with_eplb_enabled(): eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1) eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64) eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool) + eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)] router = _make_router(eplb_state=eplb_state) captured = [] diff --git a/vllm/config/utils.py b/vllm/config/utils.py index 12e0385aeb1..3df0f7210f7 100644 --- a/vllm/config/utils.py +++ b/vllm/config/utils.py @@ -203,6 +203,26 @@ class SupportsHash(Protocol): def compute_hash(self) -> str: ... +_config_hash_cache: dict[int, str] = {} + + +def compute_hash_cached(config: SupportsHash) -> str: + """Cache config.compute_hash() by object identity. + + Config objects (ModelConfig, etc.) are long-lived singletons that never + mutate after construction, but compute_hash() is expensive (JSON + serialization + SHA-256). This utility avoids recomputing the hash on + every forward pass while keeping a single consistent key type for all + lookup paths. + """ + key = id(config) + result = _config_hash_cache.get(key) + if result is None: + result = config.compute_hash() + _config_hash_cache[key] = result + return result + + class SupportsMetricsInfo(Protocol): def metrics_info(self) -> dict[str, str]: ... diff --git a/vllm/distributed/elastic_ep/elastic_execute.py b/vllm/distributed/elastic_ep/elastic_execute.py index 3cb0d603e3e..b0c3740f57e 100644 --- a/vllm/distributed/elastic_ep/elastic_execute.py +++ b/vllm/distributed/elastic_ep/elastic_execute.py @@ -458,7 +458,9 @@ class ElasticEPScalingExecutor: eplb_model_state.logical_to_physical_map, eplb_model_state.logical_replica_count, ) - eplb_state._init_should_record_tensor(model) + eplb_state._propagate_shared_tensors( + model, eplb_model_state.num_unpadded_tokens_tensors + ) model.update_physical_experts_metadata( num_physical_experts=num_physical_experts, num_local_physical_experts=num_local_experts, diff --git a/vllm/distributed/eplb/eplb_state.py b/vllm/distributed/eplb/eplb_state.py index 74f357fbdbf..feacb03d28b 100644 --- a/vllm/distributed/eplb/eplb_state.py +++ b/vllm/distributed/eplb/eplb_state.py @@ -35,6 +35,7 @@ import torch from torch.distributed import ProcessGroup, all_reduce from vllm.config import ModelConfig, ParallelConfig +from vllm.config.utils import compute_hash_cached from vllm.distributed.parallel_state import ( get_ep_group, get_eplb_group, @@ -206,6 +207,13 @@ class EplbModelState: pending_result relies on the GIL to synchronize access between the main thread and the async worker. """ + num_unpadded_tokens_tensors: list[torch.Tensor] | None = None + """ + Per-ubatch scalar int32 tensors holding the number of real (non-padding) + tokens. Allocated once in :meth:`EplbState.add_model` so that device + pointers remain stable across CUDA-graph replays. The router kernel + indexes this list with ``dbo_current_ubatch_id()``. + """ class EplbState: @@ -253,7 +261,7 @@ class EplbState: Shared scalar bool tensor for all layers. Every :class:`EplbLayerState` holds a reference to the **same** object so a single ``.fill_()`` updates all layers at once. Allocated on the - first call to :meth:`_init_should_record_tensor`. + first call to :meth:`_propagate_shared_tensors`. """ self.is_async: bool = False """ @@ -440,12 +448,19 @@ class EplbState: self.policy = EPLB_POLICIES[policy_type] logger.debug("Selected EPLB policy: %s", policy_type) + # num_ubatches is 0 when DBO is disabled. + num_ubatches = max(1, self.parallel_config.num_ubatches) + num_unpadded_tokens_tensors = [ + torch.tensor(0, dtype=torch.int32, device=self.device) + for _ in range(num_ubatches) + ] + model.set_eplb_state( expert_load_pass, logical_to_physical_map, logical_replica_count, ) - self._init_should_record_tensor(model) + self._propagate_shared_tensors(model, num_unpadded_tokens_tensors) expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]] assert self.parallel_config.eplb_config.communicator is not None, ( @@ -471,10 +486,43 @@ class EplbState: eplb_stats=None, cuda_device_index=self.cuda_device_index, communicator=communicator, + num_unpadded_tokens_tensors=num_unpadded_tokens_tensors, ) self.model_states[model_config.compute_hash()] = model_state self.num_valid_physical_experts = model.num_physical_experts + def prepare_forward( + self, + model_config: ModelConfig, + num_unpadded_tokens: int, + ubatch_slices: list | None = None, + ) -> None: + """Fill the per-[u]batch ``num_unpadded_tokens`` tensors before a + forward pass. + + Args: + model_config: Identifies which ``EplbModelState`` to update. + num_unpadded_tokens: Total number of real (non-padding) tokens + in the batch. + ubatch_slices: When DBO is active, a list of + ``UBatchSlice`` objects describing each micro-batch's + token range. When ``None``, only ``tensors[0]`` is filled. + """ + model_state = self.model_states.get(compute_hash_cached(model_config)) + if model_state is None or model_state.num_unpadded_tokens_tensors is None: + return + tensors = model_state.num_unpadded_tokens_tensors + if ubatch_slices is None: + tensors[0].fill_(num_unpadded_tokens) + else: + for i, ubatch_slice in enumerate(ubatch_slices): + ts = ubatch_slice.token_slice + # Real tokens in this ubatch: clamp the global count into + # the slice range so partially-filled ubatches get the + # correct count. + val = max(0, min(num_unpadded_tokens, ts.stop) - ts.start) + tensors[i].fill_(val) + def step( self, is_dummy: bool = False, @@ -638,11 +686,20 @@ class EplbState: self._should_record_current_step(log_stats=log_stats) ) - def _init_should_record_tensor(self, model: "MixtureOfExperts") -> None: # type: ignore[name-defined] - """Allocate (once) and propagate the shared ``should_record_tensor``. + def _propagate_shared_tensors( + self, + model: "MixtureOfExperts", # type: ignore[name-defined] + num_unpadded_tokens_tensors: list[torch.Tensor], + ) -> None: + """Propagate shared tensors to every :class:`EplbLayerState`. + + Allocates ``should_record_tensor`` on the first call and then + assigns both it and ``num_unpadded_tokens_tensors`` to every + MoE layer's :class:`EplbLayerState`. All layers reference the + **same** objects so a single update is visible everywhere. Must be called after :meth:`model.set_eplb_state` so that each - layer's ``eplb_state`` is already populated with the tensor views. + layer's ``eplb_state`` is already populated. """ layer_states = [ layer.eplb_state @@ -659,6 +716,7 @@ class EplbState: for ls in layer_states: if ls is not None: ls.should_record_tensor = self.should_record_tensor + ls.num_unpadded_tokens_tensors = num_unpadded_tokens_tensors def rearrange( self, @@ -985,6 +1043,11 @@ class EplbLayerState: sliding window before the next rearrangement, so recording them wastes GPU work. """ + num_unpadded_tokens_tensors: list[torch.Tensor] | None = None + """ + Reference to the parent :class:`EplbModelState`'s tensor list so the + router can read the correct per-[u]batch unpadded token count. + """ def set_layer_state( self, diff --git a/vllm/model_executor/layers/fused_moe/router/base_router.py b/vllm/model_executor/layers/fused_moe/router/base_router.py index 4ba855b645f..01e674b2b13 100644 --- a/vllm/model_executor/layers/fused_moe/router/base_router.py +++ b/vllm/model_executor/layers/fused_moe/router/base_router.py @@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.v1.worker.ubatching import dbo_current_ubatch_id if current_platform.is_cuda_alike(): @@ -22,11 +23,13 @@ if current_platform.is_cuda_alike(): out_ids_ptr, out_ptr, record_enabled_ptr, + num_unpadded_tokens_ptr, num_logical_experts, map_slots, out_size, numel, num_active_experts, + HAS_NUM_UNPADDED: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) @@ -50,6 +53,13 @@ if current_platform.is_cuda_alike(): token_idx = (offs // num_active_experts).to(tl.int64) hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF replica_idx = hashed % replica_count + map_index = safe_expert_id * map_slots + replica_idx + physical_id = tl.load( + logical_to_physical_ptr + map_index, + mask=mask & valid_expert, + other=-1, + ) + tl.store(out_ids_ptr + offs, physical_id, mask=mask) # 2. Record expert load metrics. @@ -64,16 +74,21 @@ if current_platform.is_cuda_alike(): # If later refactor moved all the MoE kernel calls # to the modular kernel, we can move this logic there # to achieve better efficiency. - map_index = safe_expert_id * map_slots + replica_idx - physical_id = tl.load( - logical_to_physical_ptr + map_index, - mask=mask & valid_expert, - other=-1, - ) - tl.store(out_ids_ptr + offs, physical_id, mask=mask) record_enabled = tl.load(record_enabled_ptr) != 0 - valid = mask & record_enabled & (physical_id >= 0) & (physical_id < out_size) + # Skip padded tokens when recording. + if HAS_NUM_UNPADDED: + num_unpadded_tokens = tl.load(num_unpadded_tokens_ptr) + is_unpadded = offs < num_unpadded_tokens * num_active_experts + else: + is_unpadded = True + valid = ( + mask + & record_enabled + & is_unpadded + & (physical_id >= 0) + & (physical_id < out_size) + ) safe_physical_id = tl.where(physical_id >= 0, physical_id, 0) tl.atomic_add(out_ptr + safe_physical_id, 1, mask=valid) @@ -83,6 +98,7 @@ if current_platform.is_cuda_alike(): logical_replica_count: torch.Tensor, expert_load_view: torch.Tensor, record_enabled: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None, ) -> torch.Tensor: topk_ids_in = topk_ids.contiguous().to(dtype=torch.int32) numel = topk_ids_in.numel() @@ -99,11 +115,13 @@ if current_platform.is_cuda_alike(): out_flat, expert_load_view, record_enabled, + num_unpadded_tokens, logical_replica_count.shape[0], logical_to_physical_map.shape[1], expert_load_view.shape[0], numel, num_active_experts, + HAS_NUM_UNPADDED=num_unpadded_tokens is not None, BLOCK_SIZE=256, ) return out_flat.reshape(topk_ids.shape) @@ -114,6 +132,7 @@ if current_platform.is_cuda_alike(): logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, record_enabled: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None = None, ) -> torch.Tensor: # Fused triton implementation: mapping + optional recording in one kernel. return _eplb_map_and_record_triton( @@ -122,6 +141,7 @@ if current_platform.is_cuda_alike(): logical_replica_count=logical_replica_count, expert_load_view=expert_load_view, record_enabled=record_enabled, + num_unpadded_tokens=num_unpadded_tokens, ) else: @@ -131,6 +151,7 @@ else: logical_to_physical_map: torch.Tensor, logical_replica_count: torch.Tensor, record_enabled: torch.Tensor, + num_unpadded_tokens: torch.Tensor | None = None, ) -> torch.Tensor: return topk_ids @@ -177,6 +198,8 @@ class BaseRouter(FusedMoERouter): raise ValueError("EPLB requires logical_replica_count != None") if eplb_state.should_record_tensor is None: raise ValueError("EPLB requires should_record_tensor != None") + if eplb_state.num_unpadded_tokens_tensors is None: + raise ValueError("EPLB requires num_unpadded_tokens_tensors != None") def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor: """Apply EPLB mapping to convert logical expert IDs to physical expert IDs.""" @@ -186,12 +209,16 @@ class BaseRouter(FusedMoERouter): assert eplb_state.logical_to_physical_map is not None assert eplb_state.logical_replica_count is not None assert eplb_state.should_record_tensor is not None + assert eplb_state.num_unpadded_tokens_tensors is not None return eplb_map_to_physical_and_record( topk_ids=topk_ids, logical_to_physical_map=eplb_state.logical_to_physical_map, logical_replica_count=eplb_state.logical_replica_count, expert_load_view=eplb_state.expert_load_view, record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ], ) return topk_ids diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 99373361922..f1bcd534e97 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -70,6 +70,7 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv from vllm.v1.attention.backends.registry import AttentionBackendEnum +from vllm.v1.worker.ubatching import dbo_current_ubatch_id class DeepseekV4MLP(nn.Module): @@ -464,6 +465,11 @@ class DeepseekV4MegaMoEExperts(nn.Module): logical_to_physical_map=eplb_state.logical_to_physical_map, logical_replica_count=eplb_state.logical_replica_count, record_enabled=eplb_state.should_record_tensor, + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[ + dbo_current_ubatch_id() + ] + if eplb_state.num_unpadded_tokens_tensors is not None + else None, ) prepare_megamoe_inputs( diff --git a/vllm/v1/spec_decode/extract_hidden_states.py b/vllm/v1/spec_decode/extract_hidden_states.py index b6f9eac4dfa..de7a075e2f7 100644 --- a/vllm/v1/spec_decode/extract_hidden_states.py +++ b/vllm/v1/spec_decode/extract_hidden_states.py @@ -9,6 +9,7 @@ import torch import torch.nn as nn from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config +from vllm.distributed.eplb.eplb_state import EplbState from vllm.forward_context import set_forward_context from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.model_loader import get_model @@ -43,6 +44,8 @@ class ExtractHiddenStatesProposer: self.dtype = vllm_config.model_config.dtype self.dp_rank = vllm_config.parallel_config.data_parallel_rank + self.eplb_state: EplbState | None = None + # Model and attention layer tracking (initialized in load_model) self.model: nn.Module | None = None self.attn_layer_names: list[str] = [] @@ -83,6 +86,10 @@ class ExtractHiddenStatesProposer: self.max_num_tokens, dtype=torch.int64, device=device ) + def set_eplb_state(self, eplb_state: EplbState) -> None: + """Inject EPLB state after construction.""" + self.eplb_state = eplb_state + def propose( self, num_speculative_tokens: int, @@ -145,6 +152,12 @@ class ExtractHiddenStatesProposer: if num_tokens_across_dp is not None: num_tokens_across_dp[self.dp_rank] = num_input_tokens + if self.eplb_state is not None: + assert self.vllm_config.speculative_config is not None + self.eplb_state.prepare_forward( + self.vllm_config.speculative_config.draft_model_config, + num_tokens, + ) with set_forward_context( per_layer_attn_metadata, self.vllm_config, diff --git a/vllm/v1/spec_decode/llm_base_proposer.py b/vllm/v1/spec_decode/llm_base_proposer.py index c78d0660665..4eaf6e9e4f8 100644 --- a/vllm/v1/spec_decode/llm_base_proposer.py +++ b/vllm/v1/spec_decode/llm_base_proposer.py @@ -15,6 +15,7 @@ from vllm.config import ( get_layers_from_vllm_config, replace, ) +from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger @@ -79,6 +80,7 @@ class SpecDecodeBaseProposer: self.dtype = vllm_config.model_config.dtype self.max_model_len = vllm_config.model_config.max_model_len self.dp_rank = vllm_config.parallel_config.data_parallel_rank + self.eplb_state: EplbState | None = None self.num_speculative_tokens = self.speculative_config.num_speculative_tokens # We need to get the hidden size from the draft model config because @@ -328,6 +330,10 @@ class SpecDecodeBaseProposer: "does not support M-RoPE yet" ) + def set_eplb_state(self, eplb_state: EplbState) -> None: + """Inject EPLB state after construction.""" + self.eplb_state = eplb_state + def _init_parallel_drafting_params(self): # For parallel drafting, we need the token ID to use for masked slots # And for EAGLE + parallel drafting, we need the hidden state tensor to use @@ -527,6 +533,12 @@ class SpecDecodeBaseProposer: if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"): self.model.model.set_skip_topk(False) + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.draft_model_config, + num_tokens, + ) + with set_forward_context( per_layer_attn_metadata, self.vllm_config, @@ -672,6 +684,12 @@ class SpecDecodeBaseProposer: if self.pass_hidden_states_to_model: model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size] + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.draft_model_config, + batch_size, + ) + with set_forward_context( per_layer_attn_metadata, self.vllm_config, diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index 8f04ce3577c..aea6fdeff83 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -8,6 +8,7 @@ from typing import Any import torch import torch.nn as nn +from vllm.config import ModelConfig from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( @@ -90,6 +91,7 @@ class EPLBController: draft_model, speculative_config.draft_model_config, ) + speculator.set_eplb_state(self.state) self._has_registered_models = True return True @@ -135,6 +137,16 @@ class EPLBController: log_stats=self.parallel_config.eplb_config.log_balancedness, ) + def prepare_forward( + self, + model_config: ModelConfig, + num_unpadded_tokens: int, + ubatch_slices: list | None = None, + ) -> None: + if self.state is None or not self.parallel_config.enable_eplb: + return + self.state.prepare_forward(model_config, num_unpadded_tokens, ubatch_slices) + def setup_from_mapping( self, model: nn.Module, diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ce1bb7f5504..0f57e8a31cd 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1284,6 +1284,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors) del intermediate_tensors + # Update the EPLB meta. + self.eplb.prepare_forward(self.model_config, input_batch.num_tokens) + # Run model. if batch_desc.cg_mode == CUDAGraphMode.FULL: # Use explicit cudagraph replay for FULL mode. diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index f1ab8677f75..747fb3a3905 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -213,6 +213,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): need_eager=is_profile, ) + self._prepare_eplb_forward(input_batch.num_tokens) + if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. assert self.prefill_cudagraph_manager is not None @@ -424,6 +426,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator): num_tokens_across_dp: torch.Tensor | None, cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, ) -> None: + self._prepare_eplb_forward(num_reqs) + idx_mapping = self.idx_mapping[:num_reqs] positions = self.input_buffers.positions[:num_reqs] # Run the draft model forward pass. diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 1bd130838a1..e4583967492 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py @@ -278,6 +278,9 @@ class DFlashSpeculator(DraftModelSpeculator): self.hidden_states[:num_target_tokens], self.context_positions[:num_target_tokens], ) + # DFlash processes all speculative tokens in one forward pass, + # so the real token count is num_query_tokens. + self._prepare_eplb_forward(num_query_tokens) self._generate_draft( num_reqs, num_query_tokens, @@ -354,6 +357,10 @@ class DFlashSpeculator(DraftModelSpeculator): self.kv_cache_config, ) + # DFlash processes all speculative tokens in one forward pass, + # so the real token count is num_query_tokens. + self._prepare_eplb_forward(num_query_tokens) + if batch_desc.cg_mode == CUDAGraphMode.FULL: assert self.query_cudagraph_manager is not None self.query_cudagraph_manager.run_fullgraph(batch_desc) diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index b06c9372a95..341ed715c7a 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -8,6 +8,7 @@ import torch.nn as nn from vllm.config import VllmConfig, get_layers_from_vllm_config from vllm.config.compilation import CUDAGraphMode +from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.v1.kv_cache_interface import KVCacheConfig @@ -106,6 +107,8 @@ class DraftModelSpeculator(BaseSpeculator): self.dp_size = vllm_config.parallel_config.data_parallel_size self.dp_rank = vllm_config.parallel_config.data_parallel_rank + self.eplb_state: EplbState | None = None + self.input_buffers = InputBuffers( max_num_reqs=self.max_num_reqs, max_num_tokens=self.max_num_tokens, @@ -165,6 +168,18 @@ class DraftModelSpeculator(BaseSpeculator): ) self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names + def set_eplb_state(self, eplb_state: EplbState) -> None: + """Inject EPLB state after construction.""" + self.eplb_state = eplb_state + + def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: + """Call EPLB prepare_forward if EPLB is active for the draft model.""" + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.speculative_config.draft_model_config, + num_unpadded_tokens, + ) + def set_attn( self, model_state: ModelState, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 6af53115775..ff1eba09fd0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -4319,6 +4319,13 @@ class GPUModelRunner( # When spec decode is enabled, defer connector finalization # (wait_for_save + clear metadata) until after draft model runs. defer_kv_connector_finalize = self.speculative_config is not None + # Update the EPLB meta. + if self.eplb_state is not None: + self.eplb_state.prepare_forward( + self.model_config, + num_tokens_unpadded, + ubatch_slices_padded, + ) with ( set_forward_context( attn_metadata, @@ -5215,6 +5222,8 @@ class GPUModelRunner( self.drafter.model, spec_config.draft_model_config, ) + assert hasattr(self.drafter, "set_eplb_state") + self.drafter.set_eplb_state(self.eplb_state) eplb_models += 1 self._setup_eagle3_aux_hidden_state_outputs()