[GG] Support heterogeneous per-layer routed-expert widths (deepseek_v2 + EXL3)

[malaiwah]

Purpose

Support MoE checkpoints with a different number of routed experts per layer (asymmetric-width MoE) on the DeepseekV2/GLM family. Pruning/frankenmodel pipelines naturally produce layers of different widths, but the stack currently assumes one scalar n_routed_experts for every layer.

• deepseek_v2.py: an optional n_routed_experts_per_layer list in the HF config overrides the scalar per layer on a shallow config copy inside DeepseekV2MoE (index parsed from the module prefix); the expert-name mapping table is sized by the widest layer (narrower layers simply never match the higher expert ids).
• exl3.py: the rank-sliced EXL3 loader resolves experts_per_layer per layer when the checkpoint metadata carries a list (bitrate-map validation, uniform-bitrate expansion, and the MoE create_weights expert-count check).

Checkpoints without n_routed_experts_per_layer / with scalar experts_per_layer metadata behave exactly as before.

Test Plan

Functional evidence checkpoint: malaiwah/GLM-5.2-Legume-v3 — a 20-layer, 724-expert GLM-5.2 franken with per-layer widths 56..28 (n_routed_experts_per_layer in its config, per-layer experts_per_layer list in its EXL3 metadata; ships these patches in patches/).

On RTX 5090 (SM120), gilded-gnosis r25 image with exactly these two files swapped in:

1. Negative control (stock files): boot must fail on the expert-count mismatch.
2. Patched: boot + small-prompt battery (1/2/5/8/9-token prefills) + greedy generations.

Test Result

Negative control (stock r25, list config) fails as expected in the EXL3 metadata parse:

    experts = int(self.rank_sliced_metadata["experts_per_layer"])
TypeError: int() argument must be a string, a bytes-like object or a real number, not 'list'

Patched (these two files swapped into the same image):

[config] n_routed_experts_per_layer=[0, 0, 0, 56, 52, 52, 48, 40, 48, 56, 48, 48, 36, 36, 36, 36, 28, 32, 36, 36]
[config] layers=20 distinct_widths=[0, 28, 32, 36, 40, 48, 52, 56] total_experts=724
[battery] 1-token prefill OK
[battery] 2-token prefill OK
[battery] 5-token prefill OK
[battery] 8-token prefill OK
[battery] 9-token prefill OK
LEGUME-V3-BOOT-OK

All 17 MoE layers (widths 28..56) landed on Exl3MoEMethod with per-layer runtime tiers (EXL3 mixed Trellis runtime planned: tiers=((3, 51), (4, 1)) etc. in the boot log), prefill + decode both exercised. Generation quality is a property of the zero-training franken (documented on the model card), not of this patch; held-out Apache-2.0 logprob for this exact checkpoint is on the model card (−11.70, beats the uniform-width baseline).

🤖 Generated with Claude Code

https://claude.ai/code/session_01DmTegASpyEG2T8zB5BZYTw

Notes

• MTP/NextN draft layers are unaffected (Legume-v3 carries no MTP overlay; the per-layer override only triggers when the config list is present).
• Complements qwen3_5: load rank-sliced EXL3 (Trellis) expert checkpoints #240 (rank-sliced EXL3 loading for qwen3_5); this PR is the deepseek_v2-family analog for asymmetric widths.

Summary by CodeRabbit

• Enhancements
  • Added support for models with different numbers of routed experts across layers.
  • Improved loading and processing of quantized mixture-of-experts weights for layer-specific configurations.
  • Added validation to detect incompatible expert layouts and invalid layer settings.
  • Preserved compatibility with configurations that use a single expert count across all layers.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[coderabbitai]

Warning

Review limit reached

@malaiwah, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 42 minutes

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: e6f725e8-b651-4297-8be6-b4948ba13a7b

📥 Commits

Reviewing files that changed from the base of the PR and between e15f0fe and 66135c5.

📒 Files selected for processing (4)

• tests/models/test_deepseek_v2_heterogeneous_experts.py
• vllm/model_executor/layers/quantization/exl3.py
• vllm/model_executor/models/deepseek_v2.py
• vllm/model_executor/models/utils.py

📝 Walkthrough

Walkthrough

DeepSeek V2 now supports heterogeneous routed-expert counts per layer. EXL3 rank-sliced bitrate loading, expansion, and MoE creation resolve and validate layer-specific expert widths.

Changes

Heterogeneous expert-count support

Layer / File(s)	Summary
DeepSeek per-layer expert selection vllm/model_executor/models/deepseek_v2.py	DeepseekV2MoE extracts and validates the layer index, applies the selected routed-expert count to a copied configuration, and uses the maximum configured count for weight mapping.
EXL3 layer-aware expert resolution vllm/model_executor/layers/quantization/exl3.py	EXL3 resolves scalar or per-layer expert counts for bitrate loading, uniform rank-sliced expansion, and rank-sliced MoE width validation.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• local-inference-lab/vllm#139: Introduces the EXL3 rank-sliced MoE support extended by this change.
• local-inference-lab/vllm#190: Extends the EXL3 rank-sliced MoE backend related to this change.
• local-inference-lab/vllm#240: Covers related EXL3 rank-sliced expert checkpoint handling.

Suggested reviewers: voipmonitor, brandonmmusic-max, jeejeelee

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and concisely describes the main change: support for heterogeneous per-layer routed-expert widths in DeepSeek V2 and EXL3.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

vllm/model_executor/layers/quantization/exl3.py (1)

163-173: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Use a Google-style docstring.

Add Args: and Returns: sections for metadata and layer_index.

As per coding guidelines, Python code must use Google-style docstrings with Args:/Returns:/Raises: sections.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@vllm/model_executor/layers/quantization/exl3.py` around lines 163 - 173,
Update the _experts_per_layer docstring to Google style by adding Args sections
documenting metadata and optional layer_index, plus a Returns section describing
the resolved expert count; preserve the existing behavior and wording about
scalar, per-layer, and widest-layer resolution.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 1332-1351: Update the rank-sliced validation around the resolved
`name` and `rank_sliced_layer_bitrates` so bitrate lookup uses `name`, including
the `layer.prefix` fallback when `layer.layer_name` is empty. Require a valid
`layers.<index>` match before continuing; otherwise reject the module here with
a clear error instead of allowing the later lookup to fail.

In `@vllm/model_executor/models/deepseek_v2.py`:
- Around line 1808-1815: Update the ROCm fusion shared-expert naming logic near
the shared-expert path to derive the routed-expert count from the current weight
layer index, using that layer’s entry in n_routed_experts_per_layer when
available and the scalar count otherwise. Use the resolved per-layer count
instead of self.config.n_routed_experts when constructing the mlp.experts
shared-expert name, while preserving existing behavior for uniform
configurations.

---

Nitpick comments:
In `@vllm/model_executor/layers/quantization/exl3.py`:
- Around line 163-173: Update the _experts_per_layer docstring to Google style
by adding Args sections documenting metadata and optional layer_index, plus a
Returns section describing the resolved expert count; preserve the existing
behavior and wording about scalar, per-layer, and widest-layer resolution.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: f20affc4-2ecf-4186-8238-b0502f29fdb7

📥 Commits

Reviewing files that changed from the base of the PR and between 3003860 and e15f0fe.

📒 Files selected for processing (2)

• vllm/model_executor/layers/quantization/exl3.py
• vllm/model_executor/models/deepseek_v2.py

[malaiwah]

Test Results (automated)

Host: macOS M4 Max, CPU-only (no CUDA)

No test files are included in this PR. The changes modify vllm/model_executor/layers/quantization/exl3.py and vllm/model_executor/models/deepseek_v2.py only — no test coverage added.

---

Automated test run by @malaiwah's agent. Results are from a CPU-only environment; GPU-dependent tests may behave differently on CUDA hardware.

[malaiwah]

Review findings addressed (B5, C10, C11, C14, C15, C12)

CodeRabbit's ROCm flag was real. The fused shared-expert offset in DeepseekV2Model.load_weights used the model-level scalar self.config.n_routed_experts, but each per-layer FusedMoE is built with num_experts=<layer width>. On a heterogeneous checkpoint the shared-expert weights were mapped to mlp.experts.{global_scalar + j} — the wrong slot (or a non-existent one) for any layer whose width differs from the scalar. Silent numerical corruption, no load-time error. Fixed.

Changes

B5 (blocker) — deepseek_v2.py: the shared-expert offset is now derived from the layer's own routed-expert count via a new _layer_routed_expert_count(name, self.config) helper:

shared_expert_offset = _layer_routed_expert_count(name, self.config)
chunk_name = name.replace("mlp.shared_experts",
                          f"mlp.experts.{shared_expert_offset + j}")

The name-mapping table is sized at max(n_routed_experts_per_layer) + n_shared_experts (the upper bound), and every layer's shared slot is <layer width> + j ≤ max + n_shared − 1, so the table and the per-layer offset base are now consistent.

C10 — DeepseekV2MoE.__init__ previously no-op'd (building the layer with the global scalar width) when re.search(r"layers\.(\d+)\.", prefix) returned None. It now raises a ValueError naming the prefix.

C11 — Three divergent layer-index regexes ((?:^|\.)layers\.(\d+)(?:\.|$), layers\.(\d+)\., layers\.(\d+)\b) disagreed on a prefix ending at layers.3. Extracted one _parse_layer_index(prefix) helper in vllm/model_executor/models/utils.py (regex (?:^|\.)layers\.(\d+)(?:\.|$) — matches layers.3 at end, rejects sublayers.3) and use it at all sites: exl3 create_weights, codebook_for_prefix, _moe_prefix_is_exl3, rank_sliced_layer_bitrates, and deepseek_v2 __init__.

C14 — exl3 create_weights replaced its per-layer validation with a -1 sentinel / "accept any declared width" fallback when no layer identity was on the module, so a mismatched checkpoint could load unchecked (and the operator saw nonsense expected=-1). Now raises cannot determine layer index for <name>.

C15 — _experts_per_layer did return int(value[layer_index]) with no bounds check (bare IndexError). Added a ValueError naming the index and list length, plus a docstring stating the index is the global model layer index (dense layers included).

C12 — Added tests/models/test_deepseek_v2_heterogeneous_experts.py (CPU-only, --noconftest).

Test coverage (new file, 22 tests, all passing)

• Per-layer width selection for a narrow (width 4) and a wide (width 256) layer, with the global scalar deliberately set to 256 so a fallback would be wrong.
• ROCm-FSE shared-expert name mapping for both a narrow layer (mlp.experts.4) and a wide layer (mlp.experts.256) — the B5 regression test; also asserts the fix's offset differs from the pre-fix global-scalar offset, and that n_shared_experts>1 maps to width+j.
• Mapping-table size bounds every layer's offset (the B5 consistency check).
• Unparseable prefix raises (C10).
• Out-of-range index raises (C15, both _layer_routed_expert_count and _experts_per_layer).
• Uniform (non-list) config unchanged (scalar returned for any layer).

Verified

.venv/bin/python -m pytest tests/models/test_deepseek_v2_heterogeneous_experts.py \
  tests/quantization/test_exl3.py tests/quantization/test_exl3_prefill_plan.py \
  --noconftest -q
# 47 passed

deepseek_v2, exl3, and models.utils all import cleanly on this CPU-only host (they only lazy-load CUDA/ROCm extensions); the tests exercise the extracted pure helpers directly rather than constructing the full model. All three changed source files pass ast.parse. The existing 19 test_exl3.py tests still pass.

Not fixed / out of scope

Nothing deliberately skipped. The C14 raise path inside Exl3MoEMethod.create_weights is covered indirectly (the _parse_layer_index + _experts_per_layer helpers it now calls are unit-tested); a direct end-to-end test of create_weights requires a CUDA RoutedExperts layer and a scheduler_config arena, which isn't available on CPU-only macOS.

[malaiwah]

I know, not everything is GLM-5.2 .. I have been experimenting with toy models of other architectures with GG and found this land mine. Now that we know, it is worth fixing for diversity.