common.py 22.2 KB
Newer Older
1
2
3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
4
from typing import Any
5
6
7
8
9

import torch

import vllm._custom_ops as ops
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
10
from tests.kernels.moe.utils import make_test_weights, per_token_cast_to_fp8
11
12
13
14
15
from tests.kernels.quantization.nvfp4_utils import (
    FLOAT4_E2M1_MAX,
    FLOAT8_E4M3_MAX,
    dequantize_nvfp4_to_dtype,
)
16
17
from tests.kernels.utils import torch_experts
from vllm.config import VllmConfig
18
19
20
21
22
from vllm.distributed import (
    get_dp_group,
    get_pcp_group,
    get_tensor_model_parallel_world_size,
)
23
from vllm.forward_context import set_forward_context
24
from vllm.model_executor.layers.fused_moe import fused_topk
25
from vllm.model_executor.layers.fused_moe.config import (
26
27
28
    FusedMoEConfig,
    FusedMoEParallelConfig,
    FusedMoEQuantConfig,
29
    RoutingMethodType,
30
)
31
from vllm.utils.import_utils import has_deep_ep, has_deep_gemm, has_pplx
32

33
34
35
36
37
38
39
from .mk_objects import (
    TestMoEQuantConfig,
    expert_info,
    make_fused_experts,
    make_prepare_finalize,
    prepare_finalize_info,
)
40
41
42
from .parallel_utils import ProcessGroupInfo


43
def _describe_tensor(t: torch.Tensor | None, name: str) -> str:
44
45
46
47
48
49
50
51
    if t is None:
        return f"{name} : None"
    else:
        return f"{name} : {t.shape} {t.dtype} {t.device}"


@dataclass
class Config:
52
    Ms: list[int] | int
53
54
55
    K: int
    N: int
    E: int
56
    topks: list[int] | int
57
    dtype: torch.dtype
58
    quant_config: TestMoEQuantConfig | None
59
60
61
62

    prepare_finalize_type: mk.FusedMoEPrepareAndFinalize
    fused_experts_type: mk.FusedMoEPermuteExpertsUnpermute

63
    fused_moe_chunk_size: int | None
64
65
    world_size: int

66
    torch_trace_dir_path: str | None = None
67

68
69
    def __post_init__(self):
        if self.quant_config is None:
70
            self.quant_config = TestMoEQuantConfig(None, False, False, None)
71

72
73
    def describe(self) -> str:
        s = ""
74
75
76
77
78
79
80
81
82
83
84
85
        s += "== Config:\n"
        s += f" world_size={self.world_size}\n"
        s += f" PF={self.prepare_finalize_type.__name__}\n"
        s += f" FE={self.fused_experts_type.__name__}\n"
        s += f" E={self.E}\n"
        s += f" Ms={self.Ms}\n"
        s += f" N={self.N}\n"
        s += f" K={self.K}\n"
        s += f" topk={self.topks}\n"
        s += f" dtype={self.dtype}\n"
        s += f" fused_moe_chunk_size={self.fused_moe_chunk_size}\n"
        s += " Quant:\n"
86
        if self.quant_config is not None:
87
88
89
90
            s += f"     q_dtype={self.quant_dtype}\n"
            s += f"     q_block_shape={self.quant_block_shape}\n"
            s += f"     q_per_out_ch_quant={self.is_per_out_ch_quant}\n"
            s += f"     q_per_act_token={self.is_per_act_token_quant}\n"
91
        else:
92
            s += "     quant=None\n"
93
94
95
96
97
98
99
100
        return s

    @property
    def M(self) -> int:
        assert isinstance(self.Ms, int)
        return self.Ms

    @property
101
    def quant_dtype(self) -> torch.dtype | str | None:
102
        assert self.quant_config is not None
103
104
105
106
        return self.quant_config.quant_dtype

    @property
    def is_per_act_token_quant(self) -> bool:
107
        assert self.quant_config is not None
108
109
110
111
        return self.quant_config.per_act_token_quant

    @property
    def is_per_tensor_act_quant(self) -> bool:
112
        return not self.is_per_act_token_quant and self.quant_block_shape is None
113
114
115

    @property
    def is_per_out_ch_quant(self) -> bool:
116
        assert self.quant_config is not None
117
118
119
        return self.quant_config.per_out_ch_quant

    @property
120
    def quant_block_shape(self) -> list[int] | None:
121
        assert self.quant_config is not None
122
123
124
125
126
127
128
129
130
131
132
133
134
        return self.quant_config.block_shape

    @property
    def topk(self) -> int:
        assert isinstance(self.topks, int)
        return self.topks

    @property
    def num_local_experts(self) -> int:
        return self.E // self.world_size

    def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]:
        """
135
        make env data for vllm launch.
136
137
138
139
140
141
142
143
        """
        vllm_config = VllmConfig()
        vllm_config.parallel_config.data_parallel_size = self.world_size
        vllm_config.parallel_config.enable_expert_parallel = True

        env_dict = {
            "VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())),
        }
144
145

        backend = self.all2all_backend()
146
        vllm_config.parallel_config.all2all_backend = backend
147
148
149
        if backend is not None:
            env_dict.update({"VLLM_ALL2ALL_BACKEND": backend})

150
151
        if self.fused_moe_chunk_size is not None:
            env_dict.update(
152
153
                {"VLLM_FUSED_MOE_CHUNK_SIZE": str(self.fused_moe_chunk_size)}
            )
154

155
156
157
        return vllm_config, env_dict

    def is_fp8_block_quantized(self):
158
159
160
161
        return (
            self.quant_dtype == torch.float8_e4m3fn
            and self.quant_block_shape is not None
        )
162
163

    def is_batched_prepare_finalize(self):
164
        info = prepare_finalize_info(self.prepare_finalize_type)
165
        return mk.FusedMoEActivationFormat.BatchedExperts == info.activation_format
166
167

    def is_batched_fused_experts(self):
168
        info = expert_info(self.fused_experts_type)
169
        return mk.FusedMoEActivationFormat.BatchedExperts == info.activation_format
170
171

    def is_standard_fused_experts(self):
172
173
174
175
176
177
178
179
180
181
182
183
184
185
        info = expert_info(self.fused_experts_type)
        return mk.FusedMoEActivationFormat.Standard == info.activation_format

    def fe_supported_types(self):
        info = expert_info(self.fused_experts_type)
        return info.supported_dtypes

    def pf_supported_types(self):
        info = prepare_finalize_info(self.prepare_finalize_type)
        return info.supported_dtypes

    def is_block_quant_supported(self):
        info = expert_info(self.fused_experts_type)
        return info.blocked_quantization_support
186
187

    def is_fe_supports_chunking(self):
188
189
190
191
192
193
194
195
196
197
        info = expert_info(self.fused_experts_type)
        return info.supports_chunking

    def supports_expert_map(self):
        info = expert_info(self.fused_experts_type)
        return info.supports_expert_map

    def supports_apply_weight_on_input(self):
        info = prepare_finalize_info(self.prepare_finalize_type)
        return info.supports_apply_weight_on_input
198
199

    def needs_deep_gemm(self):
200
201
        info = expert_info(self.fused_experts_type)
        return info.needs_deep_gemm
202
203

    def needs_pplx(self):
204
205
        info = prepare_finalize_info(self.prepare_finalize_type)
        return info.backend == "pplx"
206
207

    def needs_deep_ep(self):
208
        info = prepare_finalize_info(self.prepare_finalize_type)
209
210
211
212
        return (
            info.backend == "deepep_high_throughput"
            or info.backend == "deepep_low_latency"
        )
213
214

    def all2all_backend(self):
215
216
        info = prepare_finalize_info(self.prepare_finalize_type)
        return info.backend
217

218
    def is_valid(self) -> tuple[bool, str | None]:
219
220
221
        # Check prepare-finalize and fused-experts compatibility
        if self.is_batched_prepare_finalize():
            if not self.is_batched_fused_experts():
222
                return False, "Mismatched format."
223
224
        else:
            if not self.is_standard_fused_experts():
225
                return False, "Mismatched format."
226
227
228

        use_chunking = self.fused_moe_chunk_size is not None
        if use_chunking and not self.is_fe_supports_chunking():
229
            return False, "Chunking not supported."
230
231

        # Check quantization sanity
232
233
234
235
236
        if (
            int(self.is_per_act_token_quant)
            + int(self.is_per_tensor_act_quant)
            + int(self.quant_block_shape is not None)
        ) > 1:
237
            # invalid quant config
238
            return False, f"Bad quant_config {self.quant_config}."
239

240
241
        # check type support
        if self.quant_dtype is None:
242
243
244
245
            if (
                self.dtype not in self.pf_supported_types()
                or self.dtype not in self.fe_supported_types()
            ):
246
247
248
249
250
                return False, (
                    f"Unsupported type {self.dtype} not in "
                    f"{self.pf_supported_types()} and "
                    f"{self.fe_supported_types()}."
                )
251
        else:
252
253
254
255
            if (
                self.quant_dtype not in self.pf_supported_types()
                or self.quant_dtype not in self.fe_supported_types()
            ):
256
257
258
259
260
                return False, (
                    f"Unsupported quant type {self.quant_dtype} "
                    f"not in {self.pf_supported_types()} and "
                    f"{self.fe_supported_types()}."
                )
261

262
263
264
        # Check block quantization support
        is_block_quantized = self.quant_block_shape is not None
        if is_block_quantized and self.quant_dtype is None:
265
266
            return False, "No block quantization support."

267
        if is_block_quantized and not self.is_block_quant_supported():
268
            return False, "Mismatched block quantization support."
269
270

        # deep_gemm only works with block-quantized
271
        if self.needs_deep_gemm() and not is_block_quantized:
272
            return False, "Needs DeepGEMM but not block quantized."
273

274
        # Check dependencies (turn into asserts?)
275
        if self.needs_deep_ep() and not has_deep_ep():
276
            return False, "Needs DeepEP, but DeepEP not available."
277
        if self.needs_deep_gemm() and not has_deep_gemm():
278
            return False, "Needs DeepGEMM, but DeepGEMM not available."
279
        if self.needs_pplx() and not has_pplx():  # noqa: SIM103
280
            return False, "Needs PPLX, but PPLX not available."
281

282
        return True, None
283
284
285
286
287
288


@dataclass
class WeightTensors:
    w1: torch.Tensor
    w2: torch.Tensor
289
290
291
292
    w1_scale: torch.Tensor | None
    w2_scale: torch.Tensor | None
    w1_gs: torch.Tensor | None = None
    w2_gs: torch.Tensor | None = None
293
294
295
296

    def describe(self):
        s = ""
        s += "== Weight Tensors: \n"
297
298
299
300
301
302
        s += f" - {_describe_tensor(self.w1, 'w1')} \n"
        s += f" - {_describe_tensor(self.w2, 'w2')} \n"
        s += f" - {_describe_tensor(self.w1_scale, 'w1_scale')} \n"
        s += f" - {_describe_tensor(self.w2_scale, 'w2_scale')} \n"
        s += f" - {_describe_tensor(self.w1_gs, 'w1_gs')} \n"
        s += f" - {_describe_tensor(self.w2_gs, 'w2_gs')} \n"
303
304
        return s

305
306
    def is_quantized(self) -> bool:
        # or w1_scale is not None?
307
308
309
310
311
        return (
            self.w1.dtype == torch.float8_e4m3fn
            or self.w1.dtype == torch.uint8
            or self.w1.dtype == torch.int8
        )
312

313
    def to_current_device(self):
314
315
316
        device = torch.cuda.current_device()
        self.w1 = self.w1.to(device=device)
        self.w2 = self.w2.to(device=device)
317

318
319
320
321
        if self.w1_scale is not None:
            self.w1_scale = self.w1_scale.to(device=device)
        if self.w2_scale is not None:
            self.w2_scale = self.w2_scale.to(device=device)
322

323
        if self.w1_gs is not None:
324
325
326
            self.w1_gs = self.w1_gs.to(device=device)
        if self.w2_gs is not None:
            self.w2_gs = self.w2_gs.to(device=device)
327

328
    def slice_weights(self, rank: int, num_local_experts: int) -> "WeightTensors":
329
330
331
332
        s = rank * num_local_experts
        e = s + num_local_experts
        w1 = self.w1[s:e, :, :]
        w2 = self.w2[s:e, :, :]
333
334
        w1_scale = self.w1_scale[s:e, :, :] if self.w1_scale is not None else None
        w2_scale = self.w2_scale[s:e, :, :] if self.w2_scale is not None else None
335
336
        w1_gs = self.w1_gs[s:e] if self.w1_gs is not None else None
        w2_gs = self.w2_gs[s:e] if self.w2_gs is not None else None
337

338
        return WeightTensors(w1, w2, w1_scale, w2_scale, w1_gs, w2_gs)
339

340
341
342
    @staticmethod
    def make(config: Config) -> "WeightTensors":
        (_, w1, w1_scale, w1_gs), (_, w2, w2_scale, w2_gs) = make_test_weights(
343
344
345
            e=config.E,
            n=config.N,
            k=config.K,
346
347
348
            in_dtype=config.dtype,
            quant_dtype=config.quant_dtype,
            block_shape=config.quant_block_shape,
349
350
            # or config.is_per_out_ch_quant
            per_out_ch_quant=config.is_per_act_token_quant,
351
352
353
        )
        return WeightTensors(
            w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale, w1_gs=w1_gs, w2_gs=w2_gs
354
355
356
357
358
359
        )


@dataclass
class RankTensors:
    hidden_states: torch.Tensor
360
    hidden_states_scale: torch.Tensor | None
361
362
363

    topk_weights: torch.Tensor
    topk_ids: torch.Tensor
364
    expert_map: torch.Tensor | None
365
366
367
368

    def describe(self):
        s = ""
        s += "== Rank Tensors: \n"
369
370
371
372
373
        s += f" - {_describe_tensor(self.hidden_states, 'HS')} \n"
        s += f" - {_describe_tensor(self.hidden_states_scale, 'HS_scale')} \n"
        s += f" - {_describe_tensor(self.topk_weights, 'topk_weights')} \n"
        s += f" - {_describe_tensor(self.topk_ids, 'topk_ids')} \n"
        s += f" - {_describe_tensor(self.expert_map, 'expert_map')} \n"
374
375
376
377
        return s

    @staticmethod
    def make_hidden_states(
378
        config: Config,
379
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
380
381
382
383
        """
        Return hidden_states
        """
        m, k, dtype = (config.M, config.K, config.dtype)
384
        a = torch.randn((m, k), device=torch.cuda.current_device(), dtype=dtype) / 15.0
385
386
387
388
389
390
391

        if config.quant_dtype is None:
            return a, None

        # We dequant and use that as hidden_states so the tests are stable.
        # quantizing and dequantizing yield slightly different results
        # depending on the hardware. Here we, quantize and dequantize
392
        # first - so further quantize and dequantize will yield the same
393
394
        # values.
        if config.is_per_tensor_act_quant:
395
            a_q, a_scales = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=False)
396
397
398
            return a_q.float().mul(a_scales).to(dtype), a_scales

        if config.is_per_act_token_quant:
399
            a_q, a_scales = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True)
400
401
402
403
404
            return a_q.float().mul(a_scales).to(dtype), None

        assert config.quant_block_shape is not None
        block_k = config.quant_block_shape[1]
        a_q, a_scales = per_token_cast_to_fp8(a, block_size=block_k)
405
406
407
        return a_q.float().view((-1, block_k)).mul(a_scales.view(-1, 1)).view(m, k).to(
            dtype
        ), None
408
409
410
411
412

    @staticmethod
    def make(config: Config, pgi: ProcessGroupInfo):
        dtype = config.dtype
        topk, m, _ = (config.topk, config.M, config.K)
413
        hidden_states, hidden_states_scale = RankTensors.make_hidden_states(config)
414

415
416
417
        num_local_experts, global_num_experts = (config.num_local_experts, config.E)
        score = torch.randn((m, global_num_experts), device="cuda", dtype=dtype)
        topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk, False)
418
419
420
421
422
423
424

        # distribute topk_ids evenly
        for mi in range(m):
            topk_ids[mi] = torch.randperm(config.E)[:topk]
        topk_ids = topk_ids.to(device=torch.cuda.current_device())

        expert_map = None
425
        if config.world_size > 1 and config.supports_expert_map():
426
427
428
            expert_map = torch.full(
                (global_num_experts,), fill_value=-1, dtype=torch.int32
            )
429
430
431
            s = pgi.rank * num_local_experts
            e = s + num_local_experts
            expert_map[s:e] = torch.tensor(list(range(num_local_experts)))
432
433
434
            expert_map = expert_map.to(
                device=torch.cuda.current_device(), dtype=torch.int32
            )
435
436
437
438
439
440
441
442
443
444

        return RankTensors(
            hidden_states=hidden_states,
            hidden_states_scale=hidden_states_scale,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
            expert_map=expert_map,
        )


445
446
447
def reference_moe_impl(
    config: Config, weights: WeightTensors, rank_tensors: RankTensors
) -> torch.Tensor:
448
449
450
451
452
453
454
455
456
457
458
459
    if config.quant_dtype == "nvfp4":
        quant_blocksize = 16
        dtype = config.dtype

        w1_q = weights.w1
        w1_blockscale = weights.w1_scale
        w1_gs = weights.w1_gs

        w2_q = weights.w2
        w2_blockscale = weights.w2_scale
        w2_gs = weights.w2_gs

460
461
462
463
        a_global_scale = (
            (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)
            / torch.amax(rank_tensors.hidden_states.flatten(), dim=-1)
        ).to(torch.float32)
464
465
466
467
468
469
470
471
472
473
474
475

        assert w1_gs is not None
        assert w2_gs is not None
        assert w1_blockscale is not None
        assert w2_blockscale is not None

        assert w1_blockscale.shape[1] % 128 == 0
        assert w1_blockscale.shape[2] % 4 == 0
        assert w2_blockscale.shape[1] % 128 == 0
        assert w2_blockscale.shape[2] % 4 == 0

        a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
476
477
            rank_tensors.hidden_states, a_global_scale
        )
478

479
480
481
482
483
484
485
486
        a = dequantize_nvfp4_to_dtype(
            a_fp4,
            a_scale_interleaved,
            a_global_scale,
            dtype=dtype,
            device=a_fp4.device,
            block_size=quant_blocksize,
        )
487
488
489
490
491
492
493
494
495

        e = w1_q.shape[0]
        n = w1_q.shape[1] // 2
        k = w2_q.shape[1]

        w1 = torch.zeros((e, 2 * n, k), device="cuda", dtype=dtype)
        w2 = torch.zeros((e, k, n), device="cuda", dtype=dtype)

        for idx in range(0, e):
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
            w1[idx] = dequantize_nvfp4_to_dtype(
                w1_q[idx],
                w1_blockscale[idx],
                w1_gs[idx],
                dtype=dtype,
                device=w1_q.device,
                block_size=quant_blocksize,
            )
            w2[idx] = dequantize_nvfp4_to_dtype(
                w2_q[idx],
                w2_blockscale[idx],
                w2_gs[idx],
                dtype=dtype,
                device=w2_q.device,
                block_size=quant_blocksize,
            )
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
        a_scale = None
        w1_scale = None
        w2_scale = None
        quant_dtype = None
        per_act_token_quant = False
        block_shape = None
    else:
        a = rank_tensors.hidden_states
        a_scale = rank_tensors.hidden_states_scale
        w1 = weights.w1
        w1_scale = weights.w1_scale
        w2 = weights.w2
        w2_scale = weights.w2_scale
        quant_dtype = config.quant_dtype
        per_act_token_quant = config.is_per_act_token_quant
        block_shape = config.quant_block_shape

529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
    return torch_experts(
        a=a,
        w1=w1,
        w2=w2,
        topk_weight=rank_tensors.topk_weights,
        topk_ids=rank_tensors.topk_ids,
        global_num_experts=config.E,
        expert_map=None,
        w1_scale=w1_scale,
        w2_scale=w2_scale,
        a1_scale=a_scale,
        quant_dtype=quant_dtype,
        per_act_token_quant=per_act_token_quant,
        block_shape=block_shape,
        apply_router_weights_on_input=config.topk == 1
        and config.supports_apply_weight_on_input(),
    )
546
547


548
def _make_gscale(num_experts: int) -> torch.Tensor:
549
550
551
    return torch.ones(
        (num_experts,), device=torch.cuda.current_device(), dtype=torch.float32
    )
552
553


554
555
556
def make_modular_kernel(
    config: Config,
    vllm_config: VllmConfig,
557
    quant_config: FusedMoEQuantConfig,
558
) -> mk.FusedMoEModularKernel:
559
560
    def next_power_of_2(x):
        import math
561

562
563
        if x == 0:
            return 1
564
        return 2 ** math.ceil(math.log2(x))
565
566
567
568

    # make moe config
    moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make(
        tp_size_=get_tensor_model_parallel_world_size(),
569
        pcp_size_=get_pcp_group().world_size,
570
571
572
        dp_size_=get_dp_group().world_size,
        vllm_parallel_config=vllm_config.parallel_config,
    )
573

574
575
576
577
    moe = FusedMoEConfig(
        num_experts=config.E,
        experts_per_token=config.topk,
        hidden_dim=config.K,
578
        intermediate_size_per_partition=config.N,
579
580
581
582
        num_local_experts=config.num_local_experts,
        moe_parallel_config=moe_parallel_config,
        in_dtype=config.dtype,
        max_num_tokens=next_power_of_2(config.M),
583
584
585
        activation="silu",
        device=vllm_config.device_config.device,
        routing_method=RoutingMethodType.DeepSeekV3,
586
587
588
    )

    # make modular kernel
589
590
591
    prepare_finalize = make_prepare_finalize(
        config.prepare_finalize_type, config.all2all_backend(), moe, quant_config
    )
592
593
594
595

    fused_experts = make_fused_experts(
        config.fused_experts_type,
        moe,
596
        quant_config,
597
        prepare_finalize.num_dispatchers(),
598
        config.N,
599
    )
600
601

    modular_kernel = mk.FusedMoEModularKernel(
602
603
        prepare_finalize=prepare_finalize,
        fused_experts=fused_experts,
604
    )
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621

    return modular_kernel


def run_modular_kernel(
    pgi: ProcessGroupInfo,
    vllm_config: VllmConfig,
    config: Config,
    weights: WeightTensors,
    rank_tensors: RankTensors,
) -> torch.Tensor:
    assert isinstance(config.Ms, int)
    assert isinstance(config.topks, int)

    # weights for rank
    rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts)

622
623
624
625
626
627
628
629
630
631
    if config.quant_dtype == "nvfp4":
        gscale = _make_gscale(config.num_local_experts)
    else:
        gscale = None

    quant_config = FusedMoEQuantConfig.make(
        config.quant_dtype,
        w1_scale=rank_weights.w1_scale,
        w2_scale=rank_weights.w2_scale,
        a1_scale=rank_tensors.hidden_states_scale,
632
633
        g1_alphas=(1 / rank_weights.w1_gs) if rank_weights.w1_gs is not None else None,
        g2_alphas=(1 / rank_weights.w2_gs) if rank_weights.w2_gs is not None else None,
634
635
636
637
638
639
640
641
642
643
644
645
        a1_gscale=gscale,
        a2_gscale=gscale,
        block_shape=config.quant_block_shape,
        per_act_token_quant=config.is_per_act_token_quant,
        per_out_ch_quant=config.is_per_out_ch_quant,
    )

    mk = make_modular_kernel(config, vllm_config, quant_config)

    # impls might update the tensor in place
    hidden_states = rank_tensors.hidden_states.clone()

646
    topk_ids = rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype())
647
648

    mk_kwargs = {
649
650
651
652
653
654
655
656
657
        "hidden_states": hidden_states,
        "w1": rank_weights.w1,
        "w2": rank_weights.w2,
        "topk_weights": rank_tensors.topk_weights,
        "topk_ids": topk_ids,
        "expert_map": rank_tensors.expert_map,
        "global_num_experts": config.E,
        "apply_router_weight_on_input": config.topk == 1
        and config.supports_apply_weight_on_input(),
658
    }
659
660

    num_tokens = rank_tensors.hidden_states.shape[0]
661
662
663
    num_tokens_across_dp = torch.tensor(
        [num_tokens] * config.world_size, device="cuda", dtype=torch.int
    )
664
665

    with set_forward_context(
666
667
668
669
        None,
        vllm_config,
        num_tokens=num_tokens,
        num_tokens_across_dp=num_tokens_across_dp,
670
671
    ):
        out = mk.forward(**mk_kwargs)
672
673

    return out