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 5, 09:45 UTC
[Gemma4] Support quantized MoE (#39045)
Signed-off-by: Dipika Sikka <dipikasikka1@gmail.com>
Dipika Sikka committed
Apr 8, 2026 at 21:57 UTC
3aecdf08b4a896a92e2cbd11c3d5a83d3c09abc1
1 file changed
+34
-14
vllm/model_executor/models/gemma4.py
+34
-14
@@ -1248,21 +1248,27 @@ class Gemma4Model(nn.Module):
1248
("gate_up_proj", "up_proj", 1),
1249
]
1250
1251
- # MoE expert weight mapping: checkpoint 3D packed tensors are
1252
- # exploded in _weight_iterator to per-expert 2D weights like:
1251
+ # MoE expert weight mapping: checkpoint can have either:
1252
+ # 1. 3D packed tensors (exploded in _weight_iterator to per-expert 2D)
1253
+ # 2. Already per-expert 2D weights (if quantized)
1254
+ # Map to FusedMoE parameters:
1255
# moe.experts.{id}.gate_proj → FusedMoE w1 (shard of w13)
1256
# moe.experts.{id}.up_proj → FusedMoE w3 (shard of w13)
1257
# moe.experts.{id}.down_proj → FusedMoE w2
1256
- # We build the mapping directly since Gemma4 uses bare param
1257
- # names (no .weight suffix) unlike standard MoE checkpoints.
1258
+ #
1259
+ # Use prefix matching to handle both weights and
1260
+ # quantization scale parameters. The param_name is a prefix ending
1261
+ # in underscore, and weight_name ends with a dot, so that:
1262
+ # "experts.0.gate_proj.weight_scale" -> "experts.w13_weight_scale"
1263
+ # "experts.0.gate_proj.weight" -> "experts.w13_weight"
1264
num_experts = getattr(self.config, "num_experts", None) or 0
1265
expert_params_mapping = [
1266
# (param_name, weight_name, expert_id, shard_id)
1267
(
1262
- "experts.w13_weight"
1268
+ "experts.w13_"
1269
if proj_name in ["gate_proj", "up_proj"]
1264
- else "experts.w2_weight",
1265
- f"experts.{expert_id}.{proj_name}",
1270
+ else "experts.w2_",
1271
+ f"experts.{expert_id}.{proj_name}.",
1272
expert_id,
1273
shard_id,
1274
)
@@ -1322,9 +1328,21 @@ class Gemma4Model(nn.Module):
1328
expert_id,
1329
shard_id,
1330
) in expert_params_mapping:
1325
- if weight_name not in name:
1331
+ # Match both:
1332
+ # - Bare weights: "experts.0.down_proj" (from 3D explosion)
1333
+ # - With suffix: "experts.0.down_proj.weight_scale" (2D quantized)
1334
+ # weight_name has trailing dot, so check with and without it
1335
+ weight_name_base = weight_name.rstrip(".")
1336
+ if weight_name in name:
1337
+ # Has suffix (e.g., .weight_scale)
1338
+ moe_name = name.replace(weight_name, param_name)
1339
+ elif name.endswith(weight_name_base):
1340
+ # Bare weight (no suffix)
1341
+ moe_name = name.replace(
1342
+ weight_name_base, param_name.rstrip("_") + "_weight"
1343
+ )
1344
+ else:
1345
continue
1327
- moe_name = name.replace(weight_name, param_name)
1346
if moe_name not in params_dict:
1347
continue
1348
if is_pp_missing_parameter(moe_name, self):
@@ -1334,15 +1352,12 @@ class Gemma4Model(nn.Module):
1352
# orientation for FusedMoE after _weight_iterator:
1353
# gate/up: [I, H] → w1/w3 expects [I, H]
1354
# down: [H, I] → w2 expects [H, I]
1337
- assert loaded_weight.dim() == 2, (
1338
- f"Expected 2D expert weight for {weight_name}, "
1339
- f"got shape {loaded_weight.shape}"
1340
- )
1355
+ # Scales and other quantization params may be 1D or scalar.
1356
weight_loader = param.weight_loader
1357
weight_loader(
1358
param,
1359
loaded_weight,
1345
- weight_name + ".weight",
1360
+ moe_name, # Pass mapped name (handles both weights and scales)
1361
shard_id=shard_id,
1362
expert_id=expert_id,
1363
)
@@ -1499,6 +1514,11 @@ class Gemma4ForCausalLM(nn.Module, SupportsLoRA, SupportsPP, MixtureOfExperts):
1514
".moe.down_proj",
1515
)
1516
1517
+ # Remap individual 2D expert weights:
1518
+ # .experts.{id}.{proj} → .moe.experts.{id}.{proj}
1519
+ # (This handles per-expert 2D quantized weights)
1520
+ name = re.sub(r"\.experts\.(\d+)\.", r".moe.experts.\1.", name)
1521
+
1522
# MoE expert weights: checkpoint stores as 3D packed
1523
# tensors. Explode into per-expert 2D weights for
1524
# FusedMoE weight_loader.