Cleanup, format, minor fixes

csrc/moe/marlin_moe_ops.cu

@@ -23,10 +23,7 @@
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
 
-// #include "marlin_moe_ops.h"
-
 #include <iostream>
-// #include <torch/extension.h>
 
 template <typename T>
 inline std::string str(T x) {

vllm/_custom_ops.py

@@ -260,9 +260,10 @@ def gptq_marlin_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
     return torch.ops._C.gptq_marlin_repack(b_q_weight, perm, size_k, size_n,
                                            num_bits)
 
+
 def gptq_marlin_moe_repack(b_q_weight: torch.Tensor, perm: torch.Tensor,
-                       size_k: int, size_n: int,
-                       num_bits: int) -> torch.Tensor:
+                           size_k: int, size_n: int,
+                           num_bits: int) -> torch.Tensor:
     num_experts = b_q_weight.shape[0]
     output = torch.empty((num_experts, size_k // 16, size_n * 2),
                          device=b_q_weight.device,

vllm/model_executor/layers/fused_moe/fused_moe.py

@@ -737,17 +737,6 @@ def fused_marlin_moe(hidden_states: torch.Tensor,
     E = w1.shape[0]
     N = w2.shape[1] * 16
 
-    # print("hidden_states shape:", hidden_states)
-    # print("w1 shape:", w1)
-    # print("w2 shape:", w2)
-    # print("gating_output shape:", gating_output)
-    # print("g_idx1 shape:", g_idx1)
-    # print("g_idx2 shape:", g_idx2)
-    # print("w1_scale shape:", w1_scale)
-    # print("w2_scale shape:", w2_scale)
-
-    # raise ValueError("stop")
-
     topk_weights, topk_ids = fused_topk(hidden_states, gating_output, topk,
                                         renormalize)
 

vllm/model_executor/layers/linear.py

@@ -222,51 +222,6 @@ def extra_repr(self) -> str:
         s += f", output_features={self.output_size}"
         s += f", bias={self.bias is not None}"
         return s
-
-class FusedLinearMarlin(LinearBase):
-
-    """
-    Args:
-        input_size: input dimension of the linear layer.
-        output_size: output dimension of the linear layer.
-        bias: If true, add bias.
-        skip_bias_add: If true, skip adding bias but instead return it.
-        params_dtype: Data type for the parameters.
-        quant_config: Quantization configure.
-    """
-
-    def __init__(self,
-                 input_size13: int,
-                 input_size2: int,
-                 quant_config: Optional[QuantizationConfig] = None):
-        # calling with inputsize13 is a bit of an ugly workaround,
-        # it's not used for anything
-        super().__init__(input_size13, input_size13, False, None,
-                         quant_config)
-        self.input_size13 = input_size13
-        self.input_size2 = input_size2
-        self.output_size13 = input_size2
-        self.output_size2 = input_size13
-
-        # All the linear layer supports quant method.
-        assert self.quant_method is not None
-        self.quant_method.create_weights(self, self.input_size13, self.input_size2,
-                                         self.output_size13, self.output_size2,
-                                         self.input_size13, self.input_size2,
-                                         self.params_dtype)
-
-        self.register_parameter("bias", None)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        assert self.quant_method is not None
-        output = self.quant_method.apply(self, x, None)
-        return output, None
-
-    def extra_repr(self) -> str:
-        s = f"in_features={self.input_size}"
-        s += f", output_features={self.output_size}"
-        s += f", bias={self.bias is not None}"
-        return s
 
 
 class ColumnParallelLinear(LinearBase):