@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 934acddef9

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 8, 08:45 UTC

[Perf] fused_moe: add int4_w4a16 benchmark support and tuning config (#34130)

Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

Matthias Gehre committed Feb 13, 2026 at 09:14 UTC 934acddef9fa4eb1b6cc897d2e39db77385539c6

2 files changed +185 -8

benchmarks/kernels/benchmark_moe.py

+122 -8

@@ -100,13 +100,38 @@ def benchmark_config(

100 dtype: torch.dtype,

101 use_fp8_w8a8: bool,

102 use_int8_w8a16: bool,

103 + use_int4_w4a16: bool = False,

104 num_iters: int = 100,

105 block_quant_shape: list[int] = None,

106 use_deep_gemm: bool = False,

107 ) -> float:

108 init_dtype = torch.float16 if use_fp8_w8a8 else dtype

109 x = torch.randn(num_tokens, hidden_size, dtype=dtype)

109 - if use_int8_w8a16:

110 + if use_int4_w4a16:

111 + # Int4 packed weights: 2 int4 values per uint8 byte

112 + # K dimension is packed (halved)

113 + intermediate_size = shard_intermediate_size // 2 # after silu_and_mul

114 + w1 = torch.randint(

115 + 0,

116 + 255,

117 + (

118 + num_experts,

119 + shard_intermediate_size,

120 + hidden_size // 2, # int4 packing

121 + ),

122 + dtype=torch.uint8,

123 + )

124 + w2 = torch.randint(

125 + 0,

126 + 255,

127 + (

128 + num_experts,

129 + hidden_size,

130 + intermediate_size // 2, # int4 packing

131 + ),

132 + dtype=torch.uint8,

133 + )

134 + elif use_int8_w8a16:

135 w1 = torch.randint(

136 -127,

137 127,

@@ -140,7 +165,20 @@ def benchmark_config(

165 w2_scale = None

166 a1_scale = None

167 a2_scale = None

143 - if use_int8_w8a16:

168 + if use_int4_w4a16:

169 + if block_quant_shape is None:

170 + raise ValueError("block_quant_shape is required for int4_w4a16")

171 + group_size = block_quant_shape[1]

172 + # Scales shape: (E, N, K // group_size) in fp16

173 + w1_scale = torch.rand(

174 + (num_experts, shard_intermediate_size, hidden_size // group_size),

175 + dtype=dtype,

176 + )

177 + w2_scale = torch.rand(

178 + (num_experts, hidden_size, intermediate_size // group_size),

179 + dtype=dtype,

180 + )

181 + elif use_int8_w8a16:

182 w1_scale = torch.randn(

183 (num_experts, 2 * shard_intermediate_size), dtype=torch.float32

184 )

@@ -199,6 +237,7 @@ def benchmark_config(

237 a1_scale=a1_scale,

238 a2_scale=a2_scale,

239 block_shape=block_quant_shape,

240 + weight_dtype="int4" if use_int4_w4a16 else None,

241 )

242

243 deep_gemm_experts = None

@@ -481,6 +520,7 @@ class BenchmarkWorker:

520 dtype: torch.dtype,

521 use_fp8_w8a8: bool,

522 use_int8_w8a16: bool,

523 + use_int4_w4a16: bool = False,

524 block_quant_shape: list[int] = None,

525 use_deep_gemm: bool = False,

526 ) -> tuple[dict[str, int], float]:

@@ -488,7 +528,10 @@ class BenchmarkWorker:

528

529 set_random_seed(self.seed)

530 dtype_str = _get_config_dtype_str(

491 - dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8

531 + dtype,

532 + use_int8_w8a16=use_int8_w8a16,

533 + use_fp8_w8a8=use_fp8_w8a8,

534 + use_int4_w4a16=use_int4_w4a16,

535 )

536 # NOTE(woosuk): The current naming convention uses w2.shape[2], which

537 # is the intermediate size after silu_and_mul.

@@ -519,6 +562,7 @@ class BenchmarkWorker:

562 dtype,

563 use_fp8_w8a8,

564 use_int8_w8a16,

565 + use_int4_w4a16=use_int4_w4a16,

566 num_iters=100,

567 block_quant_shape=block_quant_shape,

568 use_deep_gemm=use_deep_gemm,

@@ -535,6 +579,7 @@ class BenchmarkWorker:

579 dtype: torch.dtype,

580 use_fp8_w8a8: bool,

581 use_int8_w8a16: bool,

582 + use_int4_w4a16: bool,

583 search_space: list[dict[str, int]],

584 block_quant_shape: list[int],

585 use_deep_gemm: bool,

@@ -545,7 +590,7 @@ class BenchmarkWorker:

590 best_config = None

591 best_time = float("inf")

592 if current_platform.is_rocm():

548 - is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)

593 + is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)

594 search_space = prune_rocm_search_space(

595 num_tokens,

596 shard_intermediate_size,

@@ -574,6 +619,7 @@ class BenchmarkWorker:

619 dtype,

620 use_fp8_w8a8,

621 use_int8_w8a16,

622 + use_int4_w4a16,

623 num_iters=20,

624 block_quant_shape=block_quant_shape,

625 use_deep_gemm=use_deep_gemm,

@@ -621,6 +667,7 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:

667 else {}

668 ),

669 **({"kpack": config["kpack"]} if "kpack" in config else {}),

670 + **({"SPLIT_K": config["SPLIT_K"]} if "SPLIT_K" in config else {}),

671 }

672

673

@@ -633,11 +680,15 @@ def save_configs(

680 dtype: torch.dtype,

681 use_fp8_w8a8: bool,

682 use_int8_w8a16: bool,

683 + use_int4_w4a16: bool,

684 block_quant_shape: list[int],

685 save_dir: str,

686 ) -> None:

687 dtype_str = _get_config_dtype_str(

640 - dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8

688 + dtype,

689 + use_int8_w8a16=use_int8_w8a16,

690 + use_fp8_w8a8=use_fp8_w8a8,

691 + use_int4_w4a16=use_int4_w4a16,

692 )

693

694 # NOTE(woosuk): The current naming convention uses w2.shape[2], which

@@ -739,6 +790,38 @@ def get_model_params(config):

790 return E, topk, intermediate_size, hidden_size

791

792

793 +def get_quantization_group_size(config) -> int | None:

794 + """Extract the quantization group size from the HF model config.

795 +

796 + This reads directly from the HuggingFace config object (as returned by

797 + ``get_config()``), not from vLLM's quantization config classes.

798 +

799 + Supports AWQ/GPTQ-style configs (direct 'group_size' key) and

800 + compressed-tensors configs (nested inside 'config_groups').

801 + """

802 + quantization_config = getattr(config, "quantization_config", {})

803 + if not isinstance(quantization_config, dict):

804 + return None

805 + # AWQ / GPTQ style: group_size is a top-level key

806 + gs = quantization_config.get("group_size")

807 + if gs is not None:

808 + return gs

809 + # compressed-tensors style: group_size is nested in config_groups

810 + config_groups = quantization_config.get("config_groups", {})

811 + if not isinstance(config_groups, dict):

812 + return None

813 + for group_cfg in config_groups.values():

814 + if not isinstance(group_cfg, dict):

815 + continue

816 + weights = group_cfg.get("weights", {})

817 + if not isinstance(weights, dict):

818 + continue

819 + gs = weights.get("group_size")

820 + if gs is not None:

821 + return gs

822 + return None

823 +

824 +

825 def main(args: argparse.Namespace):

826 print(args)

827

@@ -757,7 +840,20 @@ def main(args: argparse.Namespace):

840 dtype = torch.float16 if current_platform.is_rocm() else config.dtype

841 use_fp8_w8a8 = args.dtype == "fp8_w8a8"

842 use_int8_w8a16 = args.dtype == "int8_w8a16"

843 + use_int4_w4a16 = args.dtype == "int4_w4a16"

844 block_quant_shape = get_weight_block_size_safety(config)

845 + if use_int4_w4a16:

846 + group_size = get_quantization_group_size(config)

847 + if group_size is None:

848 + raise ValueError(

849 + "Could not determine group_size from model config. "

850 + "The model's quantization_config must contain a 'group_size' "

851 + "field (AWQ/GPTQ) or 'config_groups.*.weights.group_size' "

852 + "(compressed-tensors)."

853 + )

854 + # For int4_w4a16, block_shape = [0, group_size]

855 + # block_shape[0]=0 means no block quantization on N dimension

856 + block_quant_shape = [0, group_size]

857

858 if args.batch_size is None:

859 batch_sizes = [

@@ -811,8 +907,20 @@ def main(args: argparse.Namespace):

907 return ray.get(outputs)

908

909 if args.tune:

814 - is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)

815 - search_space = get_configs_compute_bound(is_fp16, block_quant_shape)

910 + # int4_w4a16 weights are uint8-packed, not fp16; treat like fp8 for

911 + # search space generation (no matrix_instr_nonkdim/kpack exploration).

912 + is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16 or use_int4_w4a16)

913 + # For int4_w4a16, the group_size constraint on BLOCK_SIZE_K does not

914 + # apply: the gptq_awq kernel handles arbitrary BLOCK_SIZE_K regardless

915 + # of group_size. Skip block_quant_shape filtering to keep the full

916 + # search space (e.g. BLOCK_SIZE_K=64 with group_size=128).

917 + tune_block_quant_shape = None if use_int4_w4a16 else block_quant_shape

918 + search_space = get_configs_compute_bound(is_fp16, tune_block_quant_shape)

919 + if use_int4_w4a16:

920 + # SPLIT_K is a required kernel constexpr for gptq_awq kernel;

921 + # only SPLIT_K=1 is used at runtime, so fix it during tuning.

922 + for cfg in search_space:

923 + cfg["SPLIT_K"] = 1

924 print(f"Start tuning over {len(search_space)} configurations...")

925 if use_deep_gemm:

926 raise ValueError(

@@ -832,6 +940,7 @@ def main(args: argparse.Namespace):

940 dtype,

941 use_fp8_w8a8,

942 use_int8_w8a16,

943 + use_int4_w4a16,

944 search_space,

945 block_quant_shape,

946 use_deep_gemm,

@@ -851,6 +960,7 @@ def main(args: argparse.Namespace):

960 dtype,

961 use_fp8_w8a8,

962 use_int8_w8a16,

963 + use_int4_w4a16,

964 block_quant_shape,

965 args.save_dir,

966 )

@@ -869,6 +979,7 @@ def main(args: argparse.Namespace):

979 dtype,

980 use_fp8_w8a8,

981 use_int8_w8a16,

982 + use_int4_w4a16,

983 block_quant_shape,

984 use_deep_gemm,

985 )

@@ -891,7 +1002,10 @@ if __name__ == "__main__":

1002 )

1003 parser.add_argument("--enable-expert-parallel", "-enable-ep", action="store_true")

1004 parser.add_argument(

894 - "--dtype", type=str, choices=["auto", "fp8_w8a8", "int8_w8a16"], default="auto"

1005 + "--dtype",

1006 + type=str,

1007 + choices=["auto", "fp8_w8a8", "int8_w8a16", "int4_w4a16"],

1008 + default="auto",

1009 )

1010 parser.add_argument("--use-deep-gemm", action="store_true")

1011 parser.add_argument(

vllm/model_executor/layers/fused_moe/configs/E=128,N=768,device_name=Radeon_8060S_Graphics,dtype=int4_w4a16.json new

+63

@@ -0,0 +1,63 @@

1 +{

2 + "triton_version": "3.6.0",

3 + "1": {

4 + "BLOCK_SIZE_M": 16,

5 + "BLOCK_SIZE_N": 16,

6 + "BLOCK_SIZE_K": 64,

7 + "GROUP_SIZE_M": 1,

8 + "SPLIT_K": 1,

9 + "num_warps": 2,

10 + "num_stages": 2,

11 + "waves_per_eu": 4

12 + },

13 + "2": {

14 + "BLOCK_SIZE_M": 16,

15 + "BLOCK_SIZE_N": 16,

16 + "BLOCK_SIZE_K": 64,

17 + "GROUP_SIZE_M": 1,

18 + "SPLIT_K": 1,

19 + "num_warps": 1,

20 + "num_stages": 2,

21 + "waves_per_eu": 4

22 + },

23 + "4": {

24 + "BLOCK_SIZE_M": 32,

25 + "BLOCK_SIZE_N": 16,

26 + "BLOCK_SIZE_K": 64,

27 + "GROUP_SIZE_M": 1,

28 + "SPLIT_K": 1,

29 + "num_warps": 2,

30 + "num_stages": 2,

31 + "waves_per_eu": 2

32 + },

33 + "8": {

34 + "BLOCK_SIZE_M": 32,

35 + "BLOCK_SIZE_N": 16,

36 + "BLOCK_SIZE_K": 64,

37 + "GROUP_SIZE_M": 1,

38 + "SPLIT_K": 1,

39 + "num_warps": 2,

40 + "num_stages": 2,

41 + "waves_per_eu": 2

42 + },

43 + "16": {

44 + "BLOCK_SIZE_M": 32,

45 + "BLOCK_SIZE_N": 16,

46 + "BLOCK_SIZE_K": 64,

47 + "GROUP_SIZE_M": 4,

48 + "SPLIT_K": 1,

49 + "num_warps": 2,

50 + "num_stages": 2,

51 + "waves_per_eu": 0

52 + },

53 + "32": {

54 + "BLOCK_SIZE_M": 32,

55 + "BLOCK_SIZE_N": 16,

56 + "BLOCK_SIZE_K": 64,

57 + "GROUP_SIZE_M": 4,

58 + "SPLIT_K": 1,

59 + "num_warps": 2,

60 + "num_stages": 2,

61 + "waves_per_eu": 0

62 + }

63 +}