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 10, 11:30 UTC
[MoE Refactoring][Bugfix]Wrap WNA16 Triton kernel into mk and change compressed tensor kernel selection (#31752)
Signed-off-by: Robert Shaw <robshaw@redhat.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Robert Shaw <robshaw@redhat.com>
Yongye Zhu committed
Jan 8, 2026 at 16:01 UTC
d62cfe546dc587925b59eebf5aa736ab70d47a1f
3 files changed
+148
-4
vllm/model_executor/layers/fused_moe/__init__.py
+2
@@ -85,6 +85,7 @@ if HAS_TRITON:
85
from vllm.model_executor.layers.fused_moe.fused_moe import (
86
GroupedTopk,
87
TritonExperts,
88
+ TritonWNA16Experts,
89
fused_experts,
90
fused_topk,
91
get_config_file_name,
@@ -103,6 +104,7 @@ if HAS_TRITON:
104
"CutlassBatchedExpertsFp8",
105
"CutlassExpertsW4A8Fp8",
106
"TritonExperts",
107
+ "TritonWNA16Experts",
108
"BatchedTritonExperts",
109
"DeepGemmExperts",
110
"BatchedDeepGemmExperts",
vllm/model_executor/layers/fused_moe/fused_moe.py
+144
-2
@@ -624,11 +624,11 @@ def invoke_fused_moe_wna16_triton_kernel(
624
compute_type: tl.dtype,
625
use_int8_w8a16: bool,
626
use_int4_w4a16: bool,
627
- block_shape: list[int],
627
+ block_shape: list[int] | None,
628
):
629
assert B_scale is not None and B_scale.ndim == 3
630
assert B_zp is None or B_zp.ndim == 3
631
- assert block_shape is None or block_shape[0] == 0
631
+ assert block_shape is not None and block_shape[0] == 0
632
633
M = A.size(0)
634
num_tokens = M * top_k
@@ -2447,6 +2447,148 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
2447
ops.moe_sum(input, output)
2448
2449
2450
+class TritonWNA16Experts(TritonExperts):
2451
+ def __init__(
2452
+ self,
2453
+ quant_config: FusedMoEQuantConfig,
2454
+ ):
2455
+ super().__init__(quant_config)
2456
+
2457
+ def apply(
2458
+ self,
2459
+ output: torch.Tensor,
2460
+ hidden_states: torch.Tensor,
2461
+ w1: torch.Tensor,
2462
+ w2: torch.Tensor,
2463
+ topk_weights: torch.Tensor,
2464
+ topk_ids: torch.Tensor,
2465
+ activation: str,
2466
+ global_num_experts: int,
2467
+ expert_map: torch.Tensor | None,
2468
+ a1q_scale: torch.Tensor | None,
2469
+ a2_scale: torch.Tensor | None,
2470
+ workspace13: torch.Tensor,
2471
+ workspace2: torch.Tensor,
2472
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
2473
+ apply_router_weight_on_input: bool,
2474
+ ):
2475
+ # Check constraints.
2476
+ if self.quant_config.use_int4_w4a16:
2477
+ assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch"
2478
+ else:
2479
+ assert hidden_states.size(-1) == w1.size(2), (
2480
+ f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}"
2481
+ )
2482
+
2483
+ assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
2484
+ assert hidden_states.dim() == 2
2485
+ assert w1.stride(-1) == 1, "Stride of last dimension must be 1"
2486
+ assert w2.stride(-1) == 1, "Stride of last dimension must be 1"
2487
+ assert hidden_states.dtype in [
2488
+ torch.float32,
2489
+ torch.float16,
2490
+ torch.bfloat16,
2491
+ torch.float8_e4m3fn,
2492
+ torch.float8_e4m3fnuz,
2493
+ ]
2494
+
2495
+ E, num_tokens, N, K, top_k_num = self.moe_problem_size(
2496
+ hidden_states, w1, w2, topk_ids
2497
+ )
2498
+
2499
+ if global_num_experts == -1:
2500
+ global_num_experts = E
2501
+
2502
+ config = try_get_optimal_moe_config(
2503
+ w1.size(),
2504
+ w2.size(),
2505
+ top_k_num,
2506
+ self.quant_config.config_name(hidden_states.dtype),
2507
+ num_tokens,
2508
+ block_shape=self.block_shape,
2509
+ )
2510
+
2511
+ if hidden_states.dtype == torch.bfloat16:
2512
+ compute_type = tl.bfloat16
2513
+ elif hidden_states.dtype == torch.float16:
2514
+ compute_type = tl.float16
2515
+ elif hidden_states.dtype == torch.float32:
2516
+ compute_type = tl.float32
2517
+ elif (
2518
+ hidden_states.dtype == torch.float8_e4m3fn
2519
+ or hidden_states.dtype == torch.float8_e4m3fnuz
2520
+ ):
2521
+ compute_type = tl.bfloat16
2522
+ else:
2523
+ raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
2524
+
2525
+ # Note that the output tensor might be in workspace1
2526
+ intermediate_cache1 = _resize_cache(workspace2, (num_tokens, top_k_num, N))
2527
+ intermediate_cache2 = _resize_cache(
2528
+ workspace13, (num_tokens * top_k_num, N // 2)
2529
+ )
2530
+ intermediate_cache3 = _resize_cache(workspace2, (num_tokens, top_k_num, K))
2531
+
2532
+ sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
2533
+ topk_ids, config["BLOCK_SIZE_M"], global_num_experts, expert_map
2534
+ )
2535
+
2536
+ invoke_fused_moe_wna16_triton_kernel(
2537
+ hidden_states,
2538
+ w1,
2539
+ intermediate_cache1,
2540
+ self.w1_scale,
2541
+ self.quant_config.w1_zp,
2542
+ None, # topk_weights
2543
+ sorted_token_ids,
2544
+ expert_ids,
2545
+ num_tokens_post_padded,
2546
+ False, # mul_routed_weights
2547
+ top_k_num,
2548
+ config,
2549
+ compute_type=compute_type,
2550
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
2551
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
2552
+ block_shape=self.block_shape,
2553
+ )
2554
+
2555
+ self.activation(
2556
+ activation, intermediate_cache2, intermediate_cache1.view(-1, N)
2557
+ )
2558
+
2559
+ a2q_scale: torch.Tensor | None = None
2560
+
2561
+ qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(
2562
+ intermediate_cache2,
2563
+ a2_scale,
2564
+ self.quant_dtype,
2565
+ self.per_act_token_quant,
2566
+ self.block_shape,
2567
+ )
2568
+
2569
+ invoke_fused_moe_wna16_triton_kernel(
2570
+ qintermediate_cache2,
2571
+ w2,
2572
+ intermediate_cache3,
2573
+ self.w2_scale,
2574
+ self.quant_config.w2_zp,
2575
+ topk_weights,
2576
+ sorted_token_ids,
2577
+ expert_ids,
2578
+ num_tokens_post_padded,
2579
+ not apply_router_weight_on_input,
2580
+ 1,
2581
+ config,
2582
+ compute_type=compute_type,
2583
+ use_int8_w8a16=self.quant_config.use_int8_w8a16,
2584
+ use_int4_w4a16=self.quant_config.use_int4_w4a16,
2585
+ block_shape=self.block_shape,
2586
+ )
2587
+
2588
+ # separate function is required for MoE + LoRA
2589
+ self.moe_sum(intermediate_cache3, output)
2590
+
2591
+
2592
def modular_triton_fused_moe(
2593
quant_config: FusedMoEQuantConfig, shared_experts: torch.nn.Module | None = None
2594
) -> mk.FusedMoEModularKernel:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+2
-2
@@ -1693,11 +1693,11 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1693
from vllm.triton_utils import HAS_TRITON
1694
1695
if HAS_TRITON:
1696
- from vllm.model_executor.layers.fused_moe import TritonExperts
1696
+ from vllm.model_executor.layers.fused_moe import TritonWNA16Experts
1697
1698
layer.w13_weight = layer.w13_weight_packed
1699
layer.w2_weight = layer.w2_weight_packed
1700
- return TritonExperts(quant_config=self.moe_quant_config)
1700
+ return TritonWNA16Experts(quant_config=self.moe_quant_config)
1701
else:
1702
raise NotImplementedError(
1703
"TritonExperts requires Triton. "