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

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 15, 21:00 UTC

[core] moe fp8 block quant tuning support (#14068)

Signed-off-by: Divakar Verma <divakar.verma@amd.com>

Divakar Verma committed Mar 3, 2025 at 19:30 UTC bb5b640359cc6695cb7818a24680e226f72a4da7

2 files changed +129 -57

benchmarks/kernels/benchmark_moe.py

+67 -31

@@ -40,6 +40,7 @@ def benchmark_config(

40 use_fp8_w8a8: bool,

41 use_int8_w8a16: bool,

42 num_iters: int = 100,

43 + block_quant_shape: List[int] = None,

44 ) -> float:

45 init_dtype = torch.float16 if use_fp8_w8a8 else dtype

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

@@ -81,8 +82,24 @@ def benchmark_config(

82 dtype=torch.float32)

83 w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32)

84 if use_fp8_w8a8:

84 - w1_scale = torch.randn(num_experts, dtype=torch.float32)

85 - w2_scale = torch.randn(num_experts, dtype=torch.float32)

85 + if block_quant_shape:

86 + block_n, block_k = block_quant_shape[0], block_quant_shape[1]

87 + E = num_experts

88 + N = shard_intermediate_size // 2

89 + K = hidden_size

90 + factor_for_scale = 1e-2

91 + n_tiles_w1 = (2 * N + block_n - 1) // block_n

92 + n_tiles_w2 = (K + block_n - 1) // block_n

93 + k_tiles_w1 = (K + block_k - 1) // block_k

94 + k_tiles_w2 = (N + block_k - 1) // block_k

95 + w1_scale = torch.rand((E, n_tiles_w1, k_tiles_w1),

96 + dtype=torch.float32) * factor_for_scale

97 + w2_scale = torch.rand((E, n_tiles_w2, k_tiles_w2),

98 + dtype=torch.float32) * factor_for_scale

99 + else:

100 + w1_scale = torch.randn(num_experts, dtype=torch.float32)

101 + w2_scale = torch.randn(num_experts, dtype=torch.float32)

102 +

103 a1_scale = torch.randn(1, dtype=torch.float32)

104 a2_scale = torch.randn(1, dtype=torch.float32)

105

@@ -111,6 +128,7 @@ def benchmark_config(

128 w2_scale=w2_scale,

129 a1_scale=a1_scale,

130 a2_scale=a2_scale,

131 + block_shape=block_quant_shape,

132 )

133

134 # JIT compilation & warmup

@@ -175,7 +193,8 @@ def get_rocm_tuning_space(use_fp16):

193 return param_ranges

194

195

178 -def get_configs_compute_bound(use_fp16) -> list[dict[str, int]]:

196 +def get_configs_compute_bound(use_fp16,

197 + block_quant_shape) -> list[dict[str, int]]:

198 configs: list[BenchmarkConfig] = []

199

200 if current_platform.is_rocm():

@@ -204,17 +223,27 @@ def get_configs_compute_bound(use_fp16) -> list[dict[str, int]]:

223 for config_values in product(*values):

224 config = dict(zip(keys, config_values))

225 configs.append(config)

226 +

227 + # Remove configs that are not compatible with fp8 block quantization

228 + # BLOCK_SIZE_K must be a multiple of block_k

229 + # BLOCK_SIZE_N must be a multiple of block_n

230 + if block_quant_shape is not None and not use_fp16:

231 + block_n, block_k = block_quant_shape[0], block_quant_shape[1]

232 + for config in configs[:]:

233 + if config["BLOCK_SIZE_K"] % block_k != 0 or config[

234 + "BLOCK_SIZE_N"] % block_n != 0:

235 + configs.remove(config)

236 return configs

237

238

239 def prune_rocm_search_space(num_tokens, shard_intermediate_size, hidden_size,

211 - search_space, is_fp16):

240 + search_space, is_fp16, topk):

241 N1, K1 = shard_intermediate_size, hidden_size

242 N2, K2 = hidden_size, shard_intermediate_size // 2

214 - pruned_space_1 = prune_rocm_configs(num_tokens * 2, N1, K1, search_space,

215 - is_fp16)

216 - pruned_space_2 = prune_rocm_configs(num_tokens * 2, N2, K2, search_space,

217 - is_fp16)

243 + pruned_space_1 = prune_rocm_configs(num_tokens * topk, N1, K1,

244 + search_space, is_fp16)

245 + pruned_space_2 = prune_rocm_configs(num_tokens * topk, N2, K2,

246 + search_space, is_fp16)

247 search_space = merge_unique_dicts(pruned_space_1, pruned_space_2)

248 return search_space

249

@@ -372,6 +401,7 @@ class BenchmarkWorker:

401 use_fp8_w8a8: bool,

402 use_int8_w8a16: bool,

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

404 + block_quant_shape: list[int],

405 ) -> dict[str, int]:

406 best_config = None

407 best_time = float("inf")

@@ -380,21 +410,23 @@ class BenchmarkWorker:

410 search_space = prune_rocm_search_space(num_tokens,

411 shard_intermediate_size,

412 hidden_size, search_space,

383 - is_fp16)

413 + is_fp16, topk)

414

415 with torch.cuda.device(self.device_id):

416 for config in tqdm(search_space):

417 try:

388 - kernel_time = benchmark_config(config,

389 - num_tokens,

390 - num_experts,

391 - shard_intermediate_size,

392 - hidden_size,

393 - topk,

394 - dtype,

395 - use_fp8_w8a8,

396 - use_int8_w8a16,

397 - num_iters=20)

418 + kernel_time = benchmark_config(

419 + config,

420 + num_tokens,

421 + num_experts,

422 + shard_intermediate_size,

423 + hidden_size,

424 + topk,

425 + dtype,

426 + use_fp8_w8a8,

427 + use_int8_w8a16,

428 + num_iters=20,

429 + block_quant_shape=block_quant_shape)

430 except triton.runtime.autotuner.OutOfResources:

431 # Some configurations may be invalid and fail to compile.

432 continue

@@ -436,8 +468,8 @@ def sort_config(config: BenchmarkConfig) -> BenchmarkConfig:

468

469 def save_configs(configs: dict[int, BenchmarkConfig], num_experts: int,

470 shard_intermediate_size: int, hidden_size: int, topk: int,

439 - dtype: torch.dtype, use_fp8_w8a8: bool,

440 - use_int8_w8a16: bool) -> None:

471 + dtype: torch.dtype, use_fp8_w8a8: bool, use_int8_w8a16: bool,

472 + block_quant_shape: List[int]) -> None:

473 dtype_str = get_config_dtype_str(dtype,

474 use_int8_w8a16=use_int8_w8a16,

475 use_fp8_w8a8=use_fp8_w8a8)

@@ -445,7 +477,7 @@ def save_configs(configs: dict[int, BenchmarkConfig], num_experts: int,

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

478 # is the intermediate size after silu_and_mul.

479 filename = get_config_file_name(num_experts, shard_intermediate_size // 2,

448 - dtype_str)

480 + dtype_str, block_quant_shape)

481

482 print(f"Writing best config to {filename}...")

483 with open(filename, "w") as f:

@@ -455,7 +487,7 @@ def save_configs(configs: dict[int, BenchmarkConfig], num_experts: int,

487

488 def main(args: argparse.Namespace):

489 print(args)

458 -

490 + block_quant_shape = None

491 config = AutoConfig.from_pretrained(

492 args.model, trust_remote_code=args.trust_remote_code)

493 if config.architectures[0] == "DbrxForCausalLM":

@@ -474,6 +506,7 @@ def main(args: argparse.Namespace):

506 topk = config.num_experts_per_tok

507 intermediate_size = config.moe_intermediate_size

508 shard_intermediate_size = 2 * intermediate_size // args.tp_size

509 + block_quant_shape = config.quantization_config['weight_block_size']

510 else:

511 # Default: Mixtral.

512 E = config.num_local_experts

@@ -511,27 +544,30 @@ def main(args: argparse.Namespace):

544

545 if args.tune:

546 is_fp16 = not (use_fp8_w8a8 or use_int8_w8a16)

514 - search_space = get_configs_compute_bound(is_fp16)

547 + search_space = get_configs_compute_bound(is_fp16, block_quant_shape)

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

549

550 start = time.time()

551 configs = _distribute(

519 - "tune", [(batch_size, E, shard_intermediate_size, hidden_size,

520 - topk, dtype, use_fp8_w8a8, use_int8_w8a16, search_space)

521 - for batch_size in batch_sizes])

552 + "tune",

553 + [(batch_size, E, shard_intermediate_size, hidden_size, topk, dtype,

554 + use_fp8_w8a8, use_int8_w8a16, search_space, block_quant_shape)

555 + for batch_size in batch_sizes])

556 best_configs = {

557 M: sort_config(config)

558 for M, config in zip(batch_sizes, configs)

559 }

560 save_configs(best_configs, E, shard_intermediate_size, hidden_size,

527 - topk, dtype, use_fp8_w8a8, use_int8_w8a16)

561 + topk, dtype, use_fp8_w8a8, use_int8_w8a16,

562 + block_quant_shape)

563 end = time.time()

564 print(f"Tuning took {end - start:.2f} seconds")

565 else:

566 outputs = _distribute(

532 - "benchmark", [(batch_size, E, shard_intermediate_size, hidden_size,

533 - topk, dtype, use_fp8_w8a8, use_int8_w8a16)

534 - for batch_size in batch_sizes])

567 + "benchmark",

568 + [(batch_size, E, shard_intermediate_size, hidden_size, topk, dtype,

569 + use_fp8_w8a8, use_int8_w8a16, block_quant_shape)

570 + for batch_size in batch_sizes])

571

572 for batch_size, (config, kernel_time) in zip(batch_sizes, outputs):

573 print(f"Batch size: {batch_size}, config: {config}")

vllm/model_executor/layers/fused_moe/configs/E=256,N=256,device_name=AMD_Instinct_MI300X,dtype=fp8_w8a8,block_shape=[128,128].json

+62 -26

@@ -1,28 +1,28 @@

1 {

2 "1": {

3 "BLOCK_SIZE_M": 16,

4 - "BLOCK_SIZE_N": 32,

4 + "BLOCK_SIZE_N": 128,

5 "BLOCK_SIZE_K": 256,

6 "GROUP_SIZE_M": 1,

7 - "num_warps": 4,

7 + "num_warps": 8,

8 "num_stages": 2,

9 "waves_per_eu": 0

10 },

11 "2": {

12 - "BLOCK_SIZE_M": 32,

13 - "BLOCK_SIZE_N": 16,

12 + "BLOCK_SIZE_M": 16,

13 + "BLOCK_SIZE_N": 128,

14 "BLOCK_SIZE_K": 256,

15 "GROUP_SIZE_M": 1,

16 - "num_warps": 2,

16 + "num_warps": 8,

17 "num_stages": 2,

18 "waves_per_eu": 0

19 },

20 "4": {

21 "BLOCK_SIZE_M": 16,

22 - "BLOCK_SIZE_N": 64,

22 + "BLOCK_SIZE_N": 128,

23 "BLOCK_SIZE_K": 256,

24 "GROUP_SIZE_M": 1,

25 - "num_warps": 4,

25 + "num_warps": 8,

26 "num_stages": 2,

27 "waves_per_eu": 0

28 },

@@ -31,15 +31,15 @@

31 "BLOCK_SIZE_N": 128,

32 "BLOCK_SIZE_K": 128,

33 "GROUP_SIZE_M": 1,

34 - "num_warps": 4,

34 + "num_warps": 8,

35 "num_stages": 2,

36 "waves_per_eu": 0

37 },

38 "16": {

39 "BLOCK_SIZE_M": 16,

40 - "BLOCK_SIZE_N": 64,

40 + "BLOCK_SIZE_N": 128,

41 "BLOCK_SIZE_K": 128,

42 - "GROUP_SIZE_M": 4,

42 + "GROUP_SIZE_M": 1,

43 "num_warps": 2,

44 "num_stages": 2,

45 "waves_per_eu": 0

@@ -49,13 +49,13 @@

49 "BLOCK_SIZE_N": 128,

50 "BLOCK_SIZE_K": 128,

51 "GROUP_SIZE_M": 1,

52 - "num_warps": 4,

52 + "num_warps": 2,

53 "num_stages": 2,

54 "waves_per_eu": 0

55 },

56 "32": {

57 "BLOCK_SIZE_M": 16,

58 - "BLOCK_SIZE_N": 64,

58 + "BLOCK_SIZE_N": 128,

59 "BLOCK_SIZE_K": 128,

60 "GROUP_SIZE_M": 4,

61 "num_warps": 2,

@@ -64,7 +64,7 @@

64 },

65 "48": {

66 "BLOCK_SIZE_M": 16,

67 - "BLOCK_SIZE_N": 64,

67 + "BLOCK_SIZE_N": 128,

68 "BLOCK_SIZE_K": 128,

69 "GROUP_SIZE_M": 4,

70 "num_warps": 2,

@@ -73,7 +73,7 @@

73 },

74 "64": {

75 "BLOCK_SIZE_M": 16,

76 - "BLOCK_SIZE_N": 64,

76 + "BLOCK_SIZE_N": 128,

77 "BLOCK_SIZE_K": 128,

78 "GROUP_SIZE_M": 1,

79 "num_warps": 2,

@@ -82,46 +82,82 @@

82 },

83 "96": {

84 "BLOCK_SIZE_M": 16,

85 - "BLOCK_SIZE_N": 64,

85 + "BLOCK_SIZE_N": 128,

86 "BLOCK_SIZE_K": 128,

87 - "GROUP_SIZE_M": 4,

88 - "num_warps": 4,

87 + "GROUP_SIZE_M": 8,

88 + "num_warps": 8,

89 "num_stages": 2,

90 "waves_per_eu": 0

91 },

92 "128": {

93 "BLOCK_SIZE_M": 16,

94 - "BLOCK_SIZE_N": 64,

95 - "BLOCK_SIZE_K": 256,

96 - "GROUP_SIZE_M": 1,

97 - "num_warps": 2,

94 + "BLOCK_SIZE_N": 128,

95 + "BLOCK_SIZE_K": 128,

96 + "GROUP_SIZE_M": 4,

97 + "num_warps": 4,

98 "num_stages": 2,

99 "waves_per_eu": 0

100 },

101 "256": {

102 "BLOCK_SIZE_M": 16,

103 - "BLOCK_SIZE_N": 64,

103 + "BLOCK_SIZE_N": 128,

104 "BLOCK_SIZE_K": 128,

105 - "GROUP_SIZE_M": 4,

105 + "GROUP_SIZE_M": 8,

106 "num_warps": 4,

107 "num_stages": 2,

108 "waves_per_eu": 0

109 },

110 "512": {

111 "BLOCK_SIZE_M": 32,

112 - "BLOCK_SIZE_N": 256,

112 + "BLOCK_SIZE_N": 128,

113 "BLOCK_SIZE_K": 128,

114 "GROUP_SIZE_M": 8,

115 - "num_warps": 8,

115 + "num_warps": 4,

116 "num_stages": 2,

117 "waves_per_eu": 0

118 },

119 "1024": {

120 "BLOCK_SIZE_M": 64,

121 + "BLOCK_SIZE_N": 128,

122 + "BLOCK_SIZE_K": 128,

123 + "GROUP_SIZE_M": 8,

124 + "num_warps": 2,

125 + "num_stages": 2,

126 + "waves_per_eu": 0

127 + },

128 + "1536": {

129 + "BLOCK_SIZE_M": 64,

130 + "BLOCK_SIZE_N": 128,

131 + "BLOCK_SIZE_K": 128,

132 + "GROUP_SIZE_M": 4,

133 + "num_warps": 2,

134 + "num_stages": 2,

135 + "waves_per_eu": 0

136 + },

137 + "2048": {

138 + "BLOCK_SIZE_M": 128,

139 "BLOCK_SIZE_N": 256,

140 "BLOCK_SIZE_K": 128,

141 "GROUP_SIZE_M": 8,

124 - "num_warps": 8,

142 + "num_warps": 4,

143 + "num_stages": 2,

144 + "waves_per_eu": 0

145 + },

146 + "3072": {

147 + "BLOCK_SIZE_M": 128,

148 + "BLOCK_SIZE_N": 256,

149 + "BLOCK_SIZE_K": 128,

150 + "GROUP_SIZE_M": 8,

151 + "num_warps": 4,

152 + "num_stages": 2,

153 + "waves_per_eu": 0

154 + },

155 + "4096": {

156 + "BLOCK_SIZE_M": 128,

157 + "BLOCK_SIZE_N": 256,

158 + "BLOCK_SIZE_K": 128,

159 + "GROUP_SIZE_M": 4,

160 + "num_warps": 4,

161 "num_stages": 2,

162 "waves_per_eu": 0

163 }