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

[MoE] Encapsulate TRT-LLM BF16 weight layout handling (#57405)

Signed-off-by: mgoin <mgoin64@gmail.com>

Misha Goin committed Sep 17, 2026 at 17:47 UTC ac78445bc9ee3bb615b8b21db22fb7dac6698ece

7 files changed +138 -117

tests/kernels/moe/test_flashinfer.py

+21 -9

@@ -23,6 +23,9 @@ from vllm.model_executor.layers.fused_moe.config import (

23 from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import (

24 FlashInferExperts,

25 )

26 +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (

27 + view_as_block_major_k,

28 +)

29 from vllm.model_executor.layers.fused_moe.experts.trtllm_fp8_moe import (

30 TrtLlmFp8ExpertsModular,

31 TrtLlmFp8ExpertsMonolithic,

@@ -570,6 +573,10 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(

573 padded = w2_shape[-1]

574 w13_ptr = layer.w13_weight.data_ptr()

575 w2_ptr = layer.w2_weight.data_ptr()

576 + load_numel = (

577 + layer.w13_weight.weight_loader_numel,

578 + layer.w2_weight.weight_loader_numel,

579 + )

580

581 for _ in range(2):

582 reloaded_w13 = torch.randn_like(layer.w13_weight)

@@ -602,19 +609,24 @@ def test_unquantized_flashinfer_trtllm_weights_can_be_reprocessed(

609 assert layer.w2_weight.shape == w2_shape

610 assert layer.w13_weight.data_ptr() == w13_ptr

611 assert layer.w2_weight.data_ptr() == w2_ptr

605 - kernel_w13, kernel_w2 = method._kernel_weights(layer)

612 + assert (

613 + layer.w13_weight.weight_loader_numel,

614 + layer.w2_weight.weight_loader_numel,

615 + ) == load_numel

616 + kernel_w13 = view_as_block_major_k(layer.w13_weight)

617 + kernel_w2 = view_as_block_major_k(layer.w2_weight)

618 assert torch.equal(kernel_w13, expected_w13)

619 assert torch.equal(kernel_w2, expected_w2)

620

621

610 -def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_index):

622 +def _check_flashinfer_ipc_weights(entries, expected, mode, device_index):

623 from vllm.model_executor.model_loader.weight_cache.ipc_loader import IpcModelLoader

624 from vllm.model_executor.utils import weights_already_processed

625

626 torch.accelerator.set_device_index(device_index)

627 with pytest.MonkeyPatch.context() as monkeypatch:

628 method, layer = _make_unquantized_flashinfer_test_layer(

617 - monkeypatch, 192, is_gated, device="meta"

629 + monkeypatch, 192, is_gated=True, device="meta"

630 )

631 loader = object.__new__(IpcModelLoader)

632 loader.mode = mode

@@ -625,7 +637,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde

637 with weights_already_processed():

638 method.process_weights_after_loading(layer)

639 assert method.moe_kernel is not None

628 - actual = method._kernel_weights(layer)

640 + actual = (

641 + view_as_block_major_k(layer.w13_weight),

642 + view_as_block_major_k(layer.w2_weight),

643 + )

644 for weight, reference, pointer in zip(actual, expected, pointers):

645 assert weight.data_ptr() == pointer

646 assert torch.equal(weight.cpu(), reference)

@@ -636,11 +651,10 @@ def _check_flashinfer_ipc_weights(entries, expected, is_gated, mode, device_inde

651 torch.accelerator.synchronize()

652

653

639 -@pytest.mark.parametrize("is_gated", [True, False])

654 @pytest.mark.parametrize("cache_ndim", [3, 4])

655 @pytest.mark.parametrize("mode", ["copy", "zero_copy"])

656 def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(

643 - monkeypatch, is_gated, cache_ndim, mode

657 + monkeypatch, cache_ndim, mode

658 ):

659 """A spawned IPC consumer restores views without transient method state."""

660 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

@@ -649,13 +663,12 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(

663 from vllm.model_executor.model_loader.weight_cache.protocol import TensorEntry

664

665 method, layer = _make_unquantized_flashinfer_test_layer(

652 - monkeypatch, 192, is_gated, device="meta"

666 + monkeypatch, 192, is_gated=True, device="meta"

667 )

668 packed = convert_moe_weights_to_flashinfer_trtllm_block_layout(

669 {},

670 torch.randn_like(layer.w13_weight, device="cuda"),

671 torch.randn_like(layer.w2_weight, device="cuda"),

658 - is_gated_act_gemm=is_gated,

672 )

673 expected = [weight.cpu() for weight in packed]

674 entries = {}

@@ -667,7 +680,6 @@ def test_unquantized_flashinfer_trtllm_cached_weights_need_no_method_state(

680 args=(

681 entries,

682 expected,

670 - is_gated,

683 mode,

684 torch.accelerator.current_device_index(),

685 ),

tests/kernels/moe/test_moe_weight_loading_padded.py

-25

@@ -12,7 +12,6 @@ correctly handles this mismatch.

12 import pytest

13 import torch

15 -from vllm.model_executor.layers.fused_moe.activation import MoEActivation

15 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

16 UnquantizedMoeBackend,

17 )

@@ -465,30 +464,6 @@ class TestUnquantizedTrtLlmPrePadding:

464 assert hidden == 64

465 assert intermediate == expected

466

468 - @pytest.mark.parametrize("is_act_and_mul", [True, False])

469 - def test_create_weights_zeroes_preallocated_padding(self, is_act_and_mul: bool):

470 - method = self._make_method(1344)

471 - if not is_act_and_mul:

472 - method.moe.activation = MoEActivation.RELU2_NO_MUL

473 - method.moe.intermediate_size_per_partition = 1408

474 - layer = torch.nn.Module()

475 -

476 - method.create_weights(

477 - layer=layer,

478 - num_experts=2,

479 - hidden_size=64,

480 - intermediate_size_per_partition=1408,

481 - params_dtype=torch.bfloat16,

482 - )

483 -

484 - up_mult = 2 if is_act_and_mul else 1

485 - assert layer.w13_weight.shape == (2, up_mult * 1408, 64)

486 - assert layer.w2_weight.shape == (2, 64, 1408)

487 - assert torch.count_nonzero(layer.w13_weight[:, 1344:1408]) == 0

488 - if is_act_and_mul:

489 - assert torch.count_nonzero(layer.w13_weight[:, 2752:]) == 0

490 - assert torch.count_nonzero(layer.w2_weight[:, :, 1344:]) == 0

491 -

467

468 class TestLoadWeightsExpertBias:

469 """Some quantized exports (e.g. GPTQ, llm-compressor NVFP4) materialize

tests/kernels/moe/test_trtllm_bf16_moe.py

+49

@@ -12,6 +12,7 @@ import pytest

12 import torch

14 import vllm.model_executor.layers.fused_moe.modular_kernel as mk

15 +from tests.kernels.moe.utils import make_dummy_moe_config

16 from tests.kernels.utils import torch_moe

17 from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

18 from vllm.model_executor.layers.fused_moe.activation import MoEActivation

@@ -27,6 +28,9 @@ from vllm.model_executor.layers.fused_moe.config import (

28 from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (

29 TrtLlmBf16ExpertsModular,

30 )

31 +from vllm.model_executor.layers.fused_moe.experts.trtllm_lora_moe import (

32 + TrtLlmBf16LoRAExperts,

33 +)

34 from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

35 UnquantizedMoeBackend,

36 convert_to_unquantized_kernel_format,

@@ -150,3 +154,48 @@ def test_trtllm_bf16_moe_modular_no_graph(

154 atol=1e-1,

155 rtol=2e-1,

156 )

157 +

158 +

159 +@pytest.mark.parametrize("has_lora_delta", [False, True])

160 +@torch.inference_mode()

161 +def test_trtllm_bf16_lora_accepts_checkpoint_shaped_weights(has_lora_delta):

162 + """LoRA dispatch accepts both 3D parameters and legacy 4D packed views."""

163 + config = make_dummy_moe_config(

164 + num_experts=128, hidden_dim=256, intermediate_size=128

165 + )

166 + experts = TrtLlmBf16LoRAExperts(config, FUSED_MOE_UNQUANTIZED_CONFIG)

167 + x = torch.randn(2, 256, device="cuda", dtype=torch.bfloat16) / 10

168 + w1, w2 = convert_to_unquantized_kernel_format(

169 + UnquantizedMoeBackend.FLASHINFER_TRTLLM,

170 + config,

171 + torch.randn(128, 256, 256, device="cuda", dtype=torch.bfloat16) / 10,

172 + torch.randn(128, 256, 128, device="cuda", dtype=torch.bfloat16) / 10,

173 + )

174 + assert w1.ndim == w2.ndim == 3

175 + topk = (

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

177 + torch.ones(2, 1, device="cuda", dtype=torch.float32),

178 + )

179 + delta = torch.randn(2, 1, 256, device="cuda", dtype=torch.bfloat16)

180 +

181 + def invoke(w1, w2):

182 + result = experts.invoke_routed_moe(

183 + hidden_states=x,

184 + w1=w1,

185 + w2=w2,

186 + topk_ids_and_weights=topk,

187 + gemm1_lora_delta=delta if has_lora_delta else None,

188 + global_num_experts=128,

189 + a1q_scale=None,

190 + output=torch.empty_like(x),

191 + )

192 + if has_lora_delta:

193 + # Only compare rows belonging to real tokens, excluding padding.

194 + indices = result[2].flatten().long()

195 + return result[0][indices], result[3][indices]

196 + return result

197 +

198 + actual = invoke(w1, w2)

199 + expected = invoke(w1.view(128, 4, 256, 64), w2.view(128, 2, 256, 64))

200 + for actual_tensor, expected_tensor in zip(actual, expected):

201 + torch.testing.assert_close(actual_tensor, expected_tensor)

vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py

+13 -4

@@ -31,6 +31,15 @@ from vllm.platforms import current_platform

31 from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

34 +def view_as_block_major_k(weight: torch.Tensor) -> torch.Tensor:

35 + """View packed storage, including legacy 4D IPC cache entries."""

36 + if weight.ndim == 4:

37 + return weight

38 + experts, rows, cols = weight.shape

39 + block_k = 128 // weight.element_size()

40 + return weight.view(experts, cols // block_k, rows, block_k)

41 +

42 +

43 class TrtLlmBf16ExpertsBase:

44 """BF16 unquantized TRTLLM-Gen MoE kernels. Shared base for modular and

45 monolithic interfaces.

@@ -200,8 +209,8 @@ class TrtLlmBf16ExpertsModular(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsModular)

209 result = flashinfer.fused_moe.trtllm_bf16_routed_moe(

210 topk_ids=(topk_ids, topk_weights),

211 hidden_states=hidden_states,

203 - gemm1_weights=w1,

204 - gemm2_weights=w2,

212 + gemm1_weights=view_as_block_major_k(w1),

213 + gemm2_weights=view_as_block_major_k(w2),

214 num_experts=global_num_experts,

215 top_k=topk_ids.size(1),

216 n_group=None,

@@ -281,8 +290,8 @@ class TrtLlmBf16ExpertsMonolithic(TrtLlmBf16ExpertsBase, mk.FusedMoEExpertsMonol

290 routing_logits=router_logits,

291 routing_bias=e_score_correction_bias,

292 hidden_states=hidden_states,

284 - gemm1_weights=w1,

285 - gemm2_weights=w2,

293 + gemm1_weights=view_as_block_major_k(w1),

294 + gemm2_weights=view_as_block_major_k(w2),

295 num_experts=global_num_experts,

296 top_k=self.topk,

297 n_group=num_expert_group,

vllm/model_executor/layers/fused_moe/experts/trtllm_lora_moe.py

+5 -2

@@ -34,6 +34,9 @@ from vllm.model_executor.layers.fused_moe.experts.lora_context import MoELoRACon

34 from vllm.model_executor.layers.fused_moe.experts.lora_experts_mixin import (

35 LoRAExpertsMixin,

36 )

37 +from vllm.model_executor.layers.fused_moe.experts.trtllm_bf16_moe import (

38 + view_as_block_major_k,

39 +)

40 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

41 TopKWeightAndReduceNoOP,

42 )

@@ -527,8 +530,8 @@ class TrtLlmBf16LoRAExperts(_TrtLlmLoRAExpertsBase):

530 ret = flashinfer.fused_moe.trtllm_bf16_routed_moe(

531 topk_ids=topk_ids_and_weights,

532 hidden_states=hidden_states,

530 - gemm1_weights=w1,

531 - gemm2_weights=w2,

533 + gemm1_weights=view_as_block_major_k(w1),

534 + gemm2_weights=view_as_block_major_k(w2),

535 gemm1_lora_delta=gemm1_lora_delta,

536 num_experts=global_num_experts,

537 top_k=self.topk,

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

+26 -1

@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

26 swap_w13_to_w31,

27 )

28 from vllm.platforms import current_platform

29 +from vllm.utils.math_utils import round_up

31 if TYPE_CHECKING:

32 from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey

@@ -349,6 +350,17 @@ def select_unquantized_moe_backend(

350 )

351

352

353 +def unquantized_round_up_hidden_size_and_intermediate_size(

354 + backend: UnquantizedMoeBackend,

355 + hidden_size: int,

356 + intermediate_size: int,

357 +) -> tuple[int, int]:

358 + """Round up dimensions before allocation to satisfy the selected kernel."""

359 + if backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:

360 + intermediate_size = round_up(intermediate_size, 128)

361 + return hidden_size, intermediate_size

362 +

363 +

364 def convert_to_unquantized_kernel_format(

365 unquantized_backend: UnquantizedMoeBackend,

366 moe_config: FusedMoEConfig,

@@ -377,13 +389,26 @@ def convert_to_unquantized_kernel_format(

389 )

390 moe_config.intermediate_size_per_partition = padded_intermediate

391

392 + # Reloads only overwrite checkpoint slices. An earlier in-place

393 + # permutation can leave nonzero values in the raw padding slots.

394 + unpadded = moe_config.intermediate_size_per_partition_unpadded

395 + assert unpadded is not None

396 + if padded_intermediate > unpadded:

397 + w13_weight[:, unpadded:padded_intermediate].zero_()

398 + if is_act_and_mul:

399 + w13_weight[:, padded_intermediate + unpadded :].zero_()

400 + w2_weight[:, :, unpadded:].zero_()

401 +

402 _cache_permute_indices: dict[torch.Size, torch.Tensor] = {}

381 - w13_weight, w2_weight = convert_moe_weights_to_flashinfer_trtllm_block_layout(

403 + convert_moe_weights_to_flashinfer_trtllm_block_layout(

404 _cache_permute_indices,

405 w13_weight,

406 w2_weight,

407 is_gated_act_gemm=is_act_and_mul,

408 )

409 + # Keep checkpoint-shaped parameters for reload/IPC. The experts create

410 + # BlockMajorK views at dispatch without changing the underlying storage.

411 + return w13_weight, w2_weight

412

413 if (

414 unquantized_backend == UnquantizedMoeBackend.TRITON

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+24 -76

@@ -24,6 +24,7 @@ from vllm.model_executor.layers.fused_moe.oracle.unquantized import (

24 convert_to_unquantized_kernel_format,

25 make_unquantized_moe_kernel,

26 select_unquantized_moe_backend,

27 + unquantized_round_up_hidden_size_and_intermediate_size,

28 )

29 from vllm.model_executor.layers.fused_moe.runner.shared_experts import (

30 SharedExperts,

@@ -34,7 +35,6 @@ from vllm.model_executor.utils import (

35 set_weight_attrs,

36 )

37 from vllm.platforms import current_platform

37 -from vllm.utils.math_utils import round_up

39 if TYPE_CHECKING:

40 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts

@@ -72,9 +72,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

72 ):

73 unpadded_intermediate = self.moe.intermediate_size_per_partition_unpadded

74 assert unpadded_intermediate is not None

75 - has_intermediate_padding = (

76 - intermediate_size_per_partition > unpadded_intermediate

77 - )

75 if self.moe.is_act_and_mul:

76 w13_up_dim = 2 * intermediate_size_per_partition

77 else:

@@ -83,22 +80,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

80 assert unpadded_hidden is not None

81 unpadded_up = unpadded_intermediate * (2 if self.moe.is_act_and_mul else 1)

82 # Fused gate_up_proj (column parallel)

86 - w13_data = torch.empty(

87 - num_experts,

88 - w13_up_dim,

89 - hidden_size,

90 - dtype=params_dtype,

83 + w13_weight = torch.nn.Parameter(

84 + torch.empty(

85 + num_experts,

86 + w13_up_dim,

87 + hidden_size,

88 + dtype=params_dtype,

89 + ),

90 + requires_grad=False,

91 )

92 - if has_intermediate_padding:

93 - # The checkpoint loader only writes the unpadded slices. Zero the

94 - # tails now so no full-size padding copy is needed after loading.

95 - w13_data[:, unpadded_intermediate:intermediate_size_per_partition].zero_()

96 - if self.moe.is_act_and_mul:

97 - w13_data[

98 - :,

99 - intermediate_size_per_partition + unpadded_intermediate :,

100 - ].zero_()

101 - w13_weight = torch.nn.Parameter(w13_data, requires_grad=False)

92 layer.register_parameter("w13_weight", w13_weight)

93 set_weight_attrs(w13_weight, extra_weight_attrs)

94 w13_weight.weight_loader_numel = num_experts * unpadded_up * unpadded_hidden

@@ -111,15 +101,15 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

101 set_weight_attrs(w13_bias, extra_weight_attrs)

102 w13_bias.weight_loader_numel = num_experts * unpadded_up

103 # down_proj (row parallel)

114 - w2_data = torch.empty(

115 - num_experts,

116 - hidden_size,

117 - intermediate_size_per_partition,

118 - dtype=params_dtype,

104 + w2_weight = torch.nn.Parameter(

105 + torch.empty(

106 + num_experts,

107 + hidden_size,

108 + intermediate_size_per_partition,

109 + dtype=params_dtype,

110 + ),

111 + requires_grad=False,

112 )

120 - if has_intermediate_padding:

121 - w2_data[:, :, unpadded_intermediate:].zero_()

122 - w2_weight = torch.nn.Parameter(w2_data, requires_grad=False)

113 layer.register_parameter("w2_weight", w2_weight)

114 set_weight_attrs(w2_weight, extra_weight_attrs)

115 w2_weight.weight_loader_numel = (

@@ -147,11 +137,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

137 act_dtype,

138 moe_parallel_config,

139 )

150 - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:

151 - intermediate_size_per_partition = round_up(

152 - intermediate_size_per_partition, 128

153 - )

154 - return hidden_size, intermediate_size_per_partition

140 + return unquantized_round_up_hidden_size_and_intermediate_size(

141 + self.unquantized_backend, hidden_size, intermediate_size_per_partition

142 + )

143

144 def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:

145 # Pad the weight tensor. This is an optimization on ROCm platform, which

@@ -177,21 +165,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

165 w13: torch.Tensor,

166 w2: torch.Tensor,

167 ) -> None:

180 - w13_parameter_shape = w13.shape

181 - w2_parameter_shape = w2.shape

182 -

183 - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:

184 - unpadded = self.moe.intermediate_size_per_partition_unpadded

185 - assert unpadded is not None

186 - intermediate = w2.shape[-1]

187 - # Reloads overwrite only checkpoint slices; the previous in-place

188 - # permutation can leave nonzero values in the raw padding slots.

189 - if intermediate > unpadded:

190 - w13[:, unpadded:intermediate].zero_()

191 - if self.moe.is_act_and_mul:

192 - w13[:, intermediate + unpadded :].zero_()

193 - w2[:, :, unpadded:].zero_()

194 -

168 # Shuffle weights to runtime format.

169 w13_new, w2_new = convert_to_unquantized_kernel_format(

170 self.unquantized_backend,

@@ -199,10 +172,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

172 w13_weight=w13,

173 w2_weight=w2,

174 )

202 - if self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM:

203 - w13_new = w13_new.view(w13_parameter_shape)

204 - w2_new = w2_new.view(w2_parameter_shape)

205 -

175 # `moe_kernel` is initialized to None in FusedMoEMethodBase.__init__;

176 # On the first call we replace the parameter normally. On subsequent

177 # calls (e.g. RL weight updates that re-trigger

@@ -231,25 +200,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

200 assert self.moe_kernel is not None

201 self.moe_kernel.fused_experts.process_weights_after_loading(layer)

202

234 - def _kernel_weights(

235 - self, layer: "RoutedExperts"

236 - ) -> tuple[torch.Tensor, torch.Tensor]:

237 - if self.unquantized_backend != UnquantizedMoeBackend.FLASHINFER_TRTLLM:

238 - return layer.w13_weight, layer.w2_weight

239 -

240 - # IPC caches restore tensor shapes/storage, not method attributes.

241 - # BlockMajorK uses 128-byte blocks, i.e. 64 BF16 elements.

242 - def block_view(weight: torch.Tensor) -> torch.Tensor:

243 - if weight.ndim == 4:

244 - return weight

245 - experts, rows, cols = weight.shape

246 - return weight.view(experts, cols // 64, rows, 64)

247 -

248 - return (

249 - block_view(layer.w13_weight),

250 - block_view(layer.w2_weight),

251 - )

252 -

203 def _init_moe_kernel(self, layer: "RoutedExperts") -> None:

204 """Build the MoE kernel from the layer's current (shuffled) weights."""

205 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

@@ -354,11 +304,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

304 shared_experts_input: torch.Tensor | None,

305 ) -> torch.Tensor:

306 assert self.moe_kernel is not None

357 - w13, w2 = self._kernel_weights(layer)

307 return self.moe_kernel.apply(

308 hidden_states=x,

360 - w1=w13,

361 - w2=w2,

309 + w1=layer.w13_weight,

310 + w2=layer.w2_weight,

311 topk_weights=topk_weights,

312 topk_ids=topk_ids,

313 activation=layer.activation,

@@ -396,11 +345,10 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

345 ) -> torch.Tensor | UnfinalizedMoEOutput:

346 assert self.is_monolithic

347 assert self.moe_kernel is not None

399 - w13, w2 = self._kernel_weights(layer)

348 return self.moe_kernel.apply_monolithic(

349 x,

402 - w13,

403 - w2,

350 + layer.w13_weight,

351 + layer.w2_weight,

352 router_logits,

353 activation=layer.activation,

354 global_num_experts=layer.global_num_experts,