A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 19:45 UTC
[EPLB] Optmize eplb mapping and record in router for prefill (#36261)
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Ilya Markov committed
Mar 30, 2026 at 21:48 UTC
12701e8af29ba20a4bfc37edf3b30901e8789d18
6 files changed
+336
-64
tests/kernels/moe/test_routing.py
+154
@@ -8,6 +8,9 @@ import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.distributed.eplb.eplb_state import EplbLayerState
11
+from vllm.model_executor.layers.fused_moe.router.base_router import (
12
+ eplb_map_to_physical_and_record,
13
+)
14
from vllm.model_executor.layers.fused_moe.router.router_factory import (
15
create_fused_moe_router,
16
)
@@ -55,11 +58,13 @@ def setup_eplb_state(enable_eplb: bool, global_num_experts: int) -> EplbLayerSta
58
logical_replica_count = torch.ones(
59
global_num_experts, dtype=torch.int64, device="cuda"
60
)
61
+ should_record_tensor = torch.ones((), dtype=torch.bool, device="cuda")
62
63
return EplbLayerState(
64
expert_load_view=expert_load_view,
65
logical_to_physical_map=logical_to_physical_map,
66
logical_replica_count=logical_replica_count,
67
+ should_record_tensor=should_record_tensor,
68
)
69
70
@@ -581,3 +586,152 @@ def test_custom(
586
587
# hidden_states, router_logits = make_test_data(m, k, global_num_experts)
588
# topk_weights, topk_ids = router.select_experts(hidden_states, router_logits)
589
+
590
+
591
+# ---------------------------------------------------------------------------
592
+# Tests for eplb_map_to_physical_and_record
593
+# ---------------------------------------------------------------------------
594
+
595
+
596
+@pytest.mark.parametrize("record_enabled", [True, False])
597
+@pytest.mark.parametrize(
598
+ "l2p_map, replica_count, num_physical, topk_ids, expected_out, expected_load",
599
+ [
600
+ pytest.param(
601
+ # logical i → physical i
602
+ [[0], [1], [2], [3]],
603
+ [1, 1, 1, 1],
604
+ 4,
605
+ [[0, 1], [2, 3], [0, 2]],
606
+ [[0, 1], [2, 3], [0, 2]],
607
+ [2, 1, 2, 1],
608
+ id="identity",
609
+ ),
610
+ pytest.param(
611
+ # logical 0→3, 1→0, 2→1, 3→2
612
+ [[3], [0], [1], [2]],
613
+ [1, 1, 1, 1],
614
+ 4,
615
+ [[0, 1], [2, 3], [0, 2]],
616
+ [[3, 0], [1, 2], [3, 1]],
617
+ [1, 2, 1, 2],
618
+ id="shuffled",
619
+ ),
620
+ pytest.param(
621
+ # logical 0→5, 1→2, 2→7, 3→0 in a larger physical space
622
+ [[5], [2], [7], [0]],
623
+ [1, 1, 1, 1],
624
+ 8,
625
+ [[0, 1], [2, 3]],
626
+ [[5, 2], [7, 0]],
627
+ [1, 0, 1, 0, 0, 1, 0, 1],
628
+ id="sparse",
629
+ ),
630
+ ],
631
+)
632
+def test_eplb_map_no_redundancy(
633
+ record_enabled,
634
+ l2p_map,
635
+ replica_count,
636
+ num_physical,
637
+ topk_ids,
638
+ expected_out,
639
+ expected_load,
640
+):
641
+ l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda")
642
+ rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda")
643
+ load = torch.zeros(num_physical, dtype=torch.int32, device="cuda")
644
+ rec = torch.tensor(record_enabled, dtype=torch.bool, device="cuda")
645
+ ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda")
646
+
647
+ out = eplb_map_to_physical_and_record(
648
+ topk_ids=ids,
649
+ expert_load_view=load,
650
+ logical_to_physical_map=l2p,
651
+ logical_replica_count=rc,
652
+ record_enabled=rec,
653
+ )
654
+
655
+ exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda")
656
+ torch.testing.assert_close(out, exp_out)
657
+
658
+ if record_enabled:
659
+ exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda")
660
+ torch.testing.assert_close(load, exp_load)
661
+ else:
662
+ assert load.sum().item() == 0
663
+
664
+
665
+@pytest.mark.parametrize("record_enabled", [True, False])
666
+@pytest.mark.parametrize(
667
+ "l2p_map, replica_count, num_physical, topk_ids, expected_out, expected_load",
668
+ [
669
+ pytest.param(
670
+ # experts 0,1 have 2 replicas; 2,3 have 1
671
+ [[0, 4], [1, 5], [2, -1], [3, -1]],
672
+ [2, 2, 1, 1],
673
+ 6,
674
+ [[0, 1], [2, 3], [0, 2]],
675
+ # offs: 0→0%2=0→p0, 1→1%2=1→p5, 2→2%1=0→p2,
676
+ # 3→3%1=0→p3, 4→4%2=0→p0, 5→5%1=0→p2
677
+ [[0, 5], [2, 3], [0, 2]],
678
+ [2, 0, 2, 1, 0, 1],
679
+ id="partial",
680
+ ),
681
+ pytest.param(
682
+ # all 4 experts have 2 replicas
683
+ [[0, 4], [1, 5], [2, 6], [3, 7]],
684
+ [2, 2, 2, 2],
685
+ 8,
686
+ [[0, 1], [2, 3], [0, 2]],
687
+ # offs: 0→0%2=0→p0, 1→1%2=1→p5, 2→2%2=0→p2,
688
+ # 3→3%2=1→p7, 4→4%2=0→p0, 5→5%2=1→p6
689
+ [[0, 5], [2, 7], [0, 6]],
690
+ [2, 0, 1, 0, 0, 1, 1, 1],
691
+ id="full",
692
+ ),
693
+ pytest.param(
694
+ # expert 0: 4 replicas, experts 1,2: 2 replicas
695
+ [[0, 3, 5, 7], [1, 4, -1, -1], [2, 6, -1, -1]],
696
+ [4, 2, 2],
697
+ 8,
698
+ [[0, 1], [2, 0], [1, 2]],
699
+ # offs: 0→0%4=0→p0, 1→1%2=1→p4, 2→2%2=0→p2,
700
+ # 3→3%4=3→p7, 4→4%2=0→p1, 5→5%2=1→p6
701
+ [[0, 4], [2, 7], [1, 6]],
702
+ [1, 1, 1, 0, 1, 0, 1, 1],
703
+ id="uneven",
704
+ ),
705
+ ],
706
+)
707
+def test_eplb_map_with_redundancy(
708
+ record_enabled,
709
+ l2p_map,
710
+ replica_count,
711
+ num_physical,
712
+ topk_ids,
713
+ expected_out,
714
+ expected_load,
715
+):
716
+ l2p = torch.tensor(l2p_map, dtype=torch.int64, device="cuda")
717
+ rc = torch.tensor(replica_count, dtype=torch.int64, device="cuda")
718
+ load = torch.zeros(num_physical, dtype=torch.int32, device="cuda")
719
+ rec = torch.tensor(record_enabled, dtype=torch.bool, device="cuda")
720
+ ids = torch.tensor(topk_ids, dtype=torch.int32, device="cuda")
721
+
722
+ out = eplb_map_to_physical_and_record(
723
+ topk_ids=ids,
724
+ expert_load_view=load,
725
+ logical_to_physical_map=l2p,
726
+ logical_replica_count=rc,
727
+ record_enabled=rec,
728
+ )
729
+
730
+ exp_out = torch.tensor(expected_out, dtype=out.dtype, device="cuda")
731
+ torch.testing.assert_close(out, exp_out)
732
+
733
+ if record_enabled:
734
+ exp_load = torch.tensor(expected_load, dtype=torch.int32, device="cuda")
735
+ torch.testing.assert_close(load, exp_load)
736
+ else:
737
+ assert load.sum().item() == 0
tests/model_executor/test_routed_experts_capture.py
+1
@@ -62,6 +62,7 @@ def test_base_router_capture_with_eplb_enabled():
62
router.eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
63
router.eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
64
router.eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
65
+ router.eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
66
67
captured = []
68
vllm/config/parallel.py
+3
-3
@@ -53,9 +53,9 @@ All2AllBackend = Literal[
53
class EPLBConfig:
54
"""Configuration for Expert Parallel Load Balancing (EP)."""
55
56
- window_size: int = 1000
56
+ window_size: int = Field(default=1000, gt=0)
57
"""Window size for expert load recording."""
58
- step_interval: int = 3000
58
+ step_interval: int = Field(default=3000, gt=0)
59
"""
60
Interval for rearranging experts in expert parallelism.
61
@@ -71,7 +71,7 @@ class EPLBConfig:
71
Log the balancedness each step of expert parallelism.
72
This is turned off by default since it will cause communication overhead.
73
"""
74
- log_balancedness_interval: int = 1
74
+ log_balancedness_interval: int = Field(default=1, gt=0)
75
"""
76
Interval for logging the balancedness.
77
"""
vllm/distributed/elastic_ep/elastic_execute.py
+1
@@ -399,6 +399,7 @@ class ElasticEPScalingExecutor:
399
eplb_model_state.logical_to_physical_map,
400
eplb_model_state.logical_replica_count,
401
)
402
+ eplb_state._init_should_record_tensor(model)
403
model.update_physical_experts_metadata(
404
num_physical_experts=num_physical_experts,
405
num_local_physical_experts=num_local_experts,
vllm/distributed/eplb/eplb_state.py
+85
-9
@@ -272,6 +272,13 @@ class EplbState:
272
Interval for expert rearrangement steps.
273
This is a constant and is taken from the config.
274
"""
275
+ self.should_record_tensor: torch.Tensor | None = None
276
+ """
277
+ Shared scalar bool tensor for all layers. Every
278
+ :class:`EplbLayerState` holds a reference to the **same** object so
279
+ a single ``.fill_()`` updates all layers at once. Allocated on the
280
+ first call to :meth:`_init_should_record_tensor`.
281
+ """
282
self.is_async: bool = False
283
"""
284
The flag indicates whether the EPLB is running in async mode.
@@ -462,7 +469,7 @@ class EplbState:
469
logical_to_physical_map,
470
logical_replica_count,
471
)
465
-
472
+ self._init_should_record_tensor(model)
473
expert_buffer = [torch.empty_like(w) for w in model.expert_weights[0]]
474
475
model_state = EplbModelState(
@@ -582,15 +589,18 @@ class EplbState:
589
590
# Update the expert load sliding window
591
if not is_dummy:
592
+ should_record = self._should_record_current_step(log_stats=log_stats)
593
for eplb_model_state in self.model_states.values():
586
- eplb_model_state.expert_load_window[self.expert_load_window_step] = (
587
- eplb_model_state.expert_load_pass.clone()
588
- )
589
- eplb_model_state.expert_load_pass.zero_()
594
+ if should_record:
595
+ eplb_model_state.expert_load_window[
596
+ self.expert_load_window_step
597
+ ].copy_(eplb_model_state.expert_load_pass)
598
+ eplb_model_state.expert_load_pass.zero_()
599
591
- self.expert_load_window_step += 1
592
- if self.expert_load_window_step >= self.expert_load_window_size:
593
- self.expert_load_window_step = 0
600
+ if should_record:
601
+ self.expert_load_window_step += 1
602
+ if self.expert_load_window_step >= self.expert_load_window_size:
603
+ self.expert_load_window_step = 0
604
605
# Step the expert rearrangement step
606
# Note that even if this is a dummy step, we still increment the
@@ -617,11 +627,66 @@ class EplbState:
627
eplb_model_state.rebalanced
628
for eplb_model_state in self.model_states.values()
629
):
620
- # Still performing asynchronous rearrangement
630
+ # Still performing asynchronous rearrangement; update
631
+ # should_record (step > step_interval, so always True) and
632
+ # bail out before the step counter is reset.
633
+ self._update_layer_should_record(log_stats=log_stats)
634
return
635
self.expert_rearrangement_step = 0
636
self.rearrange()
637
638
+ self._update_layer_should_record(log_stats=log_stats)
639
+
640
+ def _should_record_current_step(self, log_stats: bool = False) -> bool:
641
+ """Return whether expert-load recording should be enabled this step.
642
+
643
+ Recording is enabled when we are close to either:
644
+ 1) The next rearrangement step, so the sliding window is ready.
645
+ 2) The next balancedness logging step, when log_stats is enabled.
646
+ """
647
+ steps_remaining = (
648
+ self.expert_rearrangement_step_interval - self.expert_rearrangement_step
649
+ )
650
+ should_record_for_rearrange = steps_remaining <= self.expert_load_window_size
651
+
652
+ if not log_stats:
653
+ return should_record_for_rearrange
654
+
655
+ log_interval = self.parallel_config.eplb_config.log_balancedness_interval
656
+ steps_until_next_log = (
657
+ log_interval - (self.expert_rearrangement_step % log_interval)
658
+ ) % log_interval
659
+ should_record_for_log = steps_until_next_log <= self.expert_load_window_size
660
+ return should_record_for_rearrange or should_record_for_log
661
+
662
+ def _update_layer_should_record(self, log_stats: bool = False) -> None:
663
+ """Update the shared ``should_record_tensor`` for all layers."""
664
+ if self.should_record_tensor is not None:
665
+ self.should_record_tensor.fill_(
666
+ self._should_record_current_step(log_stats=log_stats)
667
+ )
668
+
669
+ def _init_should_record_tensor(self, model: "MixtureOfExperts") -> None: # type: ignore[name-defined]
670
+ """Allocate (once) and propagate the shared ``should_record_tensor``.
671
+
672
+ Must be called after :meth:`model.set_eplb_state` so that each
673
+ layer's ``eplb_state`` is already populated with the tensor views.
674
+ """
675
+ layer_states = [
676
+ layer.eplb_state
677
+ for layer in model.moe_layers
678
+ if hasattr(layer, "eplb_state")
679
+ and isinstance(layer.eplb_state, EplbLayerState)
680
+ ]
681
+
682
+ if self.should_record_tensor is None and layer_states:
683
+ self.should_record_tensor = torch.ones(
684
+ (), dtype=torch.bool, device=self.device
685
+ )
686
+
687
+ for ls in layer_states:
688
+ ls.should_record_tensor = self.should_record_tensor
689
+
690
def rearrange(
691
self,
692
is_profile: bool = False,
@@ -993,6 +1058,17 @@ class EplbLayerState:
1058
expert_load_view: torch.Tensor | None = None
1059
logical_to_physical_map: torch.Tensor | None = None
1060
logical_replica_count: torch.Tensor | None = None
1061
+ should_record_tensor: torch.Tensor | None = None
1062
+ """
1063
+ Shared scalar bool tensor controlling whether to accumulate expert load
1064
+ metrics during this forward pass. All layers reference the **same**
1065
+ tensor object, which is owned and updated by :class:`EplbState`.
1066
+
1067
+ Set to ``False`` for the first ``step_interval - window_size`` steps of
1068
+ each rearrangement period: those steps would be overwritten in the
1069
+ sliding window before the next rearrangement, so recording them wastes
1070
+ GPU work.
1071
+ """
1072
1073
1074
def _node_count_with_rank_mapping(
vllm/model_executor/layers/fused_moe/router/base_router.py
+92
-52
@@ -10,61 +10,49 @@ from vllm.model_executor.layers.fused_moe.router.fused_moe_router import (
10
FusedMoERouter,
11
)
12
from vllm.platforms import current_platform
13
+from vllm.triton_utils import tl, triton
14
15
if current_platform.is_cuda_alike():
16
16
- @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
17
- def eplb_map_to_physical_and_record(
18
- topk_ids: torch.Tensor,
19
- expert_load_view: torch.Tensor,
20
- logical_to_physical_map: torch.Tensor,
21
- logical_replica_count: torch.Tensor,
22
- ) -> torch.Tensor:
23
- """
24
- Map the logical expert ids to physical expert ids
25
- and record the expert load metrics.
26
-
27
- This will select a pseudo-random replica for each logical expert.
28
- Only used for EPLB.
29
-
30
- Args:
31
- topk_ids: The logical expert ids.
32
- expert_load_view: The expert load view.
33
- logical_to_physical_map: The logical to physical map.
34
- logical_replica_count: The logical replica count.
17
+ @triton.jit
18
+ def _eplb_map_and_record_i32_kernel(
19
+ topk_ids_ptr,
20
+ logical_replica_count_ptr,
21
+ logical_to_physical_ptr,
22
+ out_ids_ptr,
23
+ out_ptr,
24
+ record_enabled_ptr,
25
+ num_logical_experts,
26
+ map_slots,
27
+ out_size,
28
+ numel,
29
+ BLOCK_SIZE: tl.constexpr,
30
+ ):
31
+ pid = tl.program_id(0)
32
+ offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
33
+ mask = offs < numel
34
36
- Returns:
37
- The physical expert ids.
38
- """
35
+ expert_id = tl.load(topk_ids_ptr + offs, mask=mask, other=0).to(tl.int64)
36
+ valid_expert = (expert_id >= 0) & (expert_id < num_logical_experts)
37
+ safe_expert_id = tl.where(valid_expert, expert_id, 0)
38
39
# 1. Convert the logical expert ids to physical expert ids
40
# Directly select a random replica for each logical expert
42
-
43
- # In case `indices_type` is not `torch.long` or `torch.int`,
44
- # e.g. `torch.uint32` as required by dispatch/combine kernels
45
- topk_ids_long = topk_ids.long()
46
- # Use (token position) modulo (replica count)
47
- # to deterministically choose a replica
48
- replica_count = logical_replica_count[topk_ids_long]
49
- # Flatten-position based index, reshaped back to `topk_ids` shape
50
- pos_indices = torch.arange(
51
- topk_ids.numel(), device=topk_ids.device, dtype=torch.long
52
- ).reshape_as(topk_ids)
53
- # Compute pseudo-random indices by modulo
54
- replica_indices = (pos_indices % replica_count).unsqueeze(-1)
55
- physical_ids = (
56
- logical_to_physical_map[topk_ids_long]
57
- .gather(-1, replica_indices)
58
- .squeeze(-1)
41
+ replica_count = tl.load(
42
+ logical_replica_count_ptr + safe_expert_id,
43
+ mask=mask & valid_expert,
44
+ other=1,
45
)
60
-
61
- topk_ids = physical_ids
46
+ # Avoid invalid modulo/div by forcing at least 1.
47
+ replica_count = tl.maximum(replica_count, 1)
48
+ # Match torch.compile path: use flattened token position.
49
+ replica_idx = offs % replica_count
50
51
# 2. Record expert load metrics.
52
53
# TODO(bowen): When using `FusedMoEModularKernel`, this
54
# can be done in a more unified way, since
67
- # `FusedMoEPrepareAndFinalizeModular` will return the expert
55
+ # `FusedMoEPrepareAndFinalize` will return the expert
56
# token count, in some cases directly from the kernel.
57
# However, now there are many code paths not using
58
# the modular kernel, e.g. calling `fused_experts`,
@@ -73,17 +61,63 @@ if current_platform.is_cuda_alike():
61
# If later refactor moved all the MoE kernel calls
62
# to the modular kernel, we can move this logic there
63
# to achieve better efficiency.
64
+ map_index = safe_expert_id * map_slots + replica_idx
65
+ physical_id = tl.load(
66
+ logical_to_physical_ptr + map_index,
67
+ mask=mask & valid_expert,
68
+ other=-1,
69
+ )
70
+ tl.store(out_ids_ptr + offs, physical_id, mask=mask)
71
+
72
+ record_enabled = tl.load(record_enabled_ptr) != 0
73
+ valid = mask & record_enabled & (physical_id >= 0) & (physical_id < out_size)
74
+ safe_physical_id = tl.where(physical_id >= 0, physical_id, 0)
75
+ tl.atomic_add(out_ptr + safe_physical_id, 1, mask=valid)
76
77
- # `expert_load_view`: (num_physical_experts,)
77
+ def _eplb_map_and_record_triton(
78
+ topk_ids: torch.Tensor,
79
+ logical_to_physical_map: torch.Tensor,
80
+ logical_replica_count: torch.Tensor,
81
+ expert_load_view: torch.Tensor,
82
+ record_enabled: torch.Tensor,
83
+ ) -> torch.Tensor:
84
+ topk_ids_in = topk_ids.contiguous().to(dtype=torch.int32)
85
+ numel = topk_ids_in.numel()
86
+ if numel == 0:
87
+ return topk_ids
88
+ out_flat = torch.empty((numel,), device=topk_ids.device, dtype=topk_ids.dtype)
89
+ grid = lambda meta: (triton.cdiv(numel, meta["BLOCK_SIZE"]),)
90
+ assert expert_load_view.is_contiguous()
91
+ _eplb_map_and_record_i32_kernel[grid](
92
+ topk_ids_in,
93
+ logical_replica_count.contiguous(),
94
+ logical_to_physical_map.contiguous(),
95
+ out_flat,
96
+ expert_load_view,
97
+ record_enabled,
98
+ logical_replica_count.shape[0],
99
+ logical_to_physical_map.shape[1],
100
+ expert_load_view.shape[0],
101
+ numel,
102
+ BLOCK_SIZE=256,
103
+ )
104
+ return out_flat.reshape(topk_ids.shape)
105
79
- # `torch.bincount` is not compilable, so use `scatter_add_` instead.
80
- topk_ids_flatten = topk_ids.flatten()
81
- expert_load_view.scatter_add_(
82
- dim=0,
83
- index=topk_ids_flatten.long(),
84
- src=torch.ones_like(topk_ids_flatten).to(expert_load_view),
106
+ def eplb_map_to_physical_and_record(
107
+ topk_ids: torch.Tensor,
108
+ expert_load_view: torch.Tensor,
109
+ logical_to_physical_map: torch.Tensor,
110
+ logical_replica_count: torch.Tensor,
111
+ record_enabled: torch.Tensor,
112
+ ) -> torch.Tensor:
113
+ # Fused triton implementation: mapping + optional recording in one kernel.
114
+ return _eplb_map_and_record_triton(
115
+ topk_ids=topk_ids,
116
+ logical_to_physical_map=logical_to_physical_map,
117
+ logical_replica_count=logical_replica_count,
118
+ expert_load_view=expert_load_view,
119
+ record_enabled=record_enabled,
120
)
86
- return topk_ids
121
else:
122
123
def eplb_map_to_physical_and_record(
@@ -91,8 +125,8 @@ else:
125
expert_load_view: torch.Tensor,
126
logical_to_physical_map: torch.Tensor,
127
logical_replica_count: torch.Tensor,
128
+ record_enabled: torch.Tensor,
129
) -> torch.Tensor:
95
- # CPU fallback: no EPLB so just return as is
130
return topk_ids
131
132
@@ -146,6 +180,10 @@ class BaseRouter(FusedMoERouter):
180
raise ValueError(
181
"enable_eplb=True requires logical_replica_count != None"
182
)
183
+ if self.eplb_state.should_record_tensor is None:
184
+ raise ValueError(
185
+ "enable_eplb=True requires should_record_tensor != None"
186
+ )
187
188
def _get_indices_type(self) -> torch.dtype | None:
189
"""Get the desired indices dtype from the getter function."""
@@ -159,11 +197,13 @@ class BaseRouter(FusedMoERouter):
197
assert self.eplb_state.expert_load_view is not None
198
assert self.eplb_state.logical_to_physical_map is not None
199
assert self.eplb_state.logical_replica_count is not None
200
+ assert self.eplb_state.should_record_tensor is not None
201
return eplb_map_to_physical_and_record(
202
topk_ids=topk_ids,
164
- expert_load_view=self.eplb_state.expert_load_view,
203
logical_to_physical_map=self.eplb_state.logical_to_physical_map,
204
logical_replica_count=self.eplb_state.logical_replica_count,
205
+ expert_load_view=self.eplb_state.expert_load_view,
206
+ record_enabled=self.eplb_state.should_record_tensor,
207
)
208
return topk_ids
209