lint

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

Author: bnellnm

tests/kernels/moe/test_block_fp8.py

@@ -6,10 +6,9 @@
 import pytest
 import torch
 
-from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
-                                       native_w8a8_block_matmul,
-                                       per_block_cast_to_fp8)
 from tests.kernels.moe.utils import make_test_weights
+from tests.kernels.quant_utils import (native_per_token_group_quant_fp8,
+                                       native_w8a8_block_matmul)
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import fused_experts
@@ -56,7 +55,8 @@
 SEEDS = [0]
 
 
-def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids, block_shape):
+def torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weight, topk_ids,
+                             block_shape):
     """Fused moe with block-wise quantization using native torch."""
     B, D = a.shape
     topk = topk_ids.size(1)
@@ -116,7 +116,11 @@ def test_w8a8_block_fp8_fused_moe(M, N, K, E, topk, block_size, dtype, seed,
     a = torch.randn((M, K), dtype=dtype) / 10
     score = torch.randn((M, E), dtype=dtype)
 
-    _, w1, w1_s, _, w2, w2_s = make_test_weights(E, N, K, dtype, torch.float8_e4m3fn,
+    _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
+                                                 N,
+                                                 K,
+                                                 dtype,
+                                                 torch.float8_e4m3fn,
                                                  per_act_token_quant=False,
                                                  block_shape=block_size)
 
@@ -203,8 +207,8 @@ def _moe_unpermute(out, inv_perm, topk, K, topk_weight):
     return (tmp_out * topk_weight.view(M, -1, 1).to(out.dtype)).sum(dim=1)
 
 
-def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, topk_weight, topk_ids,
-                                 block_shape):
+def deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s, topk_weight,
+                                 topk_ids, block_shape):
     """Fused moe with block-wise quantization using DeepGemm grouped gemm."""
     num_groups = w1.shape[0]
     M, K = a.shape
@@ -265,7 +269,11 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
     a = torch.randn((M, K), dtype=dtype) / 10
     score = torch.randn((M, E), dtype=dtype)
 
-    _, w1, w1_s, _, w2, w2_s = make_test_weights(E, N, K, dtype, torch.float8_e4m3fn,
+    _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
+                                                 N,
+                                                 K,
+                                                 dtype,
+                                                 torch.float8_e4m3fn,
                                                  per_act_token_quant=False,
                                                  block_shape=block_size)
 
@@ -281,12 +289,14 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed,
 
     # Set the context to avoid lots of warning spam.
     with set_current_vllm_config(vllm_config):
-        if False and M >= 128:
+        if M >= 128:
             ref_out = deep_gemm_w8a8_block_fp8_moe(M, K, a, w1, w2, w1_s, w2_s,
-                                                   topk_weights, topk_ids, block_size)
+                                                   topk_weights, topk_ids,
+                                                   block_size)
         else:
-            ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s, topk_weights,
-                                               topk_ids, block_size)
+            ref_out = torch_w8a8_block_fp8_moe(a, w1, w2, w1_s, w2_s,
+                                               topk_weights, topk_ids,
+                                               block_size)
 
         if use_compile:
             deep_gemm_moe_fp8_fn = torch.compile(deep_gemm_moe_fp8,

tests/kernels/moe/test_block_int8.py

@@ -6,9 +6,9 @@
 import pytest
 import torch
 
+from tests.kernels.moe.utils import make_test_weights
 from tests.kernels.quant_utils import (native_per_token_group_quant_int8,
                                        native_w8a8_block_matmul)
-from tests.kernels.moe.utils import make_test_weights
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.fused_moe import fused_moe
@@ -84,34 +84,15 @@ def test_w8a8_block_int8_fused_moe(M, N, K, E, topk, block_size, dtype, seed):
     """Tests the fused_moe kernel with W8A8 INT8 block quantization against a
     native torch reference."""
     torch.manual_seed(seed)
-    # Use a smaller factor for scale initialization to prevent large
-    # values/overflow especially when output dtype might be float16
-    # factor_for_scale = 1e-2
-    # int8_info = torch.iinfo(torch.int8)
-    # int8_max, int8_min = int8_info.max, int8_info.min
 
     a = torch.randn((M, K), dtype=dtype) / 10
     score = torch.randn((M, E), dtype=dtype)
 
-    # w1_fp32 = (torch.rand(
-    #     (E, 2 * N, K), dtype=torch.float32) - 0.5) * 2 * int8_max
-    # w1 = w1_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
-
-    # w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2 * int8_max
-    # w2 = w2_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8)
-
-    # block_n, block_k = block_size[0], block_size[1]
-    # n_tiles_w1 = (2 * N + block_n - 1) // block_n
-    # n_tiles_w2 = (K + block_n - 1) // block_n
-    # k_tiles_w1 = (K + block_k - 1) // block_k
-    # k_tiles_w2 = (N + block_k - 1) // block_k
-
-    # w1_s = (torch.rand(
-    #     (E, n_tiles_w1, k_tiles_w1), dtype=torch.float32) * factor_for_scale)
-    # w2_s = (torch.rand(
-    #     (E, n_tiles_w2, k_tiles_w2), dtype=torch.float32) * factor_for_scale)
-
-    _, w1, w1_s, _, w2, w2_s = make_test_weights(E, N, K, dtype, torch.int8,
+    _, w1, w1_s, _, w2, w2_s = make_test_weights(E,
+                                                 N,
+                                                 K,
+                                                 dtype,
+                                                 torch.int8,
                                                  per_act_token_quant=False,
                                                  block_shape=block_size)
 

tests/kernels/moe/test_cutlass_moe.py

@@ -100,13 +100,15 @@ def make_moe_tensors_8bit(m: int, k: int, n: int, e: int,
         if False:
             _, a_scale = ops.scaled_fp8_quant(
                 moe_tensors_fp16.a, use_per_token_if_dynamic=per_act_token)
-            a_q, _ = ops.scaled_fp8_quant(moe_tensors_fp16.a,
-                                          a_scale,
-                                          use_per_token_if_dynamic=per_act_token)
+            a_q, _ = ops.scaled_fp8_quant(
+                moe_tensors_fp16.a,
+                a_scale,
+                use_per_token_if_dynamic=per_act_token)
         else:
-            a_q, a_scale = ops.scaled_fp8_quant(moe_tensors_fp16.a,
-                                                None,
-                                                use_per_token_if_dynamic=per_act_token)
+            a_q, a_scale = ops.scaled_fp8_quant(
+                moe_tensors_fp16.a,
+                None,
+                use_per_token_if_dynamic=per_act_token)
 
         w1_q = torch.empty((e, 2 * n, k), device="cuda", dtype=q_dtype)
         w2_q = torch.empty((e, k, n), device="cuda", dtype=q_dtype)
@@ -209,7 +211,7 @@ def run_8_bit(moe_tensors: MOETensors8Bit,
         'topk_ids': topk_ids,
         'w1_scale': moe_tensors.w1_scale,
         'w2_scale': moe_tensors.w2_scale,
-        'a1_scale': None #moe_tensors.a_scale
+        'a1_scale': None  #moe_tensors.a_scale
     }
 
     num_experts = moe_tensors.w1.size(0)
@@ -262,7 +264,8 @@ def test_cutlass_moe_8_bit_no_graph(
 
         cutlass_output = run_8_bit(mt, topk_weights, topk_ids)
 
-        # Note 5.5 only needed for larger problem sizes, 5 works ok for the rest.
+        # Note 5.5 only needed for larger problem sizes, 5 works ok for
+        # the rest.
         torch.testing.assert_close(triton_output,
                                    cutlass_output,
                                    atol=5.5e-2,

tests/kernels/moe/test_deepep_deepgemm_moe.py

@@ -21,9 +21,7 @@
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     per_token_group_quant_fp8)
 from vllm.platforms import current_platform
-from vllm.utils import cdiv
 
-from tests.kernels.quant_utils import per_block_cast_to_fp8
 from .deepep_utils import ProcessGroupInfo, parallel_launch
 from .utils import make_test_weights
 
@@ -73,7 +71,8 @@ def make_block_quant_fp8_weights(
     """
     Return weights w1q, w2q, w1_scale, w2_scale
     """
-    w1, w1q, w1_scale, w2, w2q, w2_scale = make_test_weights(e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_size)
+    w1, w1q, w1_scale, w2, w2q, w2_scale = make_test_weights(
+        e, n, k, torch.bfloat16, torch.float8_e4m3fn, block_size)
     return w1q, w2q, w1_scale, w2_scale
 
 

tests/kernels/moe/test_moe.py

@@ -17,8 +17,8 @@
 import vllm.model_executor.layers.fused_moe  # noqa
 from tests.kernels.utils import opcheck, stack_and_dev, torch_moe
 from vllm.config import VllmConfig, set_current_vllm_config
-from vllm.forward_context import set_forward_context
 from vllm.distributed.parallel_state import init_distributed_environment
+from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.fused_moe import fused_moe
 from vllm.model_executor.layers.fused_moe.fused_moe import (
     fused_topk, modular_triton_fused_moe)

tests/kernels/moe/utils.py

@@ -5,17 +5,17 @@
 import torch
 
 import vllm._custom_ops as ops
+from tests.kernels.quant_utils import (per_block_cast_to_fp8,
+                                       per_block_cast_to_int8)
 from vllm.model_executor.layers.fused_moe import fused_experts
 from vllm.model_executor.layers.fused_moe.fused_batched_moe import (
     BatchedPrepareAndFinalize, BatchedTritonExperts, NaiveBatchedExperts)
 from vllm.model_executor.layers.fused_moe.modular_kernel import (
     FusedMoEModularKernel)
 from vllm.model_executor.layers.fused_moe.utils import (
     moe_kernel_quantize_input)
-from vllm.model_executor.layers.quantization.utils.fp8_utils import (
-    per_token_group_quant_fp8)
 from vllm.utils import round_up
-from tests.kernels.quant_utils import per_block_cast_to_fp8, per_block_cast_to_int8
+
 
 def triton_moe(
     a: torch.Tensor,
@@ -70,7 +70,7 @@ def batched_moe(
             max_num_tokens=max_num_tokens,
             world_size=1,
             dp_size=1,
-            use_fp8_w8a8=quant_dtype==torch.float8_e4m3fn,
+            use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
             per_act_token_quant=per_act_token_quant,
             block_shape=block_shape,
         ),
@@ -112,14 +112,19 @@ def naive_batched_moe(
             max_num_tokens=max_num_tokens,
             dp_size=1,
             world_size=1,
-            use_fp8_w8a8=quant_dtype==torch.float8_e4m3fn,
+            use_fp8_w8a8=quant_dtype == torch.float8_e4m3fn,
             per_act_token_quant=per_act_token_quant,
             block_shape=block_shape,
         ),
     )
 
-    return fused_experts(a, w1, w2, topk_weight, topk_ids,
-                         w1_scale=w1_scale, w2_scale=w2_scale,
+    return fused_experts(a,
+                         w1,
+                         w2,
+                         topk_weight,
+                         topk_ids,
+                         w1_scale=w1_scale,
+                         w2_scale=w2_scale,
                          a1_scale=a1_scale,
                          a2_scale=a2_scale)
 
@@ -148,7 +153,8 @@ def make_quantized_test_activations(
     a_scale = None
 
     if quant_dtype is not None:
-        assert quant_dtype == torch.float8_e4m3fn or quant_dtype == torch.int8, "only fp8/int8 supported"
+        assert (quant_dtype == torch.float8_e4m3fn
+                or quant_dtype == torch.int8), "only fp8/int8 supported"
         a_q = torch.zeros_like(a, dtype=quant_dtype)
         a_scale = [None] * E
         for e in range(E):
@@ -169,7 +175,8 @@ def moe_quantize_weights(
     per_token_quant: bool,
     block_shape: Optional[list[int]],
 ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
-    assert quant_dtype == torch.float8_e4m3fn or quant_dtype == torch.int8, "only fp8/int8 supported"
+    assert (quant_dtype == torch.float8_e4m3fn
+            or quant_dtype == torch.int8), "only fp8/int8 supported"
 
     if block_shape is not None:
         assert not per_token_quant
@@ -179,9 +186,11 @@ def moe_quantize_weights(
             w, w_s = per_block_cast_to_fp8(w, block_shape)
     else:
         if quant_dtype == torch.int8:
-            w, w_s = ops.scaled_int8_quant(w, w_s, use_per_token_if_dynamic=per_token_quant)
+            w, w_s = ops.scaled_int8_quant(
+                w, w_s, use_per_token_if_dynamic=per_token_quant)
         else:
-            w, w_s = ops.scaled_fp8_quant(w, w_s, use_per_token_if_dynamic=per_token_quant)
+            w, w_s = ops.scaled_fp8_quant(
+                w, w_s, use_per_token_if_dynamic=per_token_quant)
 
     return w, w_s
 
@@ -233,6 +242,8 @@ def make_test_weights(
 ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor,
            torch.Tensor, Optional[torch.Tensor]]:
     return (
-        *make_test_weight(e, 2*n, k, in_dtype, quant_dtype, block_shape, per_act_token_quant),
-        *make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape, per_act_token_quant),
+        *make_test_weight(e, 2 * n, k, in_dtype, quant_dtype, block_shape,
+                          per_act_token_quant),
+        *make_test_weight(e, k, n, in_dtype, quant_dtype, block_shape,
+                          per_act_token_quant),
     )

tests/kernels/quant_utils.py

@@ -5,9 +5,10 @@
 
 import torch
 
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    group_broadcast)
 from vllm.platforms import current_platform
 from vllm.utils import round_up
-from vllm.model_executor.layers.quantization.utils.quant_utils import group_broadcast
 
 # Using the default value (240.0) from pytorch will cause accuracy
 # issue on dynamic quantization models. Here use 224.0 for rocm.
@@ -220,17 +221,17 @@ def native_per_token_group_quant_int8(x,
 
 DEFAULT_BLOCK_SHAPE = [128, 128]
 
+
 def per_block_cast_to_fp8(
     x: torch.Tensor,
     block_shape: list[int] = DEFAULT_BLOCK_SHAPE,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     block_m, block_n = block_shape
     assert x.dim() == 2
     m, n = x.shape
-    x_padded = torch.zeros(
-        (round_up(m, block_m), round_up(n, block_n)),
-        dtype=x.dtype,
-        device=x.device)
+    x_padded = torch.zeros((round_up(m, block_m), round_up(n, block_n)),
+                           dtype=x.dtype,
+                           device=x.device)
     x_padded[:m, :n] = x
     x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
     x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
@@ -248,10 +249,9 @@ def per_block_cast_to_int8(
     block_m, block_n = block_shape
     assert x.dim() == 2
     m, n = x.shape
-    x_padded = torch.zeros(
-        (round_up(m, block_m), round_up(n, block_n)),
-        dtype=x.dtype,
-        device=x.device)
+    x_padded = torch.zeros((round_up(m, block_m), round_up(n, block_n)),
+                           dtype=x.dtype,
+                           device=x.device)
     x_padded[:m, :n] = x
     x_view = x_padded.view(-1, block_m, x_padded.size(1) // block_n, block_n)
     x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4)
@@ -292,8 +292,6 @@ def native_batched_masked_quant_matmul(
     num_expert_tokens_cpu = num_expert_tokens_cpu.to(device="cpu")
     num_experts = num_expert_tokens.size(0)
 
-    f32 = torch.float32
-
     for e in range(num_experts):
         num_tokens = num_expert_tokens_cpu[e]
         if A.dtype.itemsize == 1 and block_shape is not None:
@@ -305,7 +303,8 @@ def native_batched_masked_quant_matmul(
             assert A_scale is not None and B_scale is not None
             A_dq = dequant(A[e], A_scale[e], block_shape, per_act_token_quant)
             B_dq = dequant(B[e], B_scale[e], block_shape, per_act_token_quant)
-            C[e, :num_tokens, :] = (A_dq[:num_tokens] @ B_dq.transpose(0, 1)).to(C.dtype)
+            C[e, :num_tokens, :] = (
+                A_dq[:num_tokens] @ B_dq.transpose(0, 1)).to(C.dtype)
         else:
             assert A_scale is None
             assert B_scale is None

vllm/_custom_ops.py

@@ -1276,7 +1276,8 @@ def scaled_fp8_quant(
             torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
     else:
         # num_token_padding not implemented for this case
-        assert (scale.numel() == 1 and num_token_padding is None), f"{scale.shape} {num_token_padding}"
+        assert (scale.numel() == 1 and num_token_padding
+                is None), f"{scale.shape} {num_token_padding}"
         torch.ops._C.static_scaled_fp8_quant(output, input, scale)
 
     return output, scale

vllm/model_executor/layers/fused_moe/batched_deep_gemm_moe.py

@@ -7,8 +7,6 @@
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
-from vllm.model_executor.layers.fused_moe.utils import (
-    _resize_cache, per_token_group_quant_fp8)
 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
 from vllm.triton_utils import tl, triton