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

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 2, 04:45 UTC

[compressed-tensors] Asymmetric support for MoE WNA16 marlin (#44025)

Signed-off-by: Brian Dellabetta <bdellabe@redhat.com>

Brian Dellabetta committed Jun 2, 2026 at 11:51 UTC 774e5523971beee46d0744f15fe141548969240f

5 files changed +126 -12

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

@@ -45,6 +45,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

45 kFp8StaticTensorSym,

46 kInt4Static,

47 kInt4Static32,

48 + kInt4Static32Asym,

49 + kInt4StaticAsym,

50 kInt8Static,

51 kMxfp4Static,

52 kMxfp8Static,

@@ -611,6 +613,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):

613 kInt4Static,

614 kInt8Static,

615 kInt4Static32,

616 + kInt4StaticAsym,

617 + kInt4Static32Asym,

618 ]

619 return weight_key in SUPPORTED_W

620

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

+18

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

32 marlin_moe_permute_scales,

33 marlin_permute_bias,

34 moe_awq_to_marlin_zero_points,

35 + moe_packed_to_marlin_zero_points,

36 )

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

38 QuantKey,

@@ -496,6 +497,23 @@ def _process_weights_marlin(

497 marlin_w2_scales

498 )

499

500 + # --- Permute zero points ---

501 + if w13_qzeros is not None and w2_qzeros is not None:

502 + w13_qzeros = moe_packed_to_marlin_zero_points(

503 + w13_qzeros,

504 + size_k=w13_qzeros.shape[1],

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

506 + num_bits=num_bits,

507 + is_a_8bit=is_a_8bit,

508 + )

509 + w2_qzeros = moe_packed_to_marlin_zero_points(

510 + w2_qzeros,

511 + size_k=w2_qzeros.shape[1],

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

513 + num_bits=num_bits,

514 + is_a_8bit=is_a_8bit,

515 + )

516 +

517 # --- Permute bias ---

518 if w13_bias is not None:

519 w13_bias_out = marlin_permute_bias(w13_bias)

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

+73 -12

@@ -29,6 +29,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso

29 )

30 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa

31 WNA16_SUPPORTED_TYPES_MAP,

32 + WNA16_ZP_SUPPORTED_TYPES_MAP,

33 )

34 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

35 get_marlin_input_dtype,

@@ -56,9 +57,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

57 super().__init__(moe)

58 self.weight_quant = weight_quant

59 self.input_quant = input_quant

59 - assert weight_quant.symmetric, (

60 - "Only symmetric quantization is supported for MoE"

61 - )

60 + self.symmetric = weight_quant.symmetric

61 # Extract properties from weight_quant

62 self.num_bits = weight_quant.num_bits

63 self.packed_factor = 32 // weight_quant.num_bits

@@ -66,7 +65,12 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

65 self.group_size = weight_quant.group_size

66 self.actorder = weight_quant.actorder

69 - self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]

68 + self.quant_type = (

69 + WNA16_SUPPORTED_TYPES_MAP[self.num_bits]

70 + if self.symmetric

71 + else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits]

72 + )

73 +

74 self.marlin_input_dtype = get_marlin_input_dtype(layer_name)

76 if self.num_bits == 4:

@@ -82,7 +86,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

86 "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now."

87 )

85 - weight_key = QuantKey(self.quant_type, scale)

89 + weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)

91 # Select WNA16 MoE backend via oracle.

92 self.wna16_backend, self.experts_cls = select_wna16_moe_backend(

@@ -103,15 +107,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

107 Get the shape of the weight based on the weight name, number of experts

108 hidden size, intermediate size per partition, number of groups for w2,

109 and number of groups for w13. Pass in num_groups_w2 and num_groups_w13

106 - for weight scales.

110 + for weight scales/zero_points.

111 """

108 - if weight_name == "w13_scale":

112 + if weight_name in ("w13_scale", "w13_zp"):

113 assert num_groups_w13 is not None, (

110 - "num_groups_w13 must be provided for weight scales"

114 + "num_groups_w13 must be provided for weight scales/zero_points"

115 )

112 - if weight_name == "w2_scale":

116 + if weight_name in ("w2_scale", "w2_zp"):

117 assert num_groups_w2 is not None, (

114 - "num_groups_w2 must be provided for weight scales"

118 + "num_groups_w2 must be provided for weight scales/zero_points"

119 )

120 w13_num_shards = 2 if self.moe.is_act_and_mul else 1

121 is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM

@@ -140,6 +144,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

144 w13_num_shards * intermediate_size_per_partition,

145 ),

146 },

147 + "w13_zp": {

148 + "Marlin": (

149 + num_experts,

150 + num_groups_w13,

151 + w13_num_shards

152 + * intermediate_size_per_partition

153 + // self.packed_factor,

154 + ),

155 + },

156 "w2_weight": {

157 "Flashinfer": (

158 num_experts,

@@ -156,6 +169,13 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

169 "Flashinfer": (num_experts, hidden_size, num_groups_w2),

170 "Marlin": (num_experts, num_groups_w2, hidden_size),

171 },

172 + "w2_zp": {

173 + "Marlin": (

174 + num_experts,

175 + num_groups_w2,

176 + hidden_size // self.packed_factor,

177 + ),

178 + },

179 }

180 backend_key = "Flashinfer" if is_flashinfer else "Marlin"

181 return shape_map[weight_name][backend_key]

@@ -263,6 +283,39 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

283 set_weight_attrs(w2_scale, extra_weight_attrs)

284 set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})

285

286 + if not self.symmetric:

287 + w13_zp = torch.nn.Parameter(

288 + torch.zeros(

289 + *self.get_weight_shape(

290 + "w13_zp",

291 + num_experts,

292 + hidden_size,

293 + intermediate_size_per_partition,

294 + num_groups_w13=num_groups_w13,

295 + ),

296 + dtype=torch.int32,

297 + ),

298 + requires_grad=False,

299 + )

300 + layer.register_parameter("w13_weight_zero_point", w13_zp)

301 + set_weight_attrs(w13_zp, extra_weight_attrs)

302 +

303 + w2_zp = torch.nn.Parameter(

304 + torch.zeros(

305 + *self.get_weight_shape(

306 + "w2_zp",

307 + num_experts,

308 + hidden_size,

309 + intermediate_size_per_partition,

310 + num_groups_w2=num_groups_w2,

311 + ),

312 + dtype=torch.int32,

313 + ),

314 + requires_grad=False,

315 + )

316 + layer.register_parameter("w2_weight_zero_point", w2_zp)

317 + set_weight_attrs(w2_zp, extra_weight_attrs)

318 +

319 w2_weight_shape = torch.nn.Parameter(

320 torch.empty(num_experts, 2), requires_grad=False

321 )

@@ -334,8 +387,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

387 w2_g_idx_processed,

388 w13_g_idx_sort_indices,

389 w2_g_idx_sort_indices,

337 - _, # w13_qzeros

338 - _, # w2_qzeros

390 + w13_qzeros,

391 + w2_qzeros,

392 w13_input_global_scale,

393 w2_input_global_scale,

394 _, # w13_bias

@@ -351,6 +404,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

404 w2_scale=layer.w2_weight_scale,

405 w13_g_idx=layer.w13_weight_g_idx,

406 w2_g_idx=layer.w2_weight_g_idx,

407 + w13_qzeros=getattr(layer, "w13_weight_zero_point", None),

408 + w2_qzeros=getattr(layer, "w2_weight_zero_point", None),

409 )

410

411 # Replace common parameters

@@ -359,6 +414,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

414 replace_parameter(layer, "w13_weight_scale", w13_scales)

415 replace_parameter(layer, "w2_weight_scale", w2_scales)

416

417 + if not self.symmetric:

418 + replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)

419 + replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)

420 +

421 # Marlin-specific parameters (not needed for Flashinfer)

422 if not is_flashinfer:

423 replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)

@@ -417,6 +476,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

476 w2_scale=layer.w2_weight_scale,

477 group_size=self.group_size,

478 num_bits=self.num_bits,

479 + w1_zp=getattr(layer, "w13_weight_zero_point", None),

480 + w2_zp=getattr(layer, "w2_weight_zero_point", None),

481 )

482

483 def apply_monolithic(

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+24

@@ -429,6 +429,30 @@ def maybe_warn_marlin_atomic_add(device, dtype):

429 )

430

431

432 +def moe_packed_to_marlin_zero_points(

433 + q_zp_packed: torch.Tensor,

434 + size_k: int,

435 + size_n: int,

436 + num_bits: int,

437 + is_a_8bit: bool = False,

438 +):

439 + """Convert compressed-tensors packed zero points to Marlin format.

440 +

441 + Unlike AWQ, compressed-tensors uses standard bit packing without

442 + interleaving, so we just unpack and apply Marlin permutation directly.

443 + """

444 + num_experts = q_zp_packed.shape[0]

445 + output = torch.empty(

446 + (num_experts, q_zp_packed.shape[1], q_zp_packed.shape[2]),

447 + device=q_zp_packed.device,

448 + dtype=q_zp_packed.dtype,

449 + )

450 + for e in range(num_experts):

451 + q_zp = unpack_cols(q_zp_packed[e], num_bits, size_k, size_n)

452 + output[e] = marlin_zero_points(q_zp, size_k, size_n, num_bits, is_a_8bit)

453 + return output

454 +

455 +

456 def maybe_warn_marlin_atomic_add_env():

457 if torch.compiler.is_dynamo_compiling():

458 return

vllm/model_executor/layers/quantization/utils/quant_utils.py

@@ -181,6 +181,13 @@ kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True)

181 kInt4Static32GroupScale = ScaleDesc(torch.float16, True, GroupShape(1, 32))

182 kInt4Static32 = QuantKey(INT4_DTYPE, scale=kInt4Static32GroupScale, symmetric=True)

183

184 +kInt4StaticAsym = QuantKey(

185 + scalar_types.uint4, scale=kInt4StaticGroupScale, symmetric=False

186 +)

187 +kInt4Static32Asym = QuantKey(

188 + scalar_types.uint4, scale=kInt4Static32GroupScale, symmetric=False

189 +)

190 +

191 kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)

192 kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)

193