Skip to content

[Bugfix][Gemma4] Fix quantized MoE weight loading and KV cache spec merge #39582

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
neotea wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[Bugfix][Gemma4] Fix quantized MoE weight loading and KV cache spec merge #39582

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f47faea
[Bugfix][Gemma4] Fix quantized MoE weight loading and KV cache spec m…
neotea Apr 11, 2026
118ae44
Move cache_dtype_str from MLAAttentionSpec to AttentionSpec
neotea Apr 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 40 additions & 13 deletions vllm/model_executor/models/gemma4.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,6 +168,7 @@ def __init__(
config.num_experts,
bias=False,
out_dtype=torch.float32,
quant_config=quant_config,
prefix=f"{prefix}.proj",
)

Expand Down Expand Up @@ -1342,21 +1343,34 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
expert_id,
shard_id,
) in expert_params_mapping:
# Match both:
# - Bare weights: "experts.0.down_proj" (from 3D explosion)
# - With suffix: "experts.0.down_proj.weight_scale" (2D quantized)
# weight_name has trailing dot, so check with and without it
# Match weight names in three forms:
# 1. Dot suffix: "experts.0.down_proj.weight_scale"
# 2. Bare weight: "experts.0.down_proj" (from 3D explosion)
# 3. Underscore suffix: "experts.0.down_proj_packed"
# (from compressed-tensors / AWQ quantization)
weight_name_base = weight_name.rstrip(".")
if weight_name in name:
# Has suffix (e.g., .weight_scale)
# Has dot suffix (e.g., .weight_scale)
moe_name = name.replace(weight_name, param_name)
elif name.endswith(weight_name_base):
# Bare weight (no suffix)
moe_name = name.replace(
weight_name_base, param_name.rstrip("_") + "_weight"
)
else:
continue
# Check for underscore suffix (e.g., _packed, _scale)
m_us = re.match(
re.escape(weight_name_base) + r"_(\w+)$",
name[name.find("experts."):] if "experts." in name else "",
)
if m_us:
us_suffix = "_" + m_us.group(1)
moe_name = name.replace(
weight_name_base + us_suffix,
param_name + "weight" + us_suffix,
)
else:
continue
if moe_name not in params_dict:
continue
if is_pp_missing_parameter(moe_name, self):
Expand Down Expand Up @@ -1566,22 +1580,35 @@ def _weight_iterator():
#
# No transpose needed: checkpoint orientation already
# matches FusedMoE's expected layout.
if "moe.gate_up_proj" in name and weight.dim() == 3:
#
# Handle both unquantized weights (gate_up_proj,
# down_proj) and quantized weights with suffixes
# (gate_up_proj_packed, gate_up_proj_scale,
# down_proj_packed, down_proj_scale) from
# compressed-tensors / AWQ quantization.
m_gup = re.match(r"(.*)moe\.gate_up_proj(_.*)?$", name)
if m_gup and weight.dim() == 3:
suffix = m_gup.group(2) or ""
num_experts = weight.size(0)
intermediate_size = weight.size(1) // 2
for expert_id in range(num_experts):
gate_weight = weight[expert_id, :intermediate_size, :]
up_weight = weight[expert_id, intermediate_size:, :]
base = name.replace("moe.", f"moe.experts.{expert_id}.")
yield base.replace("gate_up_proj", "gate_proj"), gate_weight
yield base.replace("gate_up_proj", "up_proj"), up_weight
prefix = name[:m_gup.end(1)]
yield (f"{prefix}moe.experts.{expert_id}"
f".gate_proj{suffix}"), gate_weight
yield (f"{prefix}moe.experts.{expert_id}"
f".up_proj{suffix}"), up_weight
continue

if "moe.down_proj" in name and weight.dim() == 3:
m_dp = re.match(r"(.*)moe\.down_proj(_.*)?$", name)
if m_dp and weight.dim() == 3:
suffix = m_dp.group(2) or ""
num_experts = weight.size(0)
for expert_id in range(num_experts):
expert_name = name.replace("moe.", f"moe.experts.{expert_id}.")
yield expert_name, weight[expert_id]
prefix = name[:m_dp.end(1)]
yield (f"{prefix}moe.experts.{expert_id}"
f".down_proj{suffix}"), weight[expert_id]
continue

# k_eq_v layers: checkpoint has k_proj but no v_proj.
Expand Down
12 changes: 10 additions & 2 deletions vllm/v1/kv_cache_interface.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -218,14 +218,18 @@ def merge(cls, specs: list[Self]) -> Self:
dtype=specs[0].dtype,
kv_quant_mode=specs[0].kv_quant_mode,
page_size_padded=specs[0].page_size_padded,
cache_dtype_str=specs[0].cache_dtype_str,

@gemini-code-assist gemini-code-assist Bot Apr 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

The FullAttentionSpec class (and its base AttentionSpec) does not define a cache_dtype_str field in its dataclass definition. Attempting to pass this keyword argument to the constructor will result in a TypeError. To fix this, you should move the cache_dtype_str field from MLAAttentionSpec up to AttentionSpec or FullAttentionSpec so it is available to all relevant attention specs during merging.

sliding_window=cls.merge_window_sizes(sliding_window),
attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
)
for spec in specs:
for f in fields(AttentionSpec):
assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
"All attention layers in the same KV cache group must have "
"the same attention spec."
"the same attention spec. "
f"Field '{f.name}': "
f"{getattr(spec, f.name)!r} != "
f"{getattr(merged_spec, f.name)!r}"
)
assert (merged_spec.sliding_window is not None) + (
merged_spec.attention_chunk_size is not None
Expand Down Expand Up @@ -413,14 +417,18 @@ def merge(cls, specs: list[Self]) -> Self:
dtype=specs[0].dtype,
kv_quant_mode=specs[0].kv_quant_mode,
page_size_padded=specs[0].page_size_padded,
cache_dtype_str=specs[0].cache_dtype_str,

@gemini-code-assist gemini-code-assist Bot Apr 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

critical

Similar to FullAttentionSpec, SinkFullAttentionSpec does not have a cache_dtype_str field. This call to the constructor will fail at runtime. Please ensure the field is defined in the class hierarchy before attempting to forward it in the merge method.

sliding_window=cls.merge_window_sizes(sliding_window),
attention_chunk_size=cls.merge_window_sizes(attention_chunk_size),
)
for spec in specs:
for f in fields(AttentionSpec):
assert getattr(spec, f.name) == getattr(merged_spec, f.name), (
"All attention layers in the same KV cache group must have "
"the same attention spec."
"the same attention spec. "
f"Field '{f.name}': "
f"{getattr(spec, f.name)!r} != "
f"{getattr(merged_spec, f.name)!r}"
)
assert (merged_spec.sliding_window is not None) + (
merged_spec.attention_chunk_size is not None
Expand Down
Loading