fixes

Signed-off-by: Bill Nell <bnell@redhat.com>

Author: bnellnm

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -207,7 +207,8 @@ def make_ll_modular_kernel(pg: ProcessGroup, pgi: ProcessGroupInfo,
     fused_experts = BatchedDeepGemmExperts(max_num_tokens=max_tokens_per_rank,
                                            world_size=pgi.world_size,
                                            dp_size=dp_size,
-                                           block_shape=test_config.block_size)
+                                           block_shape=test_config.block_size,
+                                           per_act_token_quant=True)
     mk = FusedMoEModularKernel(prepare_finalize=a2a,
                                fused_experts=fused_experts)
     return mk

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

@@ -21,7 +21,8 @@ class BatchedDeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
     DEEPGEMM_BLOCK_SHAPE: list[int] = [128, 128]
 
     def __init__(self, max_num_tokens: int, world_size: int, dp_size: int,
-                 block_shape: list[int]):
+                 block_shape: list[int],
+                 per_act_token_quant=False):
         """
         max_num_tokens: Maximum number of tokens from a DP Rank
         world_size: Number of EP ranks
@@ -31,7 +32,7 @@ def __init__(self, max_num_tokens: int, world_size: int, dp_size: int,
         super().__init__(
             FusedMoEQuantConfig(
                 quant_dtype=torch.float8_e4m3fn,
-                per_act_token_quant=False,
+                per_act_token_quant=per_act_token_quant,
                 block_shape=block_shape,
             ))
         assert self.block_shape == self.DEEPGEMM_BLOCK_SHAPE

vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py

@@ -94,12 +94,14 @@ def _do_quant(
                 ]) and quant_dtype is not None:
             # Quantization required despite none of the inputs suggesting
             # quantization. Fallback to per_token_dynamic quant.
+            #print(f"DYNAMIC")
             _per_act_token_quant = True
         else:
             _per_act_token_quant = ((block_shape is not None) or
                                    (a1_scale is not None and a1_scale.numel() != 1)
                                    or (a2_scale is not None
                                        and a2_scale.numel() != 1))
+            #print(f"{block_shape} {a1_scale} {a2_scale}")
 
         # assert per_act_token_quant == (
         #     (block_shape is not None)
@@ -108,7 +110,7 @@ def _do_quant(
 
 
         # TODO(bnell)
-        #assert per_act_token_quant == _per_act_token_quant
+        assert per_act_token_quant == _per_act_token_quant, f"{per_act_token_quant} == {_per_act_token_quant}"
 
         num_experts, max_tokens, hidden_dim = x.size()
 

vllm/model_executor/layers/fused_moe/fused_batched_moe.py

@@ -318,8 +318,8 @@ def invoke_moe_batched_triton_kernel(
         expert_num_tokens: torch.Tensor,  # [E]
         compute_type: tl.dtype,
         # Quantization data
-        A_scale: Optional[torch.Tensor],
-        B_scale: Optional[torch.Tensor],
+        A_scale: torch.Tensor, # Optional
+        B_scale: torch.Tensor, # Optional
         B_zp: torch.Tensor,
         # Quantization schemes
         use_fp8_w8a8: bool,