Skip to content

Fix NemotronH: Register "mlp" in the cache layer-type mappings - #47535

Merged
Cyrilvallez merged 2 commits into
mainfrom
fix-mlp-layer-type
Jul 27, 2026
Merged

Fix NemotronH: Register "mlp" in the cache layer-type mappings#47535
Cyrilvallez merged 2 commits into
mainfrom
fix-mlp-layer-type

Conversation

@qgallouedec

@qgallouedec qgallouedec commented Jul 25, 2026

Copy link
Copy Markdown
Member

CI

NemotronHConfig emits four layer types (linear_attention, moe, full_attention, mlp), but cache_utils.py only knows three ("moe" is registered as a cache-less placeholder) "mlp" was left out.

So building any cache for a Nemotron-H config with - (MLP) blocks raises KeyError: 'mlp'.

That breaks the whole original Nemotron-H family (nvidia/Nemotron-H-8B-Base-8K, nvidia/Nemotron-H-4B-Instruct-128K, pattern M-M-M-M*-...) — generate() fails before any forward, as well as NemotronHConfig()'s default layer_types.

Fix: register "mlp" next to "moe" in DYNAMIC_LAYER_TYPE_MAPPING and STATIC_LAYER_TYPE_MAPPING. From the cache's point of view they are the same thing: a feed-forward block with no per-token KV and no recurrent state.

Regression from #47118.

Closes #47534

@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@Cyrilvallez Cyrilvallez left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Indeed, thanks a lot!

@github-actions

Copy link
Copy Markdown
Contributor

CI recap

Dashboard: View test results in Grafana
Latest run: 30253557001:1
Result: success | Jobs: 16 | Tests: 175,157 | Failures: 0 | Duration: 14h 3m

@Cyrilvallez
Cyrilvallez merged commit 016abe1 into main Jul 27, 2026
112 checks passed
@Cyrilvallez
Cyrilvallez deleted the fix-mlp-layer-type branch July 27, 2026 09:32
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

KeyError: 'mlp' when building a cache for Nemotron-H — "mlp" missing from the layer-type mappings

3 participants