Merge tag 'v0.8.2' into v0.8.2-dev

benchmarks/kernels/benchmark_machete.py

@@ -7,9 +7,10 @@ import math
 import os
 import pickle as pkl
 import time
+from collections.abc import Iterable
 from dataclasses import dataclass
 from itertools import product
-from typing import Callable, Iterable, List, Optional, Tuple
+from typing import Callable, Optional
 
 import pandas as pd
 import torch
@@ -44,7 +45,6 @@ def terse_type_name(dt):
         torch.float16: "fp16",
         torch.int8: "int8",
         torch.float8_e4m3fn: "fp8",
-        torch.bfloat16: "bf16",
         torch.float: "float",
         torch.int: "int",
     }[dt]
@@ -102,8 +102,8 @@ def quantize_and_pack(atype: torch.dtype,
     return w_ref, w_q, w_s, w_zp
 
 
-def create_bench_tensors(shape: Tuple[int, int, int], types: TypeConfig,
-                         group_size: Optional[int]) -> List[BenchmarkTensors]:
+def create_bench_tensors(shape: tuple[int, int, int], types: TypeConfig,
+                         group_size: Optional[int]) -> list[BenchmarkTensors]:
     m, n, k = shape
 
     # we want to make sure that weights don't fit into L2 cache between runs so
@@ -114,7 +114,7 @@ def create_bench_tensors(shape: Tuple[int, int, int], types: TypeConfig,
 
     a = rand_data((m, k), types.act_type, scale=5)
 
-    benchmark_tensors: List[BenchmarkTensors] = []
+    benchmark_tensors: list[BenchmarkTensors] = []
     for _ in range(num_weights):
         w = rand_data((k, n), types.act_type, scale=5)
 
@@ -258,7 +258,7 @@ def machete_create_bench_fn(bt: BenchmarkTensors,
 
     return lambda: ops.machete_mm(
         a=bt.a,
-        b_q=bt.w_q,
+        b_q=w_q,
         b_type=bt.wtype,
         b_group_scales=bt.w_g_s,
         b_group_zeros=w_g_zp,
@@ -276,7 +276,7 @@ def machete_create_bench_fn(bt: BenchmarkTensors,
 
 
 def bench_fns(label: str, sub_label: str, description: str,
-              fns: List[Callable]):
+              fns: list[Callable]):
 
     min_run_time = 1 if not NVTX_PROFILE else 0.1
     res = TBenchmark.Timer(
@@ -311,7 +311,7 @@ def bench(types: TypeConfig,
           n: int,
           label: str,
           sub_label: str,
-          sweep_schedules: bool = True) -> List[TMeasurement]:
+          sweep_schedules: bool = True) -> list[TMeasurement]:
     benchmark_tensors = create_bench_tensors((m, n, k), types, group_size)
     sub_label += f", L={len(benchmark_tensors)}"
 
@@ -414,12 +414,12 @@ def bench(types: TypeConfig,
 
 
 # runner
-def print_timers(timers: List[TMeasurement]):
+def print_timers(timers: list[TMeasurement]):
     compare = TBenchmark.Compare(timers)
     compare.print()
 
 
-def run(args, MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
+def run(args, MKNs: Iterable[tuple[int, int, int]]) -> Iterable[TMeasurement]:
     types = TypeConfig(
         act_type=args.act_type,
         weight_type=scalar_types.uint4b8 if args.group_zero_type is None \
@@ -431,7 +431,7 @@ def run(args, MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
         token_scale_type=args.token_scale_type,
     )
 
-    results: List[TMeasurement] = []
+    results: list[TMeasurement] = []
     for m, k, n in MKNs:
         timers = bench(types,
                        args.group_size,
@@ -449,8 +449,8 @@ def run(args, MKNs: Iterable[Tuple[int, int, int]]) -> Iterable[TMeasurement]:
 
 # output makers
 def make_output(
-    data: List[TMeasurement],
-    MKNs: Iterable[Tuple[int, int, int]],
+    data: list[TMeasurement],
+    MKNs: Iterable[tuple[int, int, int]],
     base_description: str,
     timestamp=None,
 ):
@@ -497,7 +497,7 @@ def run_model_bench(args):
     for i, model in enumerate(args.models):
         print(f"[{i}]  {model}")
 
-    def model_shapes(model_name: str, tp_size: int) -> List[Tuple[int, int]]:
+    def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:
         KNs = []
         for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
             KN[tp_split_dim] = KN[tp_split_dim] // tp_size

benchmarks/kernels/benchmark_marlin.py

 # SPDX-License-Identifier: Apache-2.0
 
-from typing import List
-
 import torch
 import torch.utils.benchmark as benchmark
 from benchmark_shapes import WEIGHT_SHAPES
@@ -10,6 +8,8 @@ from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.gptq_marlin_24 import (
     GPTQ_MARLIN_24_MAX_PARALLEL, GPTQ_MARLIN_24_MIN_THREAD_N,
     GPTQ_MARLIN_24_SUPPORTED_GROUP_SIZES, GPTQ_MARLIN_24_SUPPORTED_QUANT_TYPES)
+from vllm.model_executor.layers.quantization.utils.allspark_utils import (
+    ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD, ALLSPARK_SUPPORTED_QUANT_TYPES)
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     GPTQ_MARLIN_MAX_PARALLEL, GPTQ_MARLIN_MIN_THREAD_N,
     MARLIN_SUPPORTED_GROUP_SIZES, query_marlin_supported_quant_types)
@@ -18,18 +18,18 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
 from vllm.model_executor.layers.quantization.utils.marlin_utils_test_24 import (
     marlin_24_quantize)
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
-    gptq_pack, gptq_quantize_weights, sort_weights)
+    gptq_pack, gptq_quantize_weights, quantize_weights, sort_weights)
 from vllm.scalar_type import ScalarType
 from vllm.utils import FlexibleArgumentParser
 
 DEFAULT_MODELS = ["meta-llama/Llama-2-7b-hf/TP1"]
-DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
+DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
 
 ACT_ORDER_OPTS = [False, True]
 K_FULL_OPTS = [False, True]
 
 
-def bench_run(results: List[benchmark.Measurement], model: str,
+def bench_run(results: list[benchmark.Measurement], model: str,
               act_order: bool, is_k_full: bool, quant_type: ScalarType,
               group_size: int, size_m: int, size_k: int, size_n: int):
     label = "Quant Matmul"
@@ -81,6 +81,27 @@ def bench_run(results: List[benchmark.Measurement], model: str,
                                           GPTQ_MARLIN_24_MAX_PARALLEL)
     marlin_zp = torch.zeros_like(marlin_s, dtype=torch.int)
 
+    # AllSpark W8A16 quant
+    as_supported_case = (quant_type in ALLSPARK_SUPPORTED_QUANT_TYPES
+                         and group_size == -1 and not act_order and is_k_full)
+    if as_supported_case:
+        properties = torch.cuda.get_device_properties(b.device.index)
+        sm_count = properties.multi_processor_count
+        sm_version = properties.major * 10 + properties.minor
+
+        supported_arch = (sm_version >= 80 and sm_version < 90)
+        as_supported_case = as_supported_case and supported_arch
+        if supported_arch:
+            has_zp = False
+            w_ref, qw, s, zp = quantize_weights(b, quant_type, group_size,
+                                                has_zp)
+            qw = qw.to(torch.uint8)
+
+            qw_reorder, s_reorder, zp_reorder = \
+                ops.allspark_repack_weight(
+                qw, s, zp, has_zp)
+            CUBLAS_M_THRESHOLD = ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD
+
     globals = {
         # Gen params
         "quant_type": quant_type,
@@ -109,10 +130,19 @@ def bench_run(results: List[benchmark.Measurement], model: str,
         # GPTQ params
         "q_w_gptq": q_w_gptq,
         "repack_sort_indices": repack_sort_indices,
+        # AllSpark W8A16 params
+        "qw_reorder": qw_reorder if as_supported_case else None,
+        "s_reorder": s_reorder if as_supported_case else None,
+        "zp_reorder": zp_reorder if as_supported_case else None,
+        "sm_count": sm_count if as_supported_case else None,
+        "sm_version": sm_version if as_supported_case else None,
+        "CUBLAS_M_THRESHOLD":
+        CUBLAS_M_THRESHOLD if as_supported_case else None,
         # Kernels
         "gptq_marlin_gemm": ops.gptq_marlin_gemm,
         "gptq_marlin_24_gemm": ops.gptq_marlin_24_gemm,
         "gptq_marlin_repack": ops.gptq_marlin_repack,
+        "allspark_w8a16_gemm": ops.allspark_w8a16_gemm,
     }
 
     min_run_time = 1
@@ -172,13 +202,24 @@ def bench_run(results: List[benchmark.Measurement], model: str,
             description="gptq_marlin_repack",
         ).blocked_autorange(min_run_time=min_run_time))
 
+    if as_supported_case:
+        results.append(
+            benchmark.Timer(
+                stmt=
+                "output = allspark_w8a16_gemm(a, qw_reorder, s_reorder, zp_reorder, size_n, group_size, sm_count, sm_version, CUBLAS_M_THRESHOLD, False, True)",  # noqa: E501
+                globals=globals,
+                label=label,
+                sub_label=sub_label,
+                description="allspark_w8a16_gemm_fp32",
+            ).blocked_autorange(min_run_time=min_run_time))
+
 
 def main(args):
     print("Benchmarking models:")
     for i, model in enumerate(args.models):
         print(f"[{i}]  {model}")
 
-    results: List[benchmark.Measurement] = []
+    results: list[benchmark.Measurement] = []
 
     for model in args.models:
         for layer in WEIGHT_SHAPES[model]:

benchmarks/kernels/benchmark_quant.py

@@ -40,7 +40,7 @@ def main(num_tokens: int,
 
         end_time = time.perf_counter()
         if profile:
-            torch.cuda.cudart().cudaProfilerStart()
+            torch.cuda.cudart().cudaProfilerStop()
         return (end_time - start_time) / num_iters
 
     # Warmup.

benchmarks/kernels/benchmark_rope.py

 # SPDX-License-Identifier: Apache-2.0
 
 from itertools import accumulate
-from typing import List, Optional
+from typing import Optional
 
 import nvtx
 import torch
@@ -39,7 +39,7 @@ def benchmark_rope_kernels_multi_lora(
                             })
     # non-batched RoPE takes only one scaling factor, we create multiple
     # instances to simulate the same behavior
-    non_batched_ropes: List[RotaryEmbedding] = []
+    non_batched_ropes: list[RotaryEmbedding] = []
     for scaling_factor in scaling_factors:
         non_batched_ropes.append(
             get_rope(head_size, rotary_dim, max_position, base, is_neox_style,