Disabled DeepGEMM -> batched triton fallback by default

Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

vllm/model_executor/layers/fused_moe/batched_triton_or_deep_gemm_moe.py

@@ -18,85 +18,52 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
         max_num_tokens: int,
         num_dispatchers: int,
         quant_config: FusedMoEQuantConfig,
-        allow_deep_gemm: bool = False,
+        use_deep_gemm: bool = False,
     ):
         super().__init__(quant_config)
 
-        self.batched_triton_experts = BatchedTritonExperts(
+        if use_deep_gemm:
-            max_num_tokens=max_num_tokens,
+            # Validate DeepGEMM requirements upfront
-            num_dispatchers=num_dispatchers,
+            if not self.quant_config.use_fp8_w8a8:
-            quant_config=self.quant_config,
+                raise ValueError(
-        )
+                    "DeepGEMM requires FP8 W8A8 quantization, but "
-
+                    f"quant_config.use_fp8_w8a8={self.quant_config.use_fp8_w8a8}"
-        self.allow_deep_gemm = (
+                )
-            allow_deep_gemm
+            expected_block_shape = get_mk_alignment_for_contiguous_layout()
-            and self.quant_config.use_fp8_w8a8
+            if self.block_shape != expected_block_shape:
-            and self.block_shape == get_mk_alignment_for_contiguous_layout()
+                raise ValueError(
-        )
+                    "DeepGEMM requires block_shape to match contiguous layout "
-
+                    f"alignment. Got block_shape={self.block_shape}, "
-        self.batched_deep_gemm_experts = (
+                    f"expected {expected_block_shape}"
-            BatchedDeepGemmExperts(
+                )
+            self.experts: BatchedDeepGemmExperts | BatchedTritonExperts = (
+                BatchedDeepGemmExperts(
+                    max_num_tokens=max_num_tokens,
+                    num_dispatchers=num_dispatchers,
+                    quant_config=self.quant_config,
+                )
+            )
+        else:
+            self.experts = BatchedTritonExperts(
                 max_num_tokens=max_num_tokens,
                 num_dispatchers=num_dispatchers,
                 quant_config=self.quant_config,
             )
-            if self.allow_deep_gemm
-            else None
-        )
-
-        assert (
-            self.batched_deep_gemm_experts is not None
-            or self.batched_triton_experts is not None
-        )
 
     @property
     def activation_formats(
         self,
     ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
-        if self.batched_triton_experts is not None:
+        return self.experts.activation_formats
-            assert (
-                self.batched_deep_gemm_experts is None
-                or self.batched_deep_gemm_experts.activation_formats
-                == self.batched_triton_experts.activation_formats
-            )
-            return self.batched_triton_experts.activation_formats
-        else:
-            assert self.batched_deep_gemm_experts is not None
-            return self.batched_deep_gemm_experts.activation_formats
 
     def supports_chunking(self) -> bool:
-        bdge = self.batched_deep_gemm_experts
+        return self.experts.supports_chunking()
-        bte = self.batched_triton_experts
-        return (bdge is None or bdge.supports_chunking()) and (
-            bte is None or bte.supports_chunking()
-        )
 
     def supports_expert_map(self) -> bool:
-        bdge = self.batched_deep_gemm_experts
+        return self.experts.supports_expert_map()
-        bte = self.batched_triton_experts
-        return (bdge is None or bdge.supports_expert_map()) and (
-            bte is None or bte.supports_expert_map()
-        )
 
     def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
-        bdge = self.batched_deep_gemm_experts
+        return self.experts.finalize_weight_and_reduce_impl()
-        bte = self.batched_triton_experts
-        bdge_war = bdge.finalize_weight_and_reduce_impl() if bdge else None
-        bte_war = bte.finalize_weight_and_reduce_impl() if bte else None
-        is_bdge_war = bdge_war is not None
-        is_bte_war = bte_war is not None
-
-        if is_bdge_war and is_bte_war:
-            assert bdge_war == bte_war, (
-                "Both implementations should agree on WeightAndReduce impls. "
-                f"Got bdge_war: {bdge_war}, and bte_war: {bte_war}"
-            )
-
-        if bdge_war is not None:
-            return bdge_war
-
-        assert bte_war is not None
-        return bte_war
 
     def workspace_dtype(self, act_dtype: torch.dtype) -> torch.dtype:
         return act_dtype
@@ -111,31 +78,15 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
         local_num_experts: int,
         expert_tokens_metadata: mk.ExpertTokensMetadata | None,
     ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
-        # Note: the deep gemm workspaces are strictly larger than the triton
+        return self.experts.workspace_shapes(
-        # workspaces so we can be pessimistic here and allocate for DeepGemm
+            M,
-        # even if we fall back to triton later, e.g. if expert maps are set.
+            N,
-        if self.allow_deep_gemm:
+            K,
-            assert self.batched_deep_gemm_experts is not None
+            topk,
-            return self.batched_deep_gemm_experts.workspace_shapes(
+            global_num_experts,
-                M,
+            local_num_experts,
-                N,
+            expert_tokens_metadata,
-                K,
+        )
-                topk,
-                global_num_experts,
-                local_num_experts,
-                expert_tokens_metadata,
-            )
-        else:
-            assert self.batched_triton_experts is not None
-            return self.batched_triton_experts.workspace_shapes(
-                M,
-                N,
-                K,
-                topk,
-                global_num_experts,
-                local_num_experts,
-                expert_tokens_metadata,
-            )
 
     def apply(
         self,
@@ -155,13 +106,7 @@ class BatchedTritonOrDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
         expert_tokens_meta: mk.ExpertTokensMetadata | None,
         apply_router_weight_on_input: bool,
     ):
-        experts = (
+        self.experts.apply(
-            self.batched_deep_gemm_experts
-            if self.allow_deep_gemm
-            else self.batched_triton_experts
-        )
-        assert experts is not None
-        experts.apply(
             output,
             hidden_states,
             w1,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -1086,9 +1086,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
                 max_num_tokens=max_num_tokens_per_rank,
                 num_dispatchers=prepare_finalize.num_dispatchers(),
                 quant_config=self.moe_quant_config,
-                allow_deep_gemm=(
+                use_deep_gemm=(envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM),
-                    envs.VLLM_USE_DEEP_GEMM and envs.VLLM_MOE_USE_DEEP_GEMM
-                ),
             )
         else:
             logger.debug("TritonOrDeepGemmExperts(%s)", self.__class__.__name__)