@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7c77868cdf

[Bugfix] Size and iterate w13 by shard count for non-gated MoE (#51125)

Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Matej Sirovatka <matej.sirovatka@gmail.com> Co-authored-by: Rishi Puri <riship@nvidia.com> Co-authored-by: Hsiao-Yuan Chen <littlecircle0730@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

aoshen02 committed Aug 6, 2026 at 08:24 UTC 7c77868cdf57c3f711ba971fa68d421b5b39a9c4

16 files changed +104 -80

tests/quantization/test_auto_gptq.py

@@ -65,6 +65,7 @@ def test_auto_gptq_moe_creates_zero_initialized_expert_biases():

65 method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {})

66 method.input_dtype = None

67 method.experts_cls = None

68 + method.moe = SimpleNamespace(w13_num_shards=2)

69 layer = torch.nn.Module()

71 method.create_weights(

tests/quantization/test_auto_round.py

+4 -2

@@ -8,6 +8,7 @@ Validating the configuration and printing results for manual checking.

8 Run `pytest tests/quantization/test_auto_round.py`.

9 """

11 +from types import SimpleNamespace

12 from typing import Any, cast

14 import pytest

@@ -694,7 +695,8 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(

695 INCMxfp4MoEMethod,

696 )

697

697 - method = INCMxfp4MoEMethod(moe=cast(Any, "moe-config"))

698 + expected_moe_config = SimpleNamespace(w13_num_shards=2)

699 + method = INCMxfp4MoEMethod(moe=cast(Any, expected_moe_config))

700 layer = torch.nn.Module()

701 layer._expert_routing_tables = lambda: "routing-tables"

702

@@ -723,7 +725,7 @@ def test_inc_mxfp4_moe_method_registers_weights_and_builds_kernel(

725 assert captured["quant_config_kwargs"]["w1_scale"] is layer.w13_weight_scale

726 assert captured["quant_config_kwargs"]["w2_scale"] is layer.w2_weight_scale

727 assert captured["kernel_kwargs"]["moe_quant_config"] is expected_quant_config

726 - assert captured["kernel_kwargs"]["moe_config"] == "moe-config"

728 + assert captured["kernel_kwargs"]["moe_config"] is expected_moe_config

729 assert captured["kernel_kwargs"]["experts_cls"] is expected_experts_cls

730 assert captured["kernel_kwargs"]["routing_tables"] == "routing-tables"

731 assert method.moe_kernel is expected_kernel

vllm/model_executor/layers/fused_moe/config.py

@@ -1375,6 +1375,11 @@ class FusedMoEConfig:

1375 def is_act_and_mul(self) -> bool:

1376 return self.activation.is_gated

1377

1378 + @property

1379 + def w13_num_shards(self) -> int:

1380 + """Number of shards fused into w13: gate and up for gated, up only."""

1381 + return 2 if self.is_act_and_mul else 1

1382 +

1383 @property

1384 def tp_size(self):

1385 return self.moe_parallel_config.tp_size

vllm/model_executor/layers/quantization/auto_awq.py

+7 -3

@@ -592,7 +592,9 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):

592 torch.empty(

593 num_experts,

594 hidden_size,

595 - 2 * intermediate_size_per_partition // self.quant_config.pack_factor,

595 + self.moe.w13_num_shards

596 + * intermediate_size_per_partition

597 + // self.quant_config.pack_factor,

598 dtype=torch.int32,

599 ),

600 requires_grad=False,

@@ -623,7 +625,7 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):

625 torch.empty(

626 num_experts,

627 num_groups_w13,

626 - intermediate_size_per_partition * 2,

628 + intermediate_size_per_partition * self.moe.w13_num_shards,

629 dtype=params_dtype,

630 ),

631 requires_grad=False,

@@ -644,7 +646,9 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):

646 torch.empty(

647 num_experts,

648 num_groups_w13,

647 - 2 * intermediate_size_per_partition // self.quant_config.pack_factor,

649 + self.moe.w13_num_shards

650 + * intermediate_size_per_partition

651 + // self.quant_config.pack_factor,

652 dtype=torch.int32,

653 ),

654 requires_grad=False,

vllm/model_executor/layers/quantization/auto_gptq.py

+6 -4

@@ -541,7 +541,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

541 torch.empty(

542 num_experts,

543 hidden_size // self.quant_config.pack_factor,

544 - 2 * intermediate_size_per_partition,

544 + self.moe.w13_num_shards * intermediate_size_per_partition,

545 dtype=torch.int32,

546 ),

547 requires_grad=False,

@@ -565,7 +565,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

565 torch.empty(

566 num_experts,

567 scales_size13,

568 - 2 * intermediate_size_per_partition,

568 + self.moe.w13_num_shards * intermediate_size_per_partition,

569 dtype=params_dtype,

570 ),

571 requires_grad=False,

@@ -586,7 +586,9 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

586 torch.empty(

587 num_experts,

588 scales_size13,

589 - 2 * intermediate_size_per_partition // self.quant_config.pack_factor,

589 + self.moe.w13_num_shards

590 + * intermediate_size_per_partition

591 + // self.quant_config.pack_factor,

592 dtype=torch.int32,

593 ),

594 requires_grad=False,

@@ -654,7 +656,7 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

656 w13_bias = torch.nn.Parameter(

657 torch.zeros(

658 num_experts,

657 - 2 * intermediate_size_per_partition,

659 + self.moe.w13_num_shards * intermediate_size_per_partition,

660 dtype=params_dtype,

661 ),

662 requires_grad=False,

vllm/model_executor/layers/quantization/bitsandbytes.py

+3 -3

@@ -523,7 +523,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

523 quant_ratio = calculate_quant_ratio(params_dtype)

524 # Fused gate_up_proj (column parallel)

525 w13_total_size = (

526 - hidden_size * 2 * intermediate_size_per_partition

526 + hidden_size * self.moe.w13_num_shards * intermediate_size_per_partition

527 ) // quant_ratio

528 w13_qweight = torch.nn.Parameter(

529 torch.empty(

@@ -541,10 +541,10 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

541 {

542 "num_experts": num_experts,

543 "input_dim": hidden_size,

544 - "output_dim": 2 * intermediate_size_per_partition,

544 + "output_dim": self.moe.w13_num_shards * intermediate_size_per_partition,

545 "experts_shape": (

546 num_experts,

547 - intermediate_size_per_partition * 2,

547 + intermediate_size_per_partition * self.moe.w13_num_shards,

548 hidden_size,

549 ),

550 "pack_factor": quant_ratio,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py

+2 -2

@@ -75,7 +75,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

75 w13_weight = torch.nn.Parameter(

76 torch.empty(

77 num_experts,

78 - 2 * intermediate_size_per_partition,

78 + self.moe.w13_num_shards * intermediate_size_per_partition,

79 # 2 fp4 items are packed in the input dimension

80 hidden_size // 2,

81 requires_grad=False,

@@ -102,7 +102,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

102 w13_weight_scale = torch.nn.Parameter(

103 torch.empty(

104 num_experts,

105 - 2 * intermediate_size_per_partition,

105 + self.moe.w13_num_shards * intermediate_size_per_partition,

106 # 2 fp4 items are packed in the input dimension

107 hidden_size // self.group_size,

108 dtype=torch.uint8,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_fp8.py

+3 -3

@@ -79,7 +79,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

79 w13_weight_packed = torch.nn.Parameter(

80 torch.empty(

81 num_experts,

82 - 2 * intermediate_size_per_partition,

82 + self.moe.w13_num_shards * intermediate_size_per_partition,

83 hidden_size // self.packed_factor,

84 dtype=params_dtype,

85 ),

@@ -107,7 +107,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

107 w13_weight_scale = torch.nn.Parameter(

108 torch.ones(

109 num_experts,

110 - 2 * intermediate_size_per_partition,

110 + self.moe.w13_num_shards * intermediate_size_per_partition,

111 hidden_size // self.group_size,

112 dtype=layer.orig_dtype,

113 ),

@@ -147,7 +147,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

147 w13_weight_chan_scale = torch.nn.Parameter(

148 torch.ones(

149 num_experts,

150 - 2 * intermediate_size_per_partition,

150 + self.moe.w13_num_shards * intermediate_size_per_partition,

151 dtype=torch.float32,

152 ),

153 requires_grad=False,

vllm/model_executor/layers/quantization/fp8.py

+12 -5

@@ -545,7 +545,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

545 w13_weight = torch.nn.Parameter(

546 torch.empty(

547 num_experts,

548 - 2 * intermediate_size_per_partition,

548 + self.moe.w13_num_shards * intermediate_size_per_partition,

549 hidden_size,

550 dtype=params_dtype,

551 ),

@@ -571,7 +571,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

571 w13_bias = torch.nn.Parameter(

572 torch.zeros(

573 num_experts,

574 - 2 * intermediate_size_per_partition,

574 + self.moe.w13_num_shards * intermediate_size_per_partition,

575 dtype=layer.orig_dtype,

576 ),

577 requires_grad=False,

@@ -588,13 +588,16 @@ class Fp8MoEMethod(FusedMoEMethodBase):

588 # WEIGHT_SCALES

589 if not self.block_quant:

590 # For per-tensor quant, the scales are per expert and weight.

591 - w13_scale_data = torch.ones(num_experts, 2, dtype=torch.float32)

591 + w13_scale_data = torch.ones(

592 + num_experts, self.moe.w13_num_shards, dtype=torch.float32

593 + )

594 w2_scale_data = torch.ones(num_experts, dtype=torch.float32)

595 else:

596 # For block quant, the scales are per block (typically 128x128).

597 w13_scale_data = torch.ones(

598 num_experts,

597 - 2 * ((intermediate_size_per_partition + block_n - 1) // block_n),

599 + self.moe.w13_num_shards

600 + * ((intermediate_size_per_partition + block_n - 1) // block_n),

601 (hidden_size + block_k - 1) // block_k,

602 dtype=torch.float32,

603 )

@@ -722,7 +725,11 @@ class Fp8MoEMethod(FusedMoEMethodBase):

725 if not self.block_quant:

726 shard_size = layer.intermediate_size_per_partition

727 w13, w13_scale = process_fp8_weight_tensor_strategy_moe(

725 - w13, w13_scale, shard_size, layer.local_num_experts

728 + w13,

729 + w13_scale,

730 + shard_size,

731 + layer.local_num_experts,

732 + is_act_and_mul=self.moe.is_act_and_mul,

733 )

734

735 # Shuffle weights to runtime format and setup kernel.

vllm/model_executor/layers/quantization/humming.py

+2 -2

@@ -728,10 +728,10 @@ class HummingMoEMethod(FusedMoEMethodBase):

728 # The weight names of sublayer start with the prefix "{sublayer_name}_"

729 layer.sublayer_configs = {

730 "w13": {

731 - "shape_n": intermediate_size_per_partition * 2,

731 + "shape_n": intermediate_size_per_partition * self.moe.w13_num_shards,

732 "shape_k": hidden_size,

733 "tensors_attrs": self.weight_schema.get_padded_tensors_attrs(

734 - shape_n=intermediate_size_per_partition * 2,

734 + shape_n=intermediate_size_per_partition * self.moe.w13_num_shards,

735 shape_k=hidden_size,

736 num_experts=num_experts,

737 param_dtype=params_dtype,

vllm/model_executor/layers/quantization/inc/schemes/inc_mxfp4_moe.py

+2 -2

@@ -88,7 +88,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):

88 w13_weight = torch.nn.Parameter(

89 torch.empty(

90 num_experts,

91 - 2 * intermediate_size_per_partition,

91 + self.moe.w13_num_shards * intermediate_size_per_partition,

92 hidden_size // 2,

93 dtype=torch.uint8,

94 ),

@@ -113,7 +113,7 @@ class INCMxfp4MoEMethod(FusedMoEMethodBase):

113 w13_weight_scale = torch.nn.Parameter(

114 torch.empty(

115 num_experts,

116 - 2 * intermediate_size_per_partition,

116 + self.moe.w13_num_shards * intermediate_size_per_partition,

117 hidden_size // self.group_size,

118 dtype=torch.uint8,

119 ),

vllm/model_executor/layers/quantization/moe_wna16.py

+5 -3

@@ -276,7 +276,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

276 w13_qweight = torch.nn.Parameter(

277 torch.empty(

278 num_experts,

279 - 2 * intermediate_size_per_partition,

279 + self.moe.w13_num_shards * intermediate_size_per_partition,

280 hidden_size // bit8_pack_factor,

281 dtype=torch.uint8,

282 ),

@@ -301,7 +301,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

301 w13_scales = torch.nn.Parameter(

302 torch.zeros(

303 num_experts,

304 - 2 * intermediate_size_per_partition,

304 + self.moe.w13_num_shards * intermediate_size_per_partition,

305 hidden_size // group_size,

306 dtype=params_dtype,

307 ),

@@ -326,7 +326,9 @@ class MoeWNA16Method(FusedMoEMethodBase):

326 w13_qzeros = torch.nn.Parameter(

327 torch.zeros(

328 num_experts,

329 - 2 * intermediate_size_per_partition // bit8_pack_factor,

329 + self.moe.w13_num_shards

330 + * intermediate_size_per_partition

331 + // bit8_pack_factor,

332 hidden_size // group_size,

333 dtype=torch.uint8,

334 ),

vllm/model_executor/layers/quantization/mxfp4.py

+12 -12

@@ -215,7 +215,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

215 w13_weight = torch.nn.Parameter(

216 torch.zeros(

217 num_experts,

218 - 2 * intermediate_size_per_partition,

218 + self.moe.w13_num_shards * intermediate_size_per_partition,

219 hidden_size // 2,

220 dtype=weight_dtype,

221 ),

@@ -227,7 +227,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

227 w13_weight_scale = torch.nn.Parameter(

228 torch.zeros(

229 num_experts,

230 - 2 * intermediate_size_per_partition,

230 + self.moe.w13_num_shards * intermediate_size_per_partition,

231 hidden_size // mxfp4_block,

232 dtype=scale_dtype,

233 ),

@@ -267,7 +267,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

267 w13_bias = torch.nn.Parameter(

268 torch.zeros(

269 num_experts,

270 - 2 * intermediate_size_per_partition,

270 + self.moe.w13_num_shards * intermediate_size_per_partition,

271 dtype=torch.bfloat16,

272 ),

273 requires_grad=False,

@@ -305,13 +305,13 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

305 assert (

306 w13.dim() == 3

307 and w13.shape[0] == num_experts

308 - and w13.shape[1] == intermediate_size * 2

308 + and w13.shape[1] == intermediate_size * self.moe.w13_num_shards

309 and w13.shape[2] == hidden_size // 2

310 )

311 assert (

312 w13_scale.dim() == 3

313 and w13_scale.shape[0] == num_experts

314 - and w13_scale.shape[1] == intermediate_size * 2

314 + and w13_scale.shape[1] == intermediate_size * self.moe.w13_num_shards

315 and w13_scale.shape[2] == hidden_size // sf_block_size

316 )

317 assert (

@@ -329,7 +329,7 @@ class GptOssMxfp4MoEMethod(FusedMoEMethodBase):

329 assert (

330 w13_bias.dim() == 2

331 and w13_bias.shape[0] == num_experts

332 - and w13_bias.shape[1] == intermediate_size * 2

332 + and w13_bias.shape[1] == intermediate_size * self.moe.w13_num_shards

333 )

334 if w2_bias is not None:

335 assert (

@@ -599,7 +599,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

599 w13_weight = torch.nn.Parameter(

600 torch.zeros(

601 num_experts,

602 - 2 * intermediate_size_per_partition,

602 + self.moe.w13_num_shards * intermediate_size_per_partition,

603 hidden_size // 2,

604 dtype=weight_dtype,

605 ),

@@ -611,7 +611,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

611 w13_weight_scale = torch.nn.Parameter(

612 torch.zeros(

613 num_experts,

614 - 2 * intermediate_size_per_partition,

614 + self.moe.w13_num_shards * intermediate_size_per_partition,

615 hidden_size // mxfp4_block,

616 dtype=scale_dtype,

617 ),

@@ -651,7 +651,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

651 w13_bias = torch.nn.Parameter(

652 torch.zeros(

653 num_experts,

654 - 2 * intermediate_size_per_partition,

654 + self.moe.w13_num_shards * intermediate_size_per_partition,

655 dtype=torch.bfloat16,

656 ),

657 requires_grad=False,

@@ -689,13 +689,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

689 assert (

690 w13.dim() == 3

691 and w13.shape[0] == num_experts

692 - and w13.shape[1] == intermediate_size * 2

692 + and w13.shape[1] == intermediate_size * self.moe.w13_num_shards

693 and w13.shape[2] == hidden_size // 2

694 )

695 assert (

696 w13_scale.dim() == 3

697 and w13_scale.shape[0] == num_experts

698 - and w13_scale.shape[1] == intermediate_size * 2

698 + and w13_scale.shape[1] == intermediate_size * self.moe.w13_num_shards

699 and w13_scale.shape[2] == hidden_size // sf_block_size

700 )

701 assert (

@@ -713,7 +713,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

713 assert (

714 w13_bias.dim() == 2

715 and w13_bias.shape[0] == num_experts

716 - and w13_bias.shape[1] == intermediate_size * 2

716 + and w13_bias.shape[1] == intermediate_size * self.moe.w13_num_shards

717 )

718 if w2_bias is not None:

719 assert (

vllm/model_executor/layers/quantization/online/fp8.py

+10 -12

@@ -582,12 +582,11 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):

582 hidden_size = layer.moe_config.hidden_dim_unpadded

583 intermediate_size = layer.moe_config.intermediate_size_per_partition_unpadded

584

585 - w13_half_size = layer.w13_weight.shape[1] // 2

586 - if w13_half_size > intermediate_size:

587 - layer.w13_weight[:, intermediate_size:w13_half_size, :] = 0

588 - layer.w13_weight[

589 - :, w13_half_size + intermediate_size : 2 * w13_half_size, :

590 - ] = 0

585 + w13_shard = layer.w13_weight.shape[1] // self.moe.w13_num_shards

586 + if w13_shard > intermediate_size:

587 + for shard in range(self.moe.w13_num_shards):

588 + start = shard * w13_shard + intermediate_size

589 + layer.w13_weight[:, start : (shard + 1) * w13_shard, :] = 0

590 if layer.w13_weight.shape[2] > hidden_size:

591 layer.w13_weight[:, :, hidden_size:] = 0

592

@@ -597,12 +596,11 @@ class Fp8PerBlockOnlineMoEMethod(_Fp8OnlineMoEBase):

596 layer.w2_weight[:, :, intermediate_size:] = 0

597

598 if getattr(layer, "w13_bias", None) is not None:

600 - w13_bias_half_size = layer.w13_bias.shape[1] // 2

601 - if w13_bias_half_size > intermediate_size:

602 - layer.w13_bias[:, intermediate_size:w13_bias_half_size] = 0

603 - layer.w13_bias[

604 - :, w13_bias_half_size + intermediate_size : 2 * w13_bias_half_size

605 - ] = 0

599 + w13_bias_shard = layer.w13_bias.shape[1] // self.moe.w13_num_shards

600 + if w13_bias_shard > intermediate_size:

601 + for shard in range(self.moe.w13_num_shards):

602 + start = shard * w13_bias_shard + intermediate_size

603 + layer.w13_bias[:, start : (shard + 1) * w13_bias_shard] = 0

604

605 if (

606 getattr(layer, "w2_bias", None) is not None

vllm/model_executor/layers/quantization/online/moe_base.py

+2 -2

@@ -40,7 +40,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):

40 w13_weight = torch.nn.Parameter(

41 torch.empty(

42 num_experts,

43 - 2 * intermediate_size_per_partition,

43 + self.moe.w13_num_shards * intermediate_size_per_partition,

44 hidden_size,

45 device="meta",

46 dtype=params_dtype,

@@ -69,7 +69,7 @@ class OnlineMoEMethodBase(FusedMoEMethodBase):

69 w13_bias = torch.nn.Parameter(

70 torch.zeros(

71 num_experts,

72 - 2 * intermediate_size_per_partition,

72 + self.moe.w13_num_shards * intermediate_size_per_partition,

73 device="meta",

74 dtype=layer.orig_dtype,

75 ),

vllm/model_executor/layers/quantization/quark/quark_moe.py

+28 -25

@@ -216,7 +216,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

216 w13_weight = torch.nn.Parameter(

217 torch.zeros(

218 num_experts,

219 - 2 * intermediate_size_per_partition,

219 + self.moe.w13_num_shards * intermediate_size_per_partition,

220 hidden_size,

221 dtype=params_dtype,

222 ),

@@ -243,7 +243,10 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

243 # They are combined to a single scale after weight loading.

244 if self.model_type != "gpt_oss":

245 w13_weight_scale = torch.nn.Parameter(

246 - torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False

246 + torch.ones(

247 + num_experts, self.moe.w13_num_shards, dtype=torch.float32

248 + ),

249 + requires_grad=False,

250 )

251 else:

252 # For gpt_oss, the w1(gate) & w3(up) are fused as one.

@@ -267,7 +270,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

270 w13_weight_scale = torch.nn.Parameter(

271 torch.ones(

272 num_experts,

270 - 2 * intermediate_size_per_partition,

273 + self.moe.w13_num_shards * intermediate_size_per_partition,

274 dtype=torch.float32,

275 ),

276 requires_grad=False,

@@ -306,7 +309,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

309 w13_bias = torch.nn.Parameter(

310 torch.zeros(

311 num_experts,

309 - 2 * intermediate_size_per_partition,

312 + self.moe.w13_num_shards * intermediate_size_per_partition,

313 dtype=torch.float32,

314 ),

315 requires_grad=False,

@@ -400,7 +403,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

403 # For non-gpt_oss, process w1 and w3 shards separately

404 for expert_id in range(layer.local_num_experts):

405 start = 0

403 - for shard_id in range(2):

406 + for shard_id in range(self.moe.w13_num_shards):

407 dq_weight = per_tensor_dequantize(

408 layer.w13_weight[expert_id][start : start + shard_size, :],

409 layer.w13_weight_scale[expert_id][shard_id],

@@ -560,7 +563,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

563 w13_weight = torch.nn.Parameter(

564 torch.empty(

565 num_experts,

563 - 2 * intermediate_size_per_partition,

566 + self.moe.w13_num_shards * intermediate_size_per_partition,

567 hidden_size,

568 dtype=params_dtype,

569 ),

@@ -586,7 +589,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

589 w13_weight_scale = torch.nn.Parameter(

590 torch.ones(

591 num_experts,

589 - 2 * intermediate_size_per_partition,

592 + self.moe.w13_num_shards * intermediate_size_per_partition,

593 dtype=torch.float32,

594 ),

595 requires_grad=False,

@@ -605,7 +608,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

608 else:

609 # per-tensor: one scalar per expert (two for the fused w1/w3)

610 w13_weight_scale = torch.nn.Parameter(

608 - torch.ones(num_experts, 2, dtype=torch.float32),

611 + torch.ones(num_experts, self.moe.w13_num_shards, dtype=torch.float32),

612 requires_grad=False,

613 )

614 layer.register_parameter("w13_weight_scale", w13_weight_scale)

@@ -644,7 +647,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

647

648 # ZERO POINTS (loaded but discarded after loading; kernel uses symmetric)

649 w13_input_zero_point = torch.nn.Parameter(

647 - torch.zeros(num_experts, 2, dtype=torch.int8),

650 + torch.zeros(num_experts, self.moe.w13_num_shards, dtype=torch.int8),

651 requires_grad=False,

652 )

653 layer.register_parameter("w13_input_zero_point", w13_input_zero_point)

@@ -661,7 +664,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

664 w13_weight_zero_point = torch.nn.Parameter(

665 torch.zeros(

666 num_experts,

664 - 2 * intermediate_size_per_partition,

667 + self.moe.w13_num_shards * intermediate_size_per_partition,

668 dtype=torch.int8,

669 ),

670 requires_grad=False,

@@ -672,7 +675,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

675 )

676 else:

677 w13_weight_zero_point = torch.nn.Parameter(

675 - torch.zeros(num_experts, 2, dtype=torch.int8),

678 + torch.zeros(num_experts, self.moe.w13_num_shards, dtype=torch.int8),

679 requires_grad=False,

680 )

681 w2_weight_zero_point = torch.nn.Parameter(

@@ -689,7 +692,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

692 w13_bias = torch.nn.Parameter(

693 torch.zeros(

694 num_experts,

692 - 2 * intermediate_size_per_partition,

695 + self.moe.w13_num_shards * intermediate_size_per_partition,

696 dtype=torch.float32,

697 ),

698 requires_grad=False,

@@ -762,7 +765,7 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod):

765

766 for expert_id in range(layer.local_num_experts):

767 start = 0

765 - for shard_id in range(2):

768 + for shard_id in range(self.moe.w13_num_shards):

769 dq_weight = per_tensor_dequantize(

770 layer.w13_weight[expert_id][start : start + shard_size, :],

771 layer.w13_weight_scale[expert_id][shard_id],

@@ -905,7 +908,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

908 w13_weight = torch.nn.Parameter(

909 torch.zeros(

910 num_experts,

908 - 2 * intermediate_size_per_partition,

911 + self.moe.w13_num_shards * intermediate_size_per_partition,

912 hidden_size // 8, # INT32 packing for W4

913 dtype=params_dtype,

914 ),

@@ -927,7 +930,8 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

930

931 # Per-tensor fp8 weight scales

932 w13_weight_scale = torch.nn.Parameter(

930 - torch.ones(num_experts, 2, dtype=torch.float32), requires_grad=False

933 + torch.ones(num_experts, self.moe.w13_num_shards, dtype=torch.float32),

934 + requires_grad=False,

935 )

936 w2_weight_scale = torch.nn.Parameter(

937 torch.ones(num_experts, dtype=torch.float32), requires_grad=False

@@ -944,7 +948,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

948 w13_weight_scale_2 = torch.nn.Parameter(

949 torch.ones(

950 num_experts,

947 - 2 * intermediate_size_per_partition,

951 + self.moe.w13_num_shards * intermediate_size_per_partition,

952 dtype=torch.float32,

953 ),

954 requires_grad=False,

@@ -978,7 +982,7 @@ class QuarkW4A8Fp8MoEMethod(QuarkMoEMethod):

982 for expert_id in range(layer.local_num_experts):

983 start = 0

984 max_w13_scale_fp8 = max_w13_scales[expert_id]

981 - for shard_id in range(2):

985 + for shard_id in range(self.moe.w13_num_shards):

986 if layer.w13_weight_scale[expert_id][shard_id] != max_w13_scale_fp8:

987 int4_rescale = (

988 layer.w13_weight_scale[expert_id][shard_id] / max_w13_scale_fp8

@@ -1193,7 +1197,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1197 w13_weight = torch.nn.Parameter(

1198 torch.zeros(

1199 num_experts,

1196 - 2 * intermediate_size_per_partition,

1200 + self.moe.w13_num_shards * intermediate_size_per_partition,

1201 self.get_packed_dim(hidden_size, self.weight_dtype),

1202 dtype=params_dtype,

1203 ),

@@ -1220,7 +1224,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1224 w13_weight_scale = torch.nn.Parameter(

1225 torch.ones(

1226 num_experts,

1223 - 2 * intermediate_size_per_partition,

1227 + self.moe.w13_num_shards * intermediate_size_per_partition,

1228 hidden_size // OCP_MX_BLOCK_SIZE,

1229 dtype=params_dtype,

1230 ),

@@ -1245,7 +1249,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

1249 w13_bias = torch.nn.Parameter(

1250 torch.zeros(

1251 num_experts,

1248 - 2 * intermediate_size_per_partition,

1252 + self.moe.w13_num_shards * intermediate_size_per_partition,

1253 dtype=torch.float32,

1254 ),

1255 requires_grad=False,

@@ -1499,13 +1503,12 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):

1503 layer.quant_config = self.quant_config

1504 weight_dtype = torch.uint8

1505 weight_scale_dtype = torch.float8_e4m3fn

1502 - w13_num_shards = 2 if self.moe.is_act_and_mul else 1

1506

1507 # GEMM 1 - w13 weight

1508 w13_weight = torch.nn.Parameter(

1509 torch.empty(

1510 num_experts,

1508 - w13_num_shards * intermediate_size_per_partition,

1511 + self.moe.w13_num_shards * intermediate_size_per_partition,

1512 # 2 fp4 items are packed in the input dimension

1513 hidden_size // 2,

1514 dtype=weight_dtype,

@@ -1533,7 +1536,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):

1536 w13_weight_scale = torch.nn.Parameter(

1537 torch.empty(

1538 num_experts,

1536 - w13_num_shards * intermediate_size_per_partition,

1539 + self.moe.w13_num_shards * intermediate_size_per_partition,

1540 hidden_size // self.group_size,

1541 dtype=weight_scale_dtype,

1542 ),

@@ -1563,7 +1566,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):

1566 )

1567

1568 w13_weight_scale_2 = torch.nn.Parameter(

1566 - torch.empty(num_experts, w13_num_shards, dtype=torch.float32),

1569 + torch.empty(num_experts, self.moe.w13_num_shards, dtype=torch.float32),

1570 requires_grad=False,

1571 )

1572 layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)

@@ -1578,7 +1581,7 @@ class QuarkNvfp4MoEMethod(QuarkMoEMethod):

1581

1582 # Input global scales (per-tensor FP32 scales)

1583 w13_input_scale_2 = torch.nn.Parameter(

1581 - torch.empty(num_experts, w13_num_shards, dtype=torch.float32),

1584 + torch.empty(num_experts, self.moe.w13_num_shards, dtype=torch.float32),

1585 requires_grad=False,

1586 )

1587 layer.register_parameter("w13_input_scale_2", w13_input_scale_2)