@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a2abce646f

[EPLB] Mask padding in EPLB load recording (#38128)

Signed-off-by: ilmarkov <markovilya197@gmail.com> Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com>

Ilya Markov committed Jun 29, 2026 at 04:43 UTC a2abce646f7db07f2169dfc59433d4128bc404de

17 files changed +289 -14

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):

225 logical_to_physical_map,

226 logical_replica_count,

227 )

228 + fml.router.eplb_state.should_record_tensor = torch.ones(

229 + (), dtype=torch.bool, device=device

230 + )

231 + fml.router.eplb_state.num_unpadded_tokens_tensors = [

232 + torch.tensor(0, dtype=torch.int32, device=device)

233 + ]

234

235 out_after_shuffle = []

236 with set_forward_context(

tests/kernels/moe/test_moe_layer.py

@@ -1332,6 +1332,9 @@ def _test_body_eplb(

1332 eplb_moe_layer.router.eplb_state.should_record_tensor = torch.ones(

1333 (), dtype=torch.bool, device=device

1334 )

1335 + eplb_moe_layer.router.eplb_state.num_unpadded_tokens_tensors = [

1336 + torch.tensor(0, dtype=torch.int32, device=device)

1337 + ]

1338

1339 # Get "after" output with rearranged weights and EPLB routing

1340 with set_forward_context(

tests/kernels/moe/test_routing.py

+66

@@ -61,12 +61,14 @@ def setup_eplb_state(

61 global_num_experts, dtype=torch.int64, device="cuda"

62 )

63 should_record_tensor = torch.ones((), dtype=torch.bool, device="cuda")

64 + num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32, device="cuda")]

66 return EplbLayerState(

67 expert_load_view=expert_load_view,

68 logical_to_physical_map=logical_to_physical_map,

69 logical_replica_count=logical_replica_count,

70 should_record_tensor=should_record_tensor,

71 + num_unpadded_tokens_tensors=num_unpadded_tokens_tensors,

72 )

@@ -782,3 +784,67 @@ def test_eplb_map_with_redundancy(

784 torch.testing.assert_close(load, exp_load)

785 else:

786 assert load.sum().item() == 0

787 +

788 +

789 +@pytest.mark.parametrize(

790 + "l2p_map, replica_count, num_physical, topk_ids, "

791 + "num_unpadded, expected_out, expected_load",

792 + [

793 + pytest.param(

794 + [[0], [1], [2], [3]],

795 + [1, 1, 1, 1],

796 + 4,

797 + [[0, 1], [2, 3], [0, 2], [1, 3]],

798 + 2,

799 + [[0, 1], [2, 3], [0, 2], [1, 3]],

800 + # only rows 0,1 counted: expert 0→1, 1→1, 2→1, 3→1

801 + [1, 1, 1, 1],

802 + id="half_padded",

803 + ),

804 + pytest.param(

805 + # record everything (None = no padding info)

806 + [[0], [1], [2], [3]],

807 + [1, 1, 1, 1],

808 + 4,

809 + [[0, 1], [2, 3], [0, 2], [1, 3]],

810 + None,

811 + [[0, 1], [2, 3], [0, 2], [1, 3]],

812 + [2, 2, 2, 2],

813 + id="no_padding_info",

814 + ),

815 + ],

816 +)

817 +def test_eplb_map_num_unpadded_tokens(

818 + l2p_map,

819 + replica_count,

820 + num_physical,

821 + topk_ids,

822 + num_unpadded,

823 + expected_out,

824 + expected_load,

825 +):

826 + l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda")

827 + rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda")

828 + load = torch.zeros(num_physical, dtype=torch.int32, device="cuda")

829 + rec = torch.tensor(True, dtype=torch.bool, device="cuda")

830 + ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda")

831 + num_unpadded_t = (

832 + torch.tensor(num_unpadded, dtype=torch.int32, device="cuda")

833 + if num_unpadded is not None

834 + else None

835 + )

836 +

837 + out = eplb_map_to_physical_and_record(

838 + topk_ids=ids,

839 + expert_load_view=load,

840 + logical_to_physical_map=l2p,

841 + logical_replica_count=rc,

842 + record_enabled=rec,

843 + num_unpadded_tokens=num_unpadded_t,

844 + )

845 +

846 + exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda")

847 + torch.testing.assert_close(out, exp_out)

848 +

849 + exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda")

850 + torch.testing.assert_close(load, exp_load)

tests/model_executor/test_routed_experts_capture.py

@@ -91,6 +91,7 @@ def test_base_router_capture_with_eplb_enabled():

91 eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)

92 eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)

93 eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)

94 + eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)]

95 router = _make_router(eplb_state=eplb_state)

97 captured = []

vllm/config/utils.py

+20

@@ -203,6 +203,26 @@ class SupportsHash(Protocol):

203 def compute_hash(self) -> str: ...

204

205

206 +_config_hash_cache: dict[int, str] = {}

207 +

208 +

209 +def compute_hash_cached(config: SupportsHash) -> str:

210 + """Cache config.compute_hash() by object identity.

211 +

212 + Config objects (ModelConfig, etc.) are long-lived singletons that never

213 + mutate after construction, but compute_hash() is expensive (JSON

214 + serialization + SHA-256). This utility avoids recomputing the hash on

215 + every forward pass while keeping a single consistent key type for all

216 + lookup paths.

217 + """

218 + key = id(config)

219 + result = _config_hash_cache.get(key)

220 + if result is None:

221 + result = config.compute_hash()

222 + _config_hash_cache[key] = result

223 + return result

224 +

225 +

226 class SupportsMetricsInfo(Protocol):

227 def metrics_info(self) -> dict[str, str]: ...

228

vllm/distributed/elastic_ep/elastic_execute.py

+3 -1

@@ -458,7 +458,9 @@ class ElasticEPScalingExecutor:

458 eplb_model_state.logical_to_physical_map,

459 eplb_model_state.logical_replica_count,

460 )

461 - eplb_state._init_should_record_tensor(model)

461 + eplb_state._propagate_shared_tensors(

462 + model, eplb_model_state.num_unpadded_tokens_tensors

463 + )

464 model.update_physical_experts_metadata(

465 num_physical_experts=num_physical_experts,

466 num_local_physical_experts=num_local_experts,

vllm/distributed/eplb/eplb_state.py

+68 -5

@@ -35,6 +35,7 @@ import torch

35 from torch.distributed import ProcessGroup, all_reduce

37 from vllm.config import ModelConfig, ParallelConfig

38 +from vllm.config.utils import compute_hash_cached

39 from vllm.distributed.parallel_state import (

40 get_ep_group,

41 get_eplb_group,

@@ -206,6 +207,13 @@ class EplbModelState:

207 pending_result relies on the GIL to synchronize access between the main thread and

208 the async worker.

209 """

210 + num_unpadded_tokens_tensors: list[torch.Tensor] | None = None

211 + """

212 + Per-ubatch scalar int32 tensors holding the number of real (non-padding)

213 + tokens. Allocated once in :meth:`EplbState.add_model` so that device

214 + pointers remain stable across CUDA-graph replays. The router kernel

215 + indexes this list with ``dbo_current_ubatch_id()``.

216 + """

217

218

219 class EplbState:

@@ -253,7 +261,7 @@ class EplbState:

261 Shared scalar bool tensor for all layers. Every

262 :class:`EplbLayerState` holds a reference to the **same** object so

263 a single ``.fill_()`` updates all layers at once. Allocated on the

256 - first call to :meth:`_init_should_record_tensor`.

264 + first call to :meth:`_propagate_shared_tensors`.

265 """

266 self.is_async: bool = False

267 """

@@ -440,12 +448,19 @@ class EplbState:

448 self.policy = EPLB_POLICIES[policy_type]

449 logger.debug("Selected EPLB policy: %s", policy_type)

450

451 + # num_ubatches is 0 when DBO is disabled.

452 + num_ubatches = max(1, self.parallel_config.num_ubatches)

453 + num_unpadded_tokens_tensors = [

454 + torch.tensor(0, dtype=torch.int32, device=self.device)

455 + for _ in range(num_ubatches)

456 + ]

457 +

458 model.set_eplb_state(

459 expert_load_pass,

460 logical_to_physical_map,

461 logical_replica_count,

462 )

448 - self._init_should_record_tensor(model)

463 + self._propagate_shared_tensors(model, num_unpadded_tokens_tensors)

464 expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]

465

466 assert self.parallel_config.eplb_config.communicator is not None, (

@@ -471,10 +486,43 @@ class EplbState:

486 eplb_stats=None,

487 cuda_device_index=self.cuda_device_index,

488 communicator=communicator,

489 + num_unpadded_tokens_tensors=num_unpadded_tokens_tensors,

490 )

491 self.model_states[model_config.compute_hash()] = model_state

492 self.num_valid_physical_experts = model.num_physical_experts

493

494 + def prepare_forward(

495 + self,

496 + model_config: ModelConfig,

497 + num_unpadded_tokens: int,

498 + ubatch_slices: list | None = None,

499 + ) -> None:

500 + """Fill the per-[u]batch ``num_unpadded_tokens`` tensors before a

501 + forward pass.

502 +

503 + Args:

504 + model_config: Identifies which ``EplbModelState`` to update.

505 + num_unpadded_tokens: Total number of real (non-padding) tokens

506 + in the batch.

507 + ubatch_slices: When DBO is active, a list of

508 + ``UBatchSlice`` objects describing each micro-batch's

509 + token range. When ``None``, only ``tensors[0]`` is filled.

510 + """

511 + model_state = self.model_states.get(compute_hash_cached(model_config))

512 + if model_state is None or model_state.num_unpadded_tokens_tensors is None:

513 + return

514 + tensors = model_state.num_unpadded_tokens_tensors

515 + if ubatch_slices is None:

516 + tensors[0].fill_(num_unpadded_tokens)

517 + else:

518 + for i, ubatch_slice in enumerate(ubatch_slices):

519 + ts = ubatch_slice.token_slice

520 + # Real tokens in this ubatch: clamp the global count into

521 + # the slice range so partially-filled ubatches get the

522 + # correct count.

523 + val = max(0, min(num_unpadded_tokens, ts.stop) - ts.start)

524 + tensors[i].fill_(val)

525 +

526 def step(

527 self,

528 is_dummy: bool = False,

@@ -638,11 +686,20 @@ class EplbState:

686 self._should_record_current_step(log_stats=log_stats)

687 )

688

641 - def _init_should_record_tensor(self, model: "MixtureOfExperts") -> None: # type: ignore[name-defined]

642 - """Allocate (once) and propagate the shared ``should_record_tensor``.

689 + def _propagate_shared_tensors(

690 + self,

691 + model: "MixtureOfExperts", # type: ignore[name-defined]

692 + num_unpadded_tokens_tensors: list[torch.Tensor],

693 + ) -> None:

694 + """Propagate shared tensors to every :class:`EplbLayerState`.

695 +

696 + Allocates ``should_record_tensor`` on the first call and then

697 + assigns both it and ``num_unpadded_tokens_tensors`` to every

698 + MoE layer's :class:`EplbLayerState`. All layers reference the

699 + **same** objects so a single update is visible everywhere.

700

701 Must be called after :meth:`model.set_eplb_state` so that each

645 - layer's ``eplb_state`` is already populated with the tensor views.

702 + layer's ``eplb_state`` is already populated.

703 """

704 layer_states = [

705 layer.eplb_state

@@ -659,6 +716,7 @@ class EplbState:

716 for ls in layer_states:

717 if ls is not None:

718 ls.should_record_tensor = self.should_record_tensor

719 + ls.num_unpadded_tokens_tensors = num_unpadded_tokens_tensors

720

721 def rearrange(

722 self,

@@ -985,6 +1043,11 @@ class EplbLayerState:

1043 sliding window before the next rearrangement, so recording them wastes

1044 GPU work.

1045 """

1046 + num_unpadded_tokens_tensors: list[torch.Tensor] | None = None

1047 + """

1048 + Reference to the parent :class:`EplbModelState`'s tensor list so the

1049 + router can read the correct per-[u]batch unpadded token count.

1050 + """

1051

1052 def set_layer_state(

1053 self,

vllm/model_executor/layers/fused_moe/router/base_router.py

+35 -8

@@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (

11 )

12 from vllm.platforms import current_platform

13 from vllm.triton_utils import tl, triton

14 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id

16 if current_platform.is_cuda_alike():

@@ -22,11 +23,13 @@ if current_platform.is_cuda_alike():

23 out_ids_ptr,

24 out_ptr,

25 record_enabled_ptr,

26 + num_unpadded_tokens_ptr,

27 num_logical_experts,

28 map_slots,

29 out_size,

30 numel,

31 num_active_experts,

32 + HAS_NUM_UNPADDED: tl.constexpr,

33 BLOCK_SIZE: tl.constexpr,

34 ):

35 pid = tl.program_id(0)

@@ -50,6 +53,13 @@ if current_platform.is_cuda_alike():

53 token_idx = (offs // num_active_experts).to(tl.int64)

54 hashed = (token_idx * KNUTH_MULTIPLIER) & 0xFFFFFFFF

55 replica_idx = hashed % replica_count

56 + map_index = safe_expert_id * map_slots + replica_idx

57 + physical_id = tl.load(

58 + logical_to_physical_ptr + map_index,

59 + mask=mask & valid_expert,

60 + other=-1,

61 + )

62 + tl.store(out_ids_ptr + offs, physical_id, mask=mask)

64 # 2. Record expert load metrics.

@@ -64,16 +74,21 @@ if current_platform.is_cuda_alike():

74 # If later refactor moved all the MoE kernel calls

75 # to the modular kernel, we can move this logic there

76 # to achieve better efficiency.

67 - map_index = safe_expert_id * map_slots + replica_idx

68 - physical_id = tl.load(

69 - logical_to_physical_ptr + map_index,

70 - mask=mask & valid_expert,

71 - other=-1,

72 - )

73 - tl.store(out_ids_ptr + offs, physical_id, mask=mask)

78 record_enabled = tl.load(record_enabled_ptr) != 0

76 - valid = mask & record_enabled & (physical_id >= 0) & (physical_id < out_size)

79 + # Skip padded tokens when recording.

80 + if HAS_NUM_UNPADDED:

81 + num_unpadded_tokens = tl.load(num_unpadded_tokens_ptr)

82 + is_unpadded = offs < num_unpadded_tokens * num_active_experts

83 + else:

84 + is_unpadded = True

85 + valid = (

86 + mask

87 + & record_enabled

88 + & is_unpadded

89 + & (physical_id >= 0)

90 + & (physical_id < out_size)

91 + )

92 safe_physical_id = tl.where(physical_id >= 0, physical_id, 0)

93 tl.atomic_add(out_ptr + safe_physical_id, 1, mask=valid)

@@ -83,6 +98,7 @@ if current_platform.is_cuda_alike():

98 logical_replica_count: torch.Tensor,

99 expert_load_view: torch.Tensor,

100 record_enabled: torch.Tensor,

101 + num_unpadded_tokens: torch.Tensor | None,

102 ) -> torch.Tensor:

103 topk_ids_in = topk_ids.contiguous().to(dtype=torch.int32)

104 numel = topk_ids_in.numel()

@@ -99,11 +115,13 @@ if current_platform.is_cuda_alike():

115 out_flat,

116 expert_load_view,

117 record_enabled,

118 + num_unpadded_tokens,

119 logical_replica_count.shape[0],

120 logical_to_physical_map.shape[1],

121 expert_load_view.shape[0],

122 numel,

123 num_active_experts,

124 + HAS_NUM_UNPADDED=num_unpadded_tokens is not None,

125 BLOCK_SIZE=256,

126 )

127 return out_flat.reshape(topk_ids.shape)

@@ -114,6 +132,7 @@ if current_platform.is_cuda_alike():

132 logical_to_physical_map: torch.Tensor,

133 logical_replica_count: torch.Tensor,

134 record_enabled: torch.Tensor,

135 + num_unpadded_tokens: torch.Tensor | None = None,

136 ) -> torch.Tensor:

137 # Fused triton implementation: mapping + optional recording in one kernel.

138 return _eplb_map_and_record_triton(

@@ -122,6 +141,7 @@ if current_platform.is_cuda_alike():

141 logical_replica_count=logical_replica_count,

142 expert_load_view=expert_load_view,

143 record_enabled=record_enabled,

144 + num_unpadded_tokens=num_unpadded_tokens,

145 )

146 else:

147

@@ -131,6 +151,7 @@ else:

151 logical_to_physical_map: torch.Tensor,

152 logical_replica_count: torch.Tensor,

153 record_enabled: torch.Tensor,

154 + num_unpadded_tokens: torch.Tensor | None = None,

155 ) -> torch.Tensor:

156 return topk_ids

157

@@ -177,6 +198,8 @@ class BaseRouter(FusedMoERouter):

198 raise ValueError("EPLB requires logical_replica_count != None")

199 if eplb_state.should_record_tensor is None:

200 raise ValueError("EPLB requires should_record_tensor != None")

201 + if eplb_state.num_unpadded_tokens_tensors is None:

202 + raise ValueError("EPLB requires num_unpadded_tokens_tensors != None")

203

204 def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:

205 """Apply EPLB mapping to convert logical expert IDs to physical expert IDs."""

@@ -186,12 +209,16 @@ class BaseRouter(FusedMoERouter):

209 assert eplb_state.logical_to_physical_map is not None

210 assert eplb_state.logical_replica_count is not None

211 assert eplb_state.should_record_tensor is not None

212 + assert eplb_state.num_unpadded_tokens_tensors is not None

213 return eplb_map_to_physical_and_record(

214 topk_ids=topk_ids,

215 logical_to_physical_map=eplb_state.logical_to_physical_map,

216 logical_replica_count=eplb_state.logical_replica_count,

217 expert_load_view=eplb_state.expert_load_view,

218 record_enabled=eplb_state.should_record_tensor,

219 + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[

220 + dbo_current_ubatch_id()

221 + ],

222 )

223 return topk_ids

224

vllm/models/deepseek_v4/nvidia/model.py

@@ -70,6 +70,7 @@ from vllm.platforms import current_platform

70 from vllm.sequence import IntermediateTensors

71 from vllm.utils.math_utils import cdiv

72 from vllm.v1.attention.backends.registry import AttentionBackendEnum

73 +from vllm.v1.worker.ubatching import dbo_current_ubatch_id

76 class DeepseekV4MLP(nn.Module):

@@ -464,6 +465,11 @@ class DeepseekV4MegaMoEExperts(nn.Module):

465 logical_to_physical_map=eplb_state.logical_to_physical_map,

466 logical_replica_count=eplb_state.logical_replica_count,

467 record_enabled=eplb_state.should_record_tensor,

468 + num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[

469 + dbo_current_ubatch_id()

470 + ]

471 + if eplb_state.num_unpadded_tokens_tensors is not None

472 + else None,

473 )

474

475 prepare_megamoe_inputs(

vllm/v1/spec_decode/extract_hidden_states.py

+13

@@ -9,6 +9,7 @@ import torch

9 import torch.nn as nn

11 from vllm.config import CUDAGraphMode, VllmConfig, get_layers_from_vllm_config

12 +from vllm.distributed.eplb.eplb_state import EplbState

13 from vllm.forward_context import set_forward_context

14 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

15 from vllm.model_executor.model_loader import get_model

@@ -43,6 +44,8 @@ class ExtractHiddenStatesProposer:

44 self.dtype = vllm_config.model_config.dtype

45 self.dp_rank = vllm_config.parallel_config.data_parallel_rank

47 + self.eplb_state: EplbState | None = None

48 +

49 # Model and attention layer tracking (initialized in load_model)

50 self.model: nn.Module | None = None

51 self.attn_layer_names: list[str] = []

@@ -83,6 +86,10 @@ class ExtractHiddenStatesProposer:

86 self.max_num_tokens, dtype=torch.int64, device=device

87 )

89 + def set_eplb_state(self, eplb_state: EplbState) -> None:

90 + """Inject EPLB state after construction."""

91 + self.eplb_state = eplb_state

92 +

93 def propose(

94 self,

95 num_speculative_tokens: int,

@@ -145,6 +152,12 @@ class ExtractHiddenStatesProposer:

152 if num_tokens_across_dp is not None:

153 num_tokens_across_dp[self.dp_rank] = num_input_tokens

154

155 + if self.eplb_state is not None:

156 + assert self.vllm_config.speculative_config is not None

157 + self.eplb_state.prepare_forward(

158 + self.vllm_config.speculative_config.draft_model_config,

159 + num_tokens,

160 + )

161 with set_forward_context(

162 per_layer_attn_metadata,

163 self.vllm_config,

vllm/v1/spec_decode/llm_base_proposer.py

+18

@@ -15,6 +15,7 @@ from vllm.config import (

15 get_layers_from_vllm_config,

16 replace,

17 )

18 +from vllm.distributed.eplb.eplb_state import EplbState

19 from vllm.distributed.parallel_state import get_pp_group

20 from vllm.forward_context import set_forward_context

21 from vllm.logger import init_logger

@@ -79,6 +80,7 @@ class SpecDecodeBaseProposer:

80 self.dtype = vllm_config.model_config.dtype

81 self.max_model_len = vllm_config.model_config.max_model_len

82 self.dp_rank = vllm_config.parallel_config.data_parallel_rank

83 + self.eplb_state: EplbState | None = None

84 self.num_speculative_tokens = self.speculative_config.num_speculative_tokens

86 # We need to get the hidden size from the draft model config because

@@ -328,6 +330,10 @@ class SpecDecodeBaseProposer:

330 "does not support M-RoPE yet"

331 )

332

333 + def set_eplb_state(self, eplb_state: EplbState) -> None:

334 + """Inject EPLB state after construction."""

335 + self.eplb_state = eplb_state

336 +

337 def _init_parallel_drafting_params(self):

338 # For parallel drafting, we need the token ID to use for masked slots

339 # And for EAGLE + parallel drafting, we need the hidden state tensor to use

@@ -527,6 +533,12 @@ class SpecDecodeBaseProposer:

533 if self._share_mtp_indices and hasattr(self.model.model, "set_skip_topk"):

534 self.model.model.set_skip_topk(False)

535

536 + if self.eplb_state is not None:

537 + self.eplb_state.prepare_forward(

538 + self.draft_model_config,

539 + num_tokens,

540 + )

541 +

542 with set_forward_context(

543 per_layer_attn_metadata,

544 self.vllm_config,

@@ -672,6 +684,12 @@ class SpecDecodeBaseProposer:

684 if self.pass_hidden_states_to_model:

685 model_kwargs["hidden_states"] = self.hidden_states[:input_batch_size]

686

687 + if self.eplb_state is not None:

688 + self.eplb_state.prepare_forward(

689 + self.draft_model_config,

690 + batch_size,

691 + )

692 +

693 with set_forward_context(

694 per_layer_attn_metadata,

695 self.vllm_config,

vllm/v1/worker/gpu/eplb_utils.py

+12

@@ -8,6 +8,7 @@ from typing import Any

8 import torch

9 import torch.nn as nn

11 +from vllm.config import ModelConfig

12 from vllm.distributed.eplb.eplb_state import EplbState

13 from vllm.logger import init_logger

14 from vllm.model_executor.models.interfaces import (

@@ -90,6 +91,7 @@ class EPLBController:

91 draft_model,

92 speculative_config.draft_model_config,

93 )

94 + speculator.set_eplb_state(self.state)

95 self._has_registered_models = True

96 return True

@@ -135,6 +137,16 @@ class EPLBController:

137 log_stats=self.parallel_config.eplb_config.log_balancedness,

138 )

139

140 + def prepare_forward(

141 + self,

142 + model_config: ModelConfig,

143 + num_unpadded_tokens: int,

144 + ubatch_slices: list | None = None,

145 + ) -> None:

146 + if self.state is None or not self.parallel_config.enable_eplb:

147 + return

148 + self.state.prepare_forward(model_config, num_unpadded_tokens, ubatch_slices)

149 +

150 def setup_from_mapping(

151 self,

152 model: nn.Module,

vllm/v1/worker/gpu/model_runner.py

@@ -1284,6 +1284,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1284 model_inputs["intermediate_tensors"] = IntermediateTensors(new_tensors)

1285 del intermediate_tensors

1286

1287 + # Update the EPLB meta.

1288 + self.eplb.prepare_forward(self.model_config, input_batch.num_tokens)

1289 +

1290 # Run model.

1291 if batch_desc.cg_mode == CUDAGraphMode.FULL:

1292 # Use explicit cudagraph replay for FULL mode.

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -213,6 +213,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

213 need_eager=is_profile,

214 )

215

216 + self._prepare_eplb_forward(input_batch.num_tokens)

217 +

218 if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL:

219 # Replay the full graph for draft prefill.

220 assert self.prefill_cudagraph_manager is not None

@@ -424,6 +426,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):

426 num_tokens_across_dp: torch.Tensor | None,

427 cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,

428 ) -> None:

429 + self._prepare_eplb_forward(num_reqs)

430 +

431 idx_mapping = self.idx_mapping[:num_reqs]

432 positions = self.input_buffers.positions[:num_reqs]

433 # Run the draft model forward pass.

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -278,6 +278,9 @@ class DFlashSpeculator(DraftModelSpeculator):

278 self.hidden_states[:num_target_tokens],

279 self.context_positions[:num_target_tokens],

280 )

281 + # DFlash processes all speculative tokens in one forward pass,

282 + # so the real token count is num_query_tokens.

283 + self._prepare_eplb_forward(num_query_tokens)

284 self._generate_draft(

285 num_reqs,

286 num_query_tokens,

@@ -354,6 +357,10 @@ class DFlashSpeculator(DraftModelSpeculator):

357 self.kv_cache_config,

358 )

359

360 + # DFlash processes all speculative tokens in one forward pass,

361 + # so the real token count is num_query_tokens.

362 + self._prepare_eplb_forward(num_query_tokens)

363 +

364 if batch_desc.cg_mode == CUDAGraphMode.FULL:

365 assert self.query_cudagraph_manager is not None

366 self.query_cudagraph_manager.run_fullgraph(batch_desc)

vllm/v1/worker/gpu/spec_decode/speculator.py

+15

@@ -8,6 +8,7 @@ import torch.nn as nn

9 from vllm.config import VllmConfig, get_layers_from_vllm_config

10 from vllm.config.compilation import CUDAGraphMode

11 +from vllm.distributed.eplb.eplb_state import EplbState

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

14 from vllm.v1.kv_cache_interface import KVCacheConfig

@@ -106,6 +107,8 @@ class DraftModelSpeculator(BaseSpeculator):

107 self.dp_size = vllm_config.parallel_config.data_parallel_size

108 self.dp_rank = vllm_config.parallel_config.data_parallel_rank

109

110 + self.eplb_state: EplbState | None = None

111 +

112 self.input_buffers = InputBuffers(

113 max_num_reqs=self.max_num_reqs,

114 max_num_tokens=self.max_num_tokens,

@@ -165,6 +168,18 @@ class DraftModelSpeculator(BaseSpeculator):

168 )

169 self.draft_attn_layer_names = all_attn_layers - target_attn_layer_names

170

171 + def set_eplb_state(self, eplb_state: EplbState) -> None:

172 + """Inject EPLB state after construction."""

173 + self.eplb_state = eplb_state

174 +

175 + def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None:

176 + """Call EPLB prepare_forward if EPLB is active for the draft model."""

177 + if self.eplb_state is not None:

178 + self.eplb_state.prepare_forward(

179 + self.speculative_config.draft_model_config,

180 + num_unpadded_tokens,

181 + )

182 +

183 def set_attn(

184 self,

185 model_state: ModelState,

vllm/v1/worker/gpu_model_runner.py

@@ -4319,6 +4319,13 @@ class GPUModelRunner(

4319 # When spec decode is enabled, defer connector finalization

4320 # (wait_for_save + clear metadata) until after draft model runs.

4321 defer_kv_connector_finalize = self.speculative_config is not None

4322 + # Update the EPLB meta.

4323 + if self.eplb_state is not None:

4324 + self.eplb_state.prepare_forward(

4325 + self.model_config,

4326 + num_tokens_unpadded,

4327 + ubatch_slices_padded,

4328 + )

4329 with (

4330 set_forward_context(

4331 attn_metadata,

@@ -5215,6 +5222,8 @@ class GPUModelRunner(

5222 self.drafter.model,

5223 spec_config.draft_model_config,

5224 )

5225 + assert hasattr(self.drafter, "set_eplb_state")

5226 + self.drafter.set_eplb_state(self.eplb_state)

5227 eplb_models += 1

5228

5229 self._setup_eagle3_aux_hidden_state_outputs()