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

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 3, 10:30 UTC

Refactor AWQ Marlin MoE onto modular WNA16 oracle (#42483)

Signed-off-by: Siddharth Bedekar <bedeksid@gmail.com> Signed-off-by: Siddharth Bedekar <104613085+bedeks@users.noreply.github.com> Co-authored-by: Robert Shaw <robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Siddharth Bedekar committed May 18, 2026 at 08:02 UTC 78e7a7b9b0b9c285bf6978c3fc09eeecea3ff230

6 files changed +405 -206

tests/kernels/moe/test_moe.py

+79 -13

@@ -1338,36 +1338,96 @@ def test_cpu_fused_moe_basic(

1338 torch.testing.assert_close(out, ref, atol=atol, rtol=0)

1339

1340

1341 -@pytest.mark.parametrize("m", [16, 32, 64])

1342 -@pytest.mark.parametrize("n", [128])

1343 -@pytest.mark.parametrize("k", [128])

1344 -@pytest.mark.parametrize("e", [8, 12, 16, 32])

1345 -@pytest.mark.parametrize("topk", [2, 4])

1346 -@pytest.mark.parametrize("max_tokens_per_batch", [16, 32, 64])

1341 +def _batched_fused_marlin_moe_cases() -> list[Any]:

1342 + cases = [

1343 + pytest.param(

1344 + m,

1345 + 128,

1346 + 128,

1347 + e,

1348 + topk,

1349 + max_tokens_per_batch,

1350 + torch.bfloat16,

1351 + scalar_types.float4_e2m1f,

1352 + None,

1353 + 1e-3,

1354 + id=(

1355 + f"m{m}-n128-k128-e{e}-topk{topk}-max_tokens{max_tokens_per_batch}-mxfp4"

1356 + ),

1357 + )

1358 + for m in [16, 32, 64]

1359 + for e in [8, 12, 16, 32]

1360 + for topk in [2, 4]

1361 + for max_tokens_per_batch in [16, 32, 64]

1362 + ]

1363 + cases.append(

1364 + pytest.param(

1365 + 32,

1366 + 128,

1367 + 128,

1368 + 8,

1369 + 2,

1370 + 64,

1371 + torch.float16,

1372 + scalar_types.uint4,

1373 + scalar_types.int8,

1374 + 4e-2,

1375 + id="awq-int8-activation-metadata",

1376 + )

1377 + )

1378 + return cases

1379 +

1380 +

1381 +@pytest.mark.parametrize(

1382 + ("m,n,k,e,topk,max_tokens_per_batch,dtype,quant_dtype,input_type,atol"),

1383 + _batched_fused_marlin_moe_cases(),

1384 +)

1385 @pytest.mark.skipif(current_platform.is_rocm(), reason="Skip for rocm")

1386 def test_batched_fused_marlin_moe(

1349 - m: int, n: int, k: int, e: int, topk: int, max_tokens_per_batch: int

1387 + m: int,

1388 + n: int,

1389 + k: int,

1390 + e: int,

1391 + topk: int,

1392 + max_tokens_per_batch: int,

1393 + dtype: torch.dtype,

1394 + quant_dtype: ScalarType,

1395 + input_type: ScalarType | None,

1396 + atol: float,

1397 ):

1398 print(

1399 f"testing m={m}, n={n}, k={k}, e={e}, "

1400 f"topk={topk}, "

1354 - f"max_tokens_per_batch={max_tokens_per_batch}"

1401 + f"max_tokens_per_batch={max_tokens_per_batch}, "

1402 + f"dtype={dtype}, quant_dtype={quant_dtype}, input_type={input_type}"

1403 )

1404 set_random_seed(0)

1405

1358 - dtype = torch.bfloat16

1359 - quant_dtype = scalar_types.float4_e2m1f

1406 group_size = 32

1407 + if input_type == scalar_types.int8:

1408 + input_dtype = torch.int8

1409 + elif input_type == scalar_types.float8_e4m3fn:

1410 + input_dtype = torch.float8_e4m3fn

1411 + else:

1412 + input_dtype = None

1413

1414 a = torch.randn((m, k), device="cuda", dtype=dtype) / 10

1415 w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 20

1416 w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 20

1417

1418 w1_data = MarlinMoEWeightData.make(

1367 - w=w1, quant_type=quant_dtype, group_size=group_size, act_order=None

1419 + w=w1,

1420 + quant_type=quant_dtype,

1421 + group_size=group_size,

1422 + act_order=None,

1423 + input_type=input_type,

1424 )

1425 w2_data = MarlinMoEWeightData.make(

1370 - w=w2, quant_type=quant_dtype, group_size=group_size, act_order=None

1426 + w=w2,

1427 + quant_type=quant_dtype,

1428 + group_size=group_size,

1429 + act_order=None,

1430 + input_type=input_type,

1431 )

1432

1433 score = torch.randn((m, e), device="cuda", dtype=dtype)

@@ -1487,6 +1547,12 @@ def test_batched_fused_marlin_moe(

1547 "quant_type_id": quant_dtype.id,

1548 "is_k_full": True,

1549 }

1550 + if input_dtype is not None:

1551 + kwargs["input_dtype"] = input_dtype

1552 + if w1_data.a_scales_factor is not None:

1553 + kwargs["input_global_scale1"] = w1_data.a_scales_factor

1554 + if w2_data.a_scales_factor is not None:

1555 + kwargs["input_global_scale2"] = w2_data.a_scales_factor

1556

1557 # Reference

1558 fused_marlin_moe_kwargs = kwargs | {

@@ -1502,7 +1568,7 @@ def test_batched_fused_marlin_moe(

1568 pytest.skip("Cannot represent data in Batched Format.")

1569 marlin_output = br.run(a, kwargs)

1570

1505 - torch.testing.assert_close(marlin_output, ref_marlin_output, atol=1e-3, rtol=0)

1571 + torch.testing.assert_close(marlin_output, ref_marlin_output, atol=atol, rtol=0)

1572

1573

1574 @pytest.mark.parametrize("m,n,k", [(32, 1024, 1024)])

vllm/model_executor/layers/fused_moe/config.py

+7 -2

@@ -957,9 +957,14 @@ def awq_marlin_moe_quant_config(

957 group_size: int,

958 w1_bias: torch.Tensor | None = None,

959 w2_bias: torch.Tensor | None = None,

960 + a1_gscale: torch.Tensor | None = None,

961 + a2_gscale: torch.Tensor | None = None,

962 ) -> FusedMoEQuantConfig:

963 """

964 Construct a quant config for awq marlin quantization.

965 +

966 + a1_gscale / a2_gscale are optional global scales applied to activation

967 + quantization scales when Marlin runs with 8-bit activations.

968 """

969 from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape

970

@@ -977,8 +982,8 @@ def awq_marlin_moe_quant_config(

982 raise ValueError(f"Unsupported weight_bits: {weight_bits}")

983

984 return FusedMoEQuantConfig(

980 - _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape),

981 - _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape),

985 + _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a1_gscale),

986 + _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a2_gscale),

987 _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias),

988 _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias),

989 )

vllm/model_executor/layers/fused_moe/experts/marlin_moe.py

+29

@@ -401,6 +401,8 @@ def batched_fused_marlin_moe(

401 global_num_experts: int = -1,

402 activation: MoEActivation = MoEActivation.SILU,

403 expert_map: torch.Tensor | None = None,

404 + input_global_scale1: torch.Tensor | None = None,

405 + input_global_scale2: torch.Tensor | None = None,

406 global_scale1: torch.Tensor | None = None,

407 global_scale2: torch.Tensor | None = None,

408 g_idx1: torch.Tensor | None = None,

@@ -414,6 +416,7 @@ def batched_fused_marlin_moe(

416 intermediate_cache2: torch.Tensor | None = None,

417 is_k_full: bool = True,

418 output: torch.Tensor | None = None,

419 + input_dtype: torch.dtype | None = None,

420 inplace: bool = False,

421 clamp_limit: float | None = None,

422 ) -> torch.Tensor:

@@ -489,7 +492,15 @@ def batched_fused_marlin_moe(

492 topk = 1

493

494 # TODO(varun) : Choose a decent block size like in fused_marlin_moe

495 + # Tune block_size_m based on expert capacity to reduce padding overhead.

496 block_size_m = 64

497 + for b_m in [8, 16, 32, 48, 64]:

498 + if BATCH_TOKENS_MAX / b_m < 0.9:

499 + block_size_m = b_m

500 + break

501 +

502 + if input_dtype is not None and input_dtype.itemsize == 1:

503 + block_size_m = max(block_size_m, 16)

504

505 sorted_token_ids, expert_ids, num_tokens_post_padded = batched_moe_align_block_size(

506 max_tokens_per_batch=BATCH_TOKENS_MAX,

@@ -525,6 +536,8 @@ def batched_fused_marlin_moe(

536 sorted_token_ids=sorted_token_ids,

537 expert_ids=expert_ids,

538 num_tokens_post_padded=num_tokens_post_padded,

539 + input_global_scale1=input_global_scale1,

540 + input_global_scale2=input_global_scale2,

541 global_scale1=global_scale1,

542 global_scale2=global_scale2,

543 g_idx1=g_idx1,

@@ -537,6 +550,7 @@ def batched_fused_marlin_moe(

550 intermediate_cache13=intermediate_cache13,

551 intermediate_cache2=intermediate_cache2,

552 output=output.view(-1, K) if output is not None else output,

553 + input_dtype=input_dtype,

554 is_k_full=is_k_full,

555 clamp_limit=clamp_limit,

556 )

@@ -635,6 +649,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):

649 @property

650 def quant_type_id(self) -> int:

651 if self.quant_config.use_int4_w4a16:

652 + if self.w1_zp is not None or self.w2_zp is not None:

653 + return scalar_types.uint4.id

654 return scalar_types.uint4b8.id

655 elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16:

656 return scalar_types.float4_e2m1f.id

@@ -754,6 +770,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):

770 topk_ids=topk_ids,

771 global_scale1=self.g1_alphas,

772 global_scale2=self.g2_alphas,

773 + input_global_scale1=self.a1_gscale,

774 + input_global_scale2=self.a2_gscale,

775 + w1_zeros=self.w1_zp,

776 + w2_zeros=self.w2_zp,

777 quant_type_id=self.quant_type_id,

778 apply_router_weight_on_input=apply_router_weight_on_input,

779 global_num_experts=global_num_experts,

@@ -853,6 +873,10 @@ class MarlinExperts(LoRAExpertsMixin, MarlinExpertsBase):

873 topk_ids=topk_ids,

874 global_scale1=self.g1_alphas,

875 global_scale2=self.g2_alphas,

876 + input_global_scale1=self.a1_gscale,

877 + input_global_scale2=self.a2_gscale,

878 + w1_zeros=self.w1_zp,

879 + w2_zeros=self.w2_zp,

880 quant_type_id=self.quant_type_id,

881 apply_router_weight_on_input=apply_router_weight_on_input,

882 global_num_experts=global_num_experts,

@@ -967,6 +991,8 @@ class BatchedMarlinExperts(MarlinExpertsBase):

991 global_num_experts=global_num_experts,

992 activation=activation,

993 expert_map=expert_map,

994 + input_global_scale1=self.a1_gscale,

995 + input_global_scale2=self.a2_gscale,

996 output=output,

997 intermediate_cache13=workspace13,

998 intermediate_cache2=workspace2,

@@ -974,6 +1000,9 @@ class BatchedMarlinExperts(MarlinExpertsBase):

1000 g_idx2=self.w2_g_idx,

1001 sort_indices1=self.w13_g_idx_sort_indices,

1002 sort_indices2=self.w2_g_idx_sort_indices,

1003 + w1_zeros=self.w1_zp,

1004 + w2_zeros=self.w2_zp,

1005 + input_dtype=self.input_dtype,

1006 is_k_full=self.is_k_full,

1007 clamp_limit=self.gemm1_clamp_limit,

1008 )

vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

+183 -4

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

21 marlin_act_int8_process_scales,

22 marlin_moe_permute_scales,

23 marlin_permute_bias,

24 + moe_awq_to_marlin_zero_points,

25 )

26 from vllm.model_executor.layers.quantization.utils.quant_utils import (

27 QuantKey,

@@ -28,6 +29,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

30 if TYPE_CHECKING:

31 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig

32 + from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig

34 logger = init_logger(__name__)

@@ -147,7 +149,6 @@ def make_wna16_moe_kernel(

149 moe_quant_config: FusedMoEQuantConfig,

150 moe_config: FusedMoEConfig,

151 experts_cls: type[mk.FusedMoEExperts] | None,

150 - layer: torch.nn.Module,

152 is_k_full: bool,

153 w13_g_idx: torch.Tensor | None,

154 w2_g_idx: torch.Tensor | None,

@@ -220,6 +221,8 @@ def _process_weights_marlin(

221 w2_scales: torch.Tensor,

222 w13_g_idx: torch.Tensor,

223 w2_g_idx: torch.Tensor,

224 + w13_qzeros: torch.Tensor | None = None,

225 + w2_qzeros: torch.Tensor | None = None,

226 w13_bias: torch.Tensor | None = None,

227 w2_bias: torch.Tensor | None = None,

228 ) -> tuple[

@@ -231,6 +234,8 @@ def _process_weights_marlin(

234 torch.Tensor, # w2_g_idx

235 torch.Tensor, # w13_g_idx_sort_indices

236 torch.Tensor, # w2_g_idx_sort_indices

237 + torch.Tensor | None, # w13_qzeros

238 + torch.Tensor | None, # w2_qzeros

239 torch.Tensor | None, # w13_input_global_scale

240 torch.Tensor | None, # w2_input_global_scale

241 torch.Tensor | None, # w13_bias

@@ -368,6 +373,151 @@ def _process_weights_marlin(

373 w2_g_idx,

374 w13_g_idx_sort_indices,

375 w2_g_idx_sort_indices,

376 + w13_qzeros,

377 + w2_qzeros,

378 + w13_input_global_scale,

379 + w2_input_global_scale,

380 + w13_bias_out,

381 + w2_bias_out,

382 + )

383 +

384 +

385 +def _process_awq_weights_marlin(

386 + layer: torch.nn.Module,

387 + quant_config: "AWQMarlinConfig",

388 + input_dtype: torch.dtype | None,

389 + w13_qweight: torch.Tensor,

390 + w2_qweight: torch.Tensor,

391 + w13_scales: torch.Tensor,

392 + w2_scales: torch.Tensor,

393 + w13_qzeros: torch.Tensor,

394 + w2_qzeros: torch.Tensor,

395 + w13_bias: torch.Tensor | None = None,

396 + w2_bias: torch.Tensor | None = None,

397 +) -> tuple[

398 + torch.Tensor, # w13_qweight

399 + torch.Tensor, # w2_qweight

400 + torch.Tensor, # w13_scales

401 + torch.Tensor, # w2_scales

402 + torch.Tensor | None, # w13_g_idx

403 + torch.Tensor | None, # w2_g_idx

404 + torch.Tensor | None, # w13_g_idx_sort_indices

405 + torch.Tensor | None, # w2_g_idx_sort_indices

406 + torch.Tensor | None, # w13_qzeros

407 + torch.Tensor | None, # w2_qzeros

408 + torch.Tensor | None, # w13_input_global_scale

409 + torch.Tensor | None, # w2_input_global_scale

410 + torch.Tensor | None, # w13_bias

411 + torch.Tensor | None, # w2_bias

412 +]:

413 + """AWQ-specific Marlin weight post-processing.

414 +

415 + AWQ checkpoints use a different packing order than GPTQ, so they need

416 + AWQ-specific weight repacking and zero-point conversion before Marlin runs.

417 + """

418 + num_experts = w13_qweight.shape[0]

419 + device = w13_qweight.device

420 + is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1

421 + w13_input_global_scale: torch.Tensor | None = None

422 + w2_input_global_scale: torch.Tensor | None = None

423 + w13_bias_out: torch.Tensor | None = None

424 + w2_bias_out: torch.Tensor | None = None

425 +

426 + if input_dtype == torch.float8_e4m3fn:

427 + ops.marlin_int4_fp8_preprocess(

428 + w13_qweight.view(-1, w13_qweight.size(2)),

429 + w13_qzeros.view(-1, w13_qzeros.size(2)),

430 + inplace=True,

431 + )

432 + ops.marlin_int4_fp8_preprocess(

433 + w2_qweight.view(-1, w2_qweight.size(2)),

434 + w2_qzeros.view(-1, w2_qzeros.size(2)),

435 + inplace=True,

436 + )

437 + w13_scales = w13_scales.data * 512

438 + w2_scales = w2_scales.data * 512

439 +

440 + w13_g_idx_sort_indices = torch.nn.Parameter(

441 + torch.empty((num_experts, 0), dtype=torch.int32, device=device),

442 + requires_grad=False,

443 + )

444 + w2_g_idx_sort_indices = torch.nn.Parameter(

445 + torch.empty((num_experts, 0), dtype=torch.int32, device=device),

446 + requires_grad=False,

447 + )

448 +

449 + marlin_w13_qweight = ops.awq_marlin_moe_repack(

450 + w13_qweight,

451 + w13_g_idx_sort_indices,

452 + size_k=w13_qweight.shape[1],

453 + size_n=w13_qweight.shape[2] * quant_config.pack_factor,

454 + num_bits=quant_config.weight_bits,

455 + is_a_8bit=is_a_8bit,

456 + )

457 + marlin_w2_qweight = ops.awq_marlin_moe_repack(

458 + w2_qweight,

459 + w2_g_idx_sort_indices,

460 + size_k=w2_qweight.shape[1],

461 + size_n=w2_qweight.shape[2] * quant_config.pack_factor,

462 + num_bits=quant_config.weight_bits,

463 + is_a_8bit=is_a_8bit,

464 + )

465 +

466 + marlin_w13_scales = marlin_moe_permute_scales(

467 + s=w13_scales,

468 + size_k=layer.intermediate_size_per_partition,

469 + size_n=w13_scales.shape[2],

470 + group_size=quant_config.group_size,

471 + is_a_8bit=is_a_8bit,

472 + )

473 + if input_dtype == torch.int8 and layer.num_groups_w13 > 1:

474 + marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(

475 + marlin_w13_scales

476 + )

477 +

478 + marlin_w2_scales = marlin_moe_permute_scales(

479 + s=w2_scales,

480 + size_k=layer.intermediate_size_per_partition,

481 + size_n=w2_scales.shape[2],

482 + group_size=quant_config.group_size,

483 + is_a_8bit=is_a_8bit,

484 + )

485 + if input_dtype == torch.int8 and layer.num_groups_w2 > 1:

486 + marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(

487 + marlin_w2_scales

488 + )

489 +

490 + marlin_w13_qzeros = moe_awq_to_marlin_zero_points(

491 + w13_qzeros,

492 + size_k=w13_qzeros.shape[1],

493 + size_n=w13_qzeros.shape[2] * quant_config.pack_factor,

494 + num_bits=quant_config.weight_bits,

495 + is_a_8bit=is_a_8bit,

496 + )

497 + marlin_w2_qzeros = moe_awq_to_marlin_zero_points(

498 + w2_qzeros,

499 + size_k=w2_qzeros.shape[1],

500 + size_n=w2_qzeros.shape[2] * quant_config.pack_factor,

501 + num_bits=quant_config.weight_bits,

502 + is_a_8bit=is_a_8bit,

503 + )

504 +

505 + if w13_bias is not None:

506 + w13_bias_out = marlin_permute_bias(w13_bias)

507 + if w2_bias is not None:

508 + w2_bias_out = marlin_permute_bias(w2_bias)

509 +

510 + return (

511 + marlin_w13_qweight,

512 + marlin_w2_qweight,

513 + marlin_w13_scales,

514 + marlin_w2_scales,

515 + None,

516 + None,

517 + w13_g_idx_sort_indices,

518 + w2_g_idx_sort_indices,

519 + marlin_w13_qzeros,

520 + marlin_w2_qzeros,

521 w13_input_global_scale,

522 w2_input_global_scale,

523 w13_bias_out,

@@ -384,8 +534,10 @@ def convert_to_wna16_moe_kernel_format(

534 w2: torch.Tensor,

535 w13_scale: torch.Tensor,

536 w2_scale: torch.Tensor,

387 - w13_g_idx: torch.Tensor,

388 - w2_g_idx: torch.Tensor,

537 + w13_g_idx: torch.Tensor | None = None,

538 + w2_g_idx: torch.Tensor | None = None,

539 + w13_qzeros: torch.Tensor | None = None,

540 + w2_qzeros: torch.Tensor | None = None,

541 w13_bias: torch.Tensor | None = None,

542 w2_bias: torch.Tensor | None = None,

543 ) -> tuple[

@@ -397,6 +549,8 @@ def convert_to_wna16_moe_kernel_format(

549 torch.Tensor | None, # w2_g_idx

550 torch.Tensor | None, # w13_g_idx_sort_indices

551 torch.Tensor | None, # w2_g_idx_sort_indices

552 + torch.Tensor | None, # w13_qzeros

553 + torch.Tensor | None, # w2_qzeros

554 torch.Tensor | None, # w13_input_global_scale

555 torch.Tensor | None, # w2_input_global_scale

556 torch.Tensor | None, # w13_bias

@@ -420,12 +574,35 @@ def convert_to_wna16_moe_kernel_format(

574 from vllm.model_executor.layers.quantization.auto_gptq import (

575 AutoGPTQConfig,

576 )

577 + from vllm.model_executor.layers.quantization.awq_marlin import (

578 + AWQMarlinConfig,

579 + )

580 +

581 + if isinstance(quant_config, AWQMarlinConfig):

582 + if w13_qzeros is None or w2_qzeros is None:

583 + raise ValueError("AWQ Marlin MoE requires zero-point tensors.")

584 + return _process_awq_weights_marlin(

585 + layer,

586 + quant_config,

587 + input_dtype,

588 + w13,

589 + w2,

590 + w13_scale,

591 + w2_scale,

592 + w13_qzeros,

593 + w2_qzeros,

594 + w13_bias,

595 + w2_bias,

596 + )

597

598 if not isinstance(quant_config, AutoGPTQConfig):

599 raise TypeError(

426 - "Marlin WNA16 MoE backend requires AutoGPTQConfig, got "

600 + "Marlin WNA16 MoE backend requires AutoGPTQConfig or "

601 + "AWQMarlinConfig, got "

602 f"{type(quant_config).__name__}."

603 )

604 + if w13_g_idx is None or w2_g_idx is None:

605 + raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")

606 return _process_weights_marlin(

607 layer,

608 quant_config,

@@ -436,6 +613,8 @@ def convert_to_wna16_moe_kernel_format(

613 w2_scale,

614 w13_g_idx,

615 w2_g_idx,

616 + w13_qzeros,

617 + w2_qzeros,

618 w13_bias,

619 w2_bias,

620 )

vllm/model_executor/layers/quantization/auto_gptq.py

+2 -1

@@ -659,6 +659,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

659 w2_g_idx,

660 w13_g_idx_sort_indices,

661 w2_g_idx_sort_indices,

662 + _w13_qzeros,

663 + _w2_qzeros,

664 w13_input_global_scale,

665 w2_input_global_scale,

666 w13_bias,

@@ -729,7 +731,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

731 moe_quant_config=self.moe_quant_config,

732 moe_config=self.moe,

733 experts_cls=self.experts_cls,

732 - layer=layer,

734 is_k_full=self.is_k_full,

735 w13_g_idx=layer.w13_g_idx,

736 w2_g_idx=layer.w2_g_idx,

vllm/model_executor/layers/quantization/awq_marlin.py

+105 -186

@@ -9,7 +9,6 @@ from torch.nn import Parameter

9 from transformers import PretrainedConfig

11 import vllm.model_executor.layers.fused_moe # noqa

12 -from vllm import _custom_ops as ops

12 from vllm import envs

13 from vllm.logger import init_logger

14 from vllm.model_executor.kernels.linear import (

@@ -27,7 +26,11 @@ from vllm.model_executor.layers.fused_moe.config import (

26 FusedMoEConfig,

27 FusedMoEQuantConfig,

28 )

30 -from vllm.model_executor.layers.fused_moe.experts.marlin_moe import fused_marlin_moe

29 +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (

30 + convert_to_wna16_moe_kernel_format,

31 + make_wna16_moe_kernel,

32 + select_wna16_moe_backend,

33 +)

34 from vllm.model_executor.layers.linear import (

35 LinearBase,

36 LinearMethodBase,

@@ -45,14 +48,13 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

48 check_marlin_supports_layer,

49 check_moe_marlin_supports_layer,

50 get_marlin_input_dtype,

48 - marlin_act_int8_process_scales,

51 marlin_make_workspace_new,

50 - marlin_moe_permute_scales,

51 - marlin_permute_bias,

52 - moe_awq_to_marlin_zero_points,

52 verify_marlin_supported,

53 )

55 -from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped

54 +from vllm.model_executor.layers.quantization.utils.quant_utils import (

55 + is_layer_skipped,

56 + kInt4Static,

57 +)

58 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

59 from vllm.model_executor.parameter import GroupQuantScaleParameter, PackedvLLMParameter

60 from vllm.platforms import current_platform

@@ -72,6 +74,19 @@ logger = init_logger(__name__)

74 _REVERSE_AWQ_PACK_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]

77 +def _replace_or_register_parameter(

78 + layer: torch.nn.Module,

79 + name: str,

80 + value: torch.Tensor | None,

81 +) -> None:

82 + if value is None:

83 + return

84 + if hasattr(layer, name):

85 + replace_parameter(layer, name, value)

86 + else:

87 + layer.register_parameter(name, Parameter(value, requires_grad=False))

88 +

89 +

90 def _convert_awq_to_standard_format(

91 layer: torch.nn.Module,

92 w_q_name: str,

@@ -505,6 +520,9 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

520 self.quant_type = scalar_types.uint4

521 self.input_dtype = None

522 self.use_marlin = True

523 + self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(

524 + moe, kInt4Static, quant_config.weight_bits

525 + )

526

527 def create_weights(

528 self,

@@ -608,52 +626,38 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

626 layer.workspace = marlin_make_workspace_new(device, 4)

627

628 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

611 - num_experts = layer.w13_qweight.shape[0]

612 - device = layer.w13_qweight.device

613 - is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1

614 -

615 - if self.input_dtype == torch.float8_e4m3fn:

616 - ops.marlin_int4_fp8_preprocess(

617 - layer.w13_qweight.view(-1, layer.w13_qweight.size(2)),

618 - layer.w13_qzeros.view(-1, layer.w13_qzeros.size(2)),

619 - inplace=True,

620 - )

621 - ops.marlin_int4_fp8_preprocess(

622 - layer.w2_qweight.view(-1, layer.w2_qweight.size(2)),

623 - layer.w2_qzeros.view(-1, layer.w2_qzeros.size(2)),

624 - inplace=True,

625 - )

626 - layer.w13_scales.data = layer.w13_scales.data * 512

627 - layer.w2_scales.data = layer.w2_scales.data * 512

628 -

629 - layer.w13_g_idx_sort_indices = torch.nn.Parameter(

630 - torch.empty((num_experts, 0), dtype=torch.int32, device=device),

631 - requires_grad=False,

632 - )

633 - layer.w2_g_idx_sort_indices = torch.nn.Parameter(

634 - torch.empty((num_experts, 0), dtype=torch.int32, device=device),

635 - requires_grad=False,

636 - )

637 -

638 - marlin_w13_qweight = ops.awq_marlin_moe_repack(

639 - layer.w13_qweight,

640 - layer.w13_g_idx_sort_indices,

641 - size_k=layer.w13_qweight.shape[1],

642 - size_n=layer.w13_qweight.shape[2] * self.quant_config.pack_factor,

643 - num_bits=self.quant_config.weight_bits,

644 - is_a_8bit=is_a_8bit,

629 + (

630 + w13,

631 + w2,

632 + w13_scale,

633 + w2_scale,

634 + w13_g_idx,

635 + w2_g_idx,

636 + w13_g_idx_sort_indices,

637 + w2_g_idx_sort_indices,

638 + w13_qzeros,

639 + w2_qzeros,

640 + w13_input_global_scale,

641 + w2_input_global_scale,

642 + w13_bias,

643 + w2_bias,

644 + ) = convert_to_wna16_moe_kernel_format(

645 + backend=self.wna16_moe_backend,

646 + layer=layer,

647 + quant_config=self.quant_config,

648 + input_dtype=self.input_dtype,

649 + w13=layer.w13_qweight,

650 + w2=layer.w2_qweight,

651 + w13_scale=layer.w13_scales,

652 + w2_scale=layer.w2_scales,

653 + w13_qzeros=layer.w13_qzeros,

654 + w2_qzeros=layer.w2_qzeros,

655 + w13_bias=getattr(layer, "w13_bias", None),

656 + w2_bias=getattr(layer, "w2_bias", None),

657 )

646 - replace_parameter(layer, "w13_qweight", marlin_w13_qweight)

658

648 - marlin_w2_qweight = ops.awq_marlin_moe_repack(

649 - layer.w2_qweight,

650 - layer.w2_g_idx_sort_indices,

651 - size_k=layer.w2_qweight.shape[1],

652 - size_n=layer.w2_qweight.shape[2] * self.quant_config.pack_factor,

653 - num_bits=self.quant_config.weight_bits,

654 - is_a_8bit=is_a_8bit,

655 - )

656 - replace_parameter(layer, "w2_qweight", marlin_w2_qweight)

659 + replace_parameter(layer, "w13_qweight", w13)

660 + replace_parameter(layer, "w2_qweight", w2)

661

662 # The modular kernel expects w13_weight and w2_weight,

663 # but AWQ uses w13_qweight and w2_qweight

@@ -662,70 +666,46 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

666 # Alias for modular kernel

667 layer.w2_weight = layer.w2_qweight

668

665 - # Why does this take the intermediate size for size_k?

666 - marlin_w13_scales = marlin_moe_permute_scales(

667 - s=layer.w13_scales,

668 - size_k=layer.intermediate_size_per_partition,

669 - size_n=layer.w13_scales.shape[2],

670 - group_size=self.quant_config.group_size,

671 - is_a_8bit=is_a_8bit,

669 + replace_parameter(layer, "w13_scales", w13_scale)

670 + replace_parameter(layer, "w2_scales", w2_scale)

671 + _replace_or_register_parameter(

672 + layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices

673 )

673 - if self.input_dtype == torch.int8 and layer.num_groups_w13 > 1:

674 - marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(

675 - marlin_w13_scales

676 - )

677 - layer.register_parameter(

678 - "w13_input_global_scale",

679 - Parameter(w13_input_global_scale, requires_grad=False),

680 - )

681 -

682 - replace_parameter(layer, "w13_scales", marlin_w13_scales)

683 -

684 - marlin_w2_scales = marlin_moe_permute_scales(

685 - s=layer.w2_scales,

686 - size_k=layer.intermediate_size_per_partition,

687 - size_n=layer.w2_scales.shape[2],

688 - group_size=self.quant_config.group_size,

689 - is_a_8bit=is_a_8bit,

674 + _replace_or_register_parameter(

675 + layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices

676 )

691 - if self.input_dtype == torch.int8 and layer.num_groups_w2 > 1:

692 - marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(

693 - marlin_w2_scales

694 - )

695 - layer.register_parameter(

696 - "w2_input_global_scale",

697 - Parameter(w2_input_global_scale, requires_grad=False),

698 - )

699 -

700 - replace_parameter(layer, "w2_scales", marlin_w2_scales)

701 -

702 - marlin_w13_zp = moe_awq_to_marlin_zero_points(

703 - layer.w13_qzeros,

704 - size_k=layer.w13_qzeros.shape[1],

705 - size_n=layer.w13_qzeros.shape[2] * self.quant_config.pack_factor,

706 - num_bits=self.quant_config.weight_bits,

707 - is_a_8bit=is_a_8bit,

677 + _replace_or_register_parameter(layer, "w13_g_idx", w13_g_idx)

678 + _replace_or_register_parameter(layer, "w2_g_idx", w2_g_idx)

679 + _replace_or_register_parameter(layer, "w13_qzeros", w13_qzeros)

680 + _replace_or_register_parameter(layer, "w2_qzeros", w2_qzeros)

681 + _replace_or_register_parameter(

682 + layer, "w13_input_global_scale", w13_input_global_scale

683 )

709 - replace_parameter(layer, "w13_qzeros", marlin_w13_zp)

710 -

711 - marlin_w2_zp = moe_awq_to_marlin_zero_points(

712 - layer.w2_qzeros,

713 - size_k=layer.w2_qzeros.shape[1],

714 - size_n=layer.w2_qzeros.shape[2] * self.quant_config.pack_factor,

715 - num_bits=self.quant_config.weight_bits,

716 - is_a_8bit=is_a_8bit,

684 + _replace_or_register_parameter(

685 + layer, "w2_input_global_scale", w2_input_global_scale

686 )

718 - replace_parameter(layer, "w2_qzeros", marlin_w2_zp)

687 + _replace_or_register_parameter(layer, "w13_bias", w13_bias)

688 + _replace_or_register_parameter(layer, "w2_bias", w2_bias)

689 +

690 + self._setup_kernel(layer)

691

720 - if hasattr(layer, "w13_bias") and layer.w13_bias is not None:

721 - layer.w13_bias.data = marlin_permute_bias(layer.w13_bias)

692 + def _setup_kernel(self, layer: RoutedExperts) -> None:

693 + """Build the FusedMoEKernel for this layer."""

694

723 - if hasattr(layer, "w2_bias") and layer.w2_bias is not None:

724 - layer.w2_bias.data = marlin_permute_bias(layer.w2_bias)

695 + self.moe_quant_config = self.get_fused_moe_quant_config(layer)

696 + self.moe_kernel = make_wna16_moe_kernel(

697 + moe_quant_config=self.moe_quant_config,

698 + moe_config=self.moe,

699 + experts_cls=self.experts_cls,

700 + is_k_full=self.is_k_full,

701 + w13_g_idx=getattr(layer, "w13_g_idx", None),

702 + w2_g_idx=getattr(layer, "w2_g_idx", None),

703 + w13_g_idx_sort_indices=layer.w13_g_idx_sort_indices,

704 + w2_g_idx_sort_indices=layer.w2_g_idx_sort_indices,

705 + routing_tables=layer._expert_routing_tables(),

706 + )

707

726 - def get_fused_moe_quant_config(

727 - self, layer: RoutedExperts

728 - ) -> FusedMoEQuantConfig | None:

708 + def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:

709 from vllm.model_executor.layers.fused_moe.config import (

710 awq_marlin_moe_quant_config,

711 )

@@ -743,6 +723,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

723 else None,

724 w1_bias=getattr(layer, "w13_bias", None),

725 w2_bias=getattr(layer, "w2_bias", None),

726 + a1_gscale=getattr(layer, "w13_input_global_scale", None),

727 + a2_gscale=getattr(layer, "w2_input_global_scale", None),

728 )

729

730 def select_gemm_impl(

@@ -750,67 +732,11 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

732 prepare_finalize,

733 layer: RoutedExperts,

734 ):

753 - """

754 - Select the GEMM implementation for AWQ-Marlin MoE.

755 - Returns MarlinExperts configured for AWQ quantization.

756 - This is ONLY used when LoRA is enabled.

757 - Without LoRA, AWQ uses its own apply() method.

758 - """

759 - # Only use modular kernels when LoRA is enabled

760 - # Without LoRA, AWQ's own apply() method works fine and is more efficient

761 - if not self.moe.is_lora_enabled:

762 - raise NotImplementedError(

763 - "AWQ-Marlin uses its own apply() method when LoRA is not enabled. "

764 - "Modular kernels are only used for LoRA support."

765 - )

766 -

767 - from vllm.model_executor.layers.fused_moe import modular_kernel as mk

768 - from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (

769 - BatchedMarlinExperts,

770 - MarlinExperts,

735 + raise ValueError(

736 + f"{self.__class__.__name__} uses the new modular kernel "

737 + "initialization logic. This function should not be called."

738 )

739

773 - # Ensure quant config is initialized

774 - assert self.moe_quant_config is not None, (

775 - "moe_quant_config must be initialized before select_gemm_impl"

776 - )

777 -

778 - w13_g_idx = getattr(layer, "w13_g_idx", None)

779 - w2_g_idx = getattr(layer, "w2_g_idx", None)

780 - w13_g_idx_sort_indices = getattr(layer, "w13_g_idx_sort_indices", None)

781 - w2_g_idx_sort_indices = getattr(layer, "w2_g_idx_sort_indices", None)

782 -

783 - # Check if using batched expert format (for Expert Parallelism)

784 - if (

785 - prepare_finalize.activation_format

786 - == mk.FusedMoEActivationFormat.BatchedExperts

787 - ):

788 - # For batched format, use BatchedMarlinExperts

789 - max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()

790 - assert max_num_tokens_per_rank is not None

791 - return BatchedMarlinExperts(

792 - max_num_tokens=max_num_tokens_per_rank,

793 - num_dispatchers=prepare_finalize.num_dispatchers(),

794 - moe_config=self.moe,

795 - quant_config=self.moe_quant_config,

796 - w13_g_idx=w13_g_idx,

797 - w2_g_idx=w2_g_idx,

798 - w13_g_idx_sort_indices=w13_g_idx_sort_indices,

799 - w2_g_idx_sort_indices=w2_g_idx_sort_indices,

800 - is_k_full=self.is_k_full,

801 - )

802 - else:

803 - # Standard Marlin experts for AWQ

804 - return MarlinExperts(

805 - moe_config=self.moe,

806 - quant_config=self.moe_quant_config,

807 - w13_g_idx=w13_g_idx,

808 - w2_g_idx=w2_g_idx,

809 - w13_g_idx_sort_indices=w13_g_idx_sort_indices,

810 - w2_g_idx_sort_indices=w2_g_idx_sort_indices,

811 - is_k_full=self.is_k_full,

812 - )

813 -

740 def apply(

741 self,

742 layer: RoutedExperts,

@@ -820,25 +746,18 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

746 shared_experts: SharedExperts | None,

747 shared_experts_input: torch.Tensor | None,

748 ) -> torch.Tensor:

823 - return fused_marlin_moe(

824 - x,

825 - layer.w13_qweight,

826 - layer.w2_qweight,

827 - getattr(layer, "w13_bias", None),

828 - getattr(layer, "w2_bias", None),

829 - layer.w13_scales,

830 - layer.w2_scales,

831 - topk_weights,

832 - topk_ids,

833 - input_global_scale1=getattr(layer, "w13_input_global_scale", None),

834 - input_global_scale2=getattr(layer, "w2_input_global_scale", None),

835 - quant_type_id=self.quant_type.id,

836 - apply_router_weight_on_input=layer.apply_router_weight_on_input,

749 + assert not self.is_monolithic

750 + assert self.moe_kernel is not None

751 + return self.moe_kernel.apply(

752 + hidden_states=x,

753 + w1=layer.w13_qweight,

754 + w2=layer.w2_qweight,

755 + topk_weights=topk_weights,

756 + topk_ids=topk_ids,

757 + activation=layer.activation,

758 global_num_experts=layer.global_num_experts,

759 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

760 expert_map=layer.expert_map,

839 - w1_zeros=layer.w13_qzeros,

840 - w2_zeros=layer.w2_qzeros,

841 - workspace=layer.workspace,

842 - input_dtype=self.input_dtype,

843 - inplace=not self.moe.disable_inplace,

761 + shared_experts=shared_experts,

762 + shared_experts_input=shared_experts_input,

763 )