[models] fix: accept Gemma4 router input ids

[yaoyu-33]

Mirrors #3820 for MB-365 and adds the Gemma4 router compatibility fix needed for the MCore dev bump. MCore dev now forwards input_ids into TopKRouter.routing(), so Gemma4TopKRouter needs to accept and pass that argument through to the parent routing implementation.

Original bump PR: #3820
Linear: MB-365

Validation:

• python3 -m py_compile src/megatron/bridge/models/gemma/gemma4_provider.py
• UV_CACHE_DIR=/tmp/uv-cache-yuya uv run --no-sync ruff check src/megatron/bridge/models/gemma/gemma4_provider.py
• cw job 11767705: Gemma/Gemma-VL unit subset passed, 128 passed, 34 warnings in 2.38s
• cw job 11767727: tests/functional_tests/launch_scripts/h100/active/L0_Launch_recipes_gemma4_vl.sh passed, 1 passed, 929 warnings in 119.55s; Slurm state COMPLETED, exit 0:0

Note: cw did not expose a GB200 partition in sinfo, so the Gemma4-VL functional rerun used the H100 active launch script for the same failing recipe group.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[claude]

Light Code Review

MCore dev bump + Gemma4 router compatibility fix. The change is minimal and correct: Gemma4TopKRouter.routing() now accepts and forwards the new input_ids kwarg that MCore TopKRouter.routing() started passing through.

No issues found. The code change is straightforward:

• input_ids=None default preserves backward compatibility.
• The kwarg is passed through to super().routing() correctly.
• LoRATopKRouter (the only other router subclass) already uses *args, **kwargs in its forward(), so it is not affected by this MCore change.

Test coverage gap: Gemma4TopKRouter.routing() has no unit test. There are no tests in tests/unit_tests/models/gemma/ that instantiate Gemma4TopKRouter or call its routing() method. While this specific change is safe (just passing through a kwarg), the renormalization + per_expert_scale logic in routing() is also untested. Consider adding a unit test that verifies: (1) routing() accepts input_ids without error, and (2) the renormalization and per_expert_scale application produce correct outputs.

Lock file: The uv.lock diff is large but consists entirely of transitive dependency version bumps from the MCore submodule update. Notable changes: torch 2.11 to 2.12, wandb 0.26 to 0.27, langchain-openai removed (replaced by langchain-nvidia-ai-endpoints via nvidia-resiliency-ext), ssm and te extras folded into dev/lts. No pyproject.toml changes in this PR, so the lock changes are expected to come from the submodule.

Suggested test cases: No perf tests impacted.

[yaoyu-33]

/ok to test c36259c

[yaoyu-33]

/ok to test e4d4802