lint

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

Author: bnellnm

tests/kernels/moe/test_batched_moe.py

@@ -9,8 +9,8 @@
 import triton.language as tl
 
 from tests.kernels.moe.utils import (batched_moe, make_test_weights,
-                                     torch_moe2, triton_moe,
-                                     per_block_cast_to_fp8)
+                                     per_block_cast_to_fp8, torch_moe2,
+                                     triton_moe)
 from tests.kernels.quant_utils import native_w8a8_block_matmul
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

tests/kernels/moe/test_pplx_moe.py

@@ -23,13 +23,12 @@
                                      torch_moe2)
 from tests.pplx_utils import ProcessGroupInfo, parallel_launch
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.model_executor.layers.fused_moe import (override_config,
+from vllm.model_executor.layers.fused_moe import (BatchedTritonExperts,
                                                   FusedMoEConfig,
+                                                  FusedMoEModularKernel,
                                                   fused_topk,
                                                   get_default_config,
-                                                  FusedMoEModularKernel,
-                                                  BatchedTritonExperts,
-                                                  FusedMoEModularKernel)
+                                                  override_config)
 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
     BatchedPrepareAndFinalize, NaiveBatchedExperts)
 from vllm.platforms import current_platform
@@ -561,8 +560,6 @@ def _pplx_moe(
 
     moe_config = get_default_config(m, e, n, k, topk, a.dtype, False)
 
-    use_fp8_w8a8 = qtype == torch.float8_e4m3fn
-
     device = torch.device("cuda", pgi.rank)
     a = a.to(device)
     w1 = w1.to(device)

tests/kernels/moe/utils.py

@@ -33,10 +33,7 @@ def Xnative_w8a8_block_matmul(A: torch.Tensor,
     `Bs` (float32).
     The output is returned in the specified `output_dtype`.
     """
-    if A.dtype.itemsize <= 2:
-        compute_type = torch.bfloat16
-    else:
-        compute_type = torch.float32
+    compute_type = torch.bfloat16 if A.dtype.itemsize <= 2 else torch.float32
 
     A = A.to(compute_type)
     B = B.to(compute_type).contiguous()
@@ -101,7 +98,7 @@ def torch_moe2(
     block_shape: Optional[list[int]] = None,
 ) -> torch.Tensor:
     M, K = a.shape
-    N = w1.shape[1]
+    #N = w1.shape[1]
     topk = topk_ids.shape[1]
 
     a = a.view(M, -1, K).repeat(1, topk, 1).reshape(-1, K)

vllm/model_executor/layers/fused_moe/config.py

@@ -10,7 +10,7 @@
 
 import vllm.envs as envs
 from vllm.config import ParallelConfig
-from vllm.distributed import get_dp_group
+from vllm.distributed import get_dp_group, get_tensor_model_parallel_rank
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig)
 

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -12,7 +12,8 @@
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm import _custom_ops as ops
 from vllm.logger import init_logger
-from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEQuantConfig, get_config_quant_dtype)
 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (
     _valid_deep_gemm, deep_gemm_moe_fp8)
 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (
@@ -475,7 +476,7 @@ def prepare_scales(
     from vllm.utils import round_up
     max_num_tokens = round_up(a1.shape[0], 64)
     num_tokens, hidden_dim = a1.size()
-    topk = topk_ids.size(1)
+    #topk = topk_ids.size(1)
 
     tokens_per_expert = torch.zeros(num_experts,
                                     dtype=torch.int,

vllm/model_executor/layers/fused_moe/layer.py

@@ -24,18 +24,19 @@
     is_rocm_aiter_moe_enabled)
 from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig, QuantizeMethodBase)
+from vllm.model_executor.layers.fused_moe.modular_kernel import (
+    FusedMoEModularKernel,
+    FusedMoEPermuteExpertsUnpermute,
+    FusedMoEPrepareAndFinalize)
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.platforms.interface import CpuArchEnum
 from vllm.utils import direct_register_custom_op
 
+
 has_pplx = importlib.util.find_spec("pplx_kernels") is not None
 has_deepep = importlib.util.find_spec("deep_ep") is not None
 
-from .modular_kernel import (FusedMoEModularKernel,
-                             FusedMoEPermuteExpertsUnpermute,
-                             FusedMoEPrepareAndFinalize)
-
 if current_platform.is_cuda_alike():
     from .fused_batched_moe import BatchedTritonExperts
     from .fused_moe import TritonExperts, fused_experts