4bdf400218 · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

+1

@@ -286,6 +286,7 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

286 global_num_experts: int,

287 expert_map: Optional[torch.Tensor],

288 a1q_scale: Optional[torch.Tensor],

289 + a2_scale: Optional[torch.Tensor],

290 workspace13: torch.Tensor,

291 workspace2: torch.Tensor,

292 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py

+2 -1

@@ -126,6 +126,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

126 global_num_experts: int,

127 expert_map: Optional[torch.Tensor],

128 a1q_scale: Optional[torch.Tensor],

129 + a2_scale: Optional[torch.Tensor],

130 workspace13: torch.Tensor,

131 workspace2: torch.Tensor,

132 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

@@ -136,5 +137,5 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

137 assert experts is not None

138 experts.apply(output, hidden_states, w1, w2, topk_weights, topk_ids,

139 activation, global_num_experts, expert_map, a1q_scale,

139 - workspace13, workspace2, expert_tokens_meta,

140 + a2_scale, workspace13, workspace2, expert_tokens_meta,

141 apply_router_weight_on_input)

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+3 -1

@@ -241,6 +241,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

241 global_num_experts: int,

242 expert_map: Optional[torch.Tensor],

243 a1q_scale: Optional[torch.Tensor],

244 + a2_scale: Optional[torch.Tensor],

245 workspace13: torch.Tensor,

246 workspace2: torch.Tensor,

247 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

@@ -262,7 +263,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEPermuteExpertsUnpermute):

263 run_cutlass_moe_fp8(

264 output, hidden_states, w1, w2, topk_ids, activation_callable,

265 global_num_experts, expert_map, self.w1_scale, self.w2_scale,

265 - a1q_scale, self.a2_scale, self.ab_strides1, self.ab_strides2,

266 + a1q_scale, a2_scale, self.ab_strides1, self.ab_strides2,

267 self.c_strides1, self.c_strides2, workspace13, workspace2,

268 expert_num_tokens,

269 self.out_dtype if self.out_dtype is not None else in_dtype,

@@ -705,6 +706,7 @@ class CutlassExpertsFp4(mk.FusedMoEPermuteExpertsUnpermute):

706 global_num_experts: int,

707 expert_map: Optional[torch.Tensor],

708 a1q_scale: Optional[torch.Tensor], # unused

709 + a2_scale: Optional[torch.Tensor], # unused

710 workspace13: Optional[torch.Tensor],

711 workspace2: Optional[torch.Tensor],

712 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

vllm/model_executor/layers/fused_moe/deep_gemm_moe.py

+2 -1

@@ -214,13 +214,14 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

214 global_num_experts: int,

215 expert_map: Optional[torch.Tensor],

216 a1q_scale: Optional[torch.Tensor],

217 + a2_scale: Optional[torch.Tensor],

218 workspace13: torch.Tensor,

219 workspace2: torch.Tensor,

220 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

221 apply_router_weight_on_input: bool,

222 ):

223 assert a1q_scale is not None

223 - assert self.a2_scale is None

224 + assert a2_scale is None

225 assert self.block_shape is not None

226 assert self.w1_scale is not None

227 assert self.w2_scale is not None

vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py

+1

@@ -129,6 +129,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):

129 global_num_experts: int,

130 expert_map: Optional[torch.Tensor],

131 a1q_scale: Optional[torch.Tensor],

132 + a2_scale: Optional[torch.Tensor],

133 workspace13: Optional[torch.Tensor],

134 workspace2: Optional[torch.Tensor],

135 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

+3 -1

@@ -688,6 +688,7 @@ class NaiveBatchedExperts(mk.FusedMoEPermuteExpertsUnpermute):

688 global_num_experts: int,

689 expert_map: Optional[torch.Tensor],

690 a1q_scale: Optional[torch.Tensor],

691 + a2_scale: Optional[torch.Tensor],

692 workspace13: torch.Tensor,

693 workspace2: torch.Tensor,

694 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

@@ -879,6 +880,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

880 global_num_experts: int,

881 expert_map: Optional[torch.Tensor],

882 a1q_scale: Optional[torch.Tensor],

883 + a2_scale: Optional[torch.Tensor],

884 workspace13: torch.Tensor,

885 workspace2: torch.Tensor,

886 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

@@ -970,7 +972,7 @@ class BatchedTritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

972 intermediate_cache1.view(-1, N))

973

974 qintermediate_cache2, a2q_scale = batched_moe_kernel_quantize_input(

973 - intermediate_cache2, self.a2_scale, max_num_tokens, E, N,

975 + intermediate_cache2, a2_scale, max_num_tokens, E, N,

976 expert_num_tokens, self.quant_dtype, self.per_act_token_quant,

977 self.block_shape)

978

vllm/model_executor/layers/fused_moe/fused_moe.py

+2 -1

@@ -1598,6 +1598,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1598 global_num_experts: int,

1599 expert_map: Optional[torch.Tensor],

1600 a1q_scale: Optional[torch.Tensor],

1601 + a2_scale: Optional[torch.Tensor],

1602 workspace13: torch.Tensor,

1603 workspace2: torch.Tensor,

1604 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

@@ -1690,7 +1691,7 @@ class TritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

1691 a2q_scale: Optional[torch.Tensor] = None

1692

1693 qintermediate_cache2, a2q_scale = moe_kernel_quantize_input(

1693 - intermediate_cache2, self.a2_scale, self.quant_dtype,

1694 + intermediate_cache2, a2_scale, self.quant_dtype,

1695 self.per_act_token_quant, self.block_shape)

1696

1697 invoke_fused_moe_kernel(

vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py

+1

@@ -179,6 +179,7 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):

179 global_num_experts: int,

180 expert_map: Optional[torch.Tensor],

181 a1q_scale: Optional[torch.Tensor],

182 + a2_scale: Optional[torch.Tensor],

183 workspace13: torch.Tensor,

184 workspace2: torch.Tensor,

185 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

vllm/model_executor/layers/fused_moe/modular_kernel.py

+5

@@ -519,6 +519,7 @@ class FusedMoEPermuteExpertsUnpermute(ABC):

519 global_num_experts: int,

520 expert_map: Optional[torch.Tensor],

521 a1q_scale: Optional[torch.Tensor],

522 + a2_scale: Optional[torch.Tensor],

523 workspace13: torch.Tensor,

524 workspace2: torch.Tensor,

525 expert_tokens_meta: Optional[ExpertTokensMetadata],

@@ -634,6 +635,7 @@ class FusedMoEModularKernel(torch.nn.Module):

635 local_num_experts: int,

636 expert_map: Optional[torch.Tensor],

637 a1q_scale: Optional[torch.Tensor],

638 + a2_scale: Optional[torch.Tensor],

639 expert_tokens_meta: Optional[ExpertTokensMetadata],

640 apply_router_weight_on_input: bool,

641 ) -> torch.Tensor:

@@ -671,6 +673,7 @@ class FusedMoEModularKernel(torch.nn.Module):

673 global_num_experts=global_num_experts,

674 expert_map=expert_map,

675 a1q_scale=a1q_scale,

676 + a2_scale=a2_scale,

677 workspace13=workspace13,

678 workspace2=workspace2,

679 expert_tokens_meta=expert_tokens_meta,

@@ -718,6 +721,7 @@ class FusedMoEModularKernel(torch.nn.Module):

721 local_num_experts=local_num_experts,

722 expert_map=expert_map,

723 a1q_scale=a1q_scale,

724 + a2_scale=self.fused_experts.a2_scale,

725 expert_tokens_meta=expert_tokens_meta,

726 apply_router_weight_on_input=apply_router_weight_on_input,

727 )

@@ -803,6 +807,7 @@ class FusedMoEModularKernel(torch.nn.Module):

807 local_num_experts=local_num_experts,

808 expert_map=expert_map,

809 a1q_scale=c_a1q_scale,

810 + a2_scale=c_a2_scale,

811 expert_tokens_meta=c_expert_tokens_meta,

812 apply_router_weight_on_input=apply_router_weight_on_input,

813 )

vllm/model_executor/layers/fused_moe/triton_deep_gemm_moe.py

+2

@@ -111,6 +111,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

111 global_num_experts: int,

112 expert_map: Optional[torch.Tensor],

113 a1q_scale: Optional[torch.Tensor],

114 + a2_scale: Optional[torch.Tensor],

115 workspace13: torch.Tensor,

116 workspace2: torch.Tensor,

117 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

@@ -134,6 +135,7 @@ class TritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):

135 global_num_experts,

136 expert_map,

137 a1q_scale,

138 + a2_scale,

139 workspace13,

140 workspace2,

141 expert_tokens_meta,

vllm/model_executor/layers/fused_moe/trtllm_moe.py

+1

@@ -103,6 +103,7 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute):

103 global_num_experts: int,

104 expert_map: Optional[torch.Tensor],

105 a1q_scale: Optional[torch.Tensor],

106 + a2_scale: Optional[torch.Tensor],

107 workspace13: torch.Tensor,

108 workspace2: torch.Tensor,

109 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],

[Bugfix] Fix chunked a2_scales in modular kernels (#25264)