Description

This PR fixes multiple issues causing gibberish outputs and initialization crashes when using scan_layers=True with Gemma 3 and 4 in vLLM workflows.

Context and Details:

• Why is this change being made: To enable successful RL training and generation with Gemma 3 and 4 by ensuring KV caches and scanned layer weights are correctly handled in the
  vLLM rollout engine.
• The problems solved:
  1. vLLM Decoder Gibberish: Fixed missing kv_caches argument in _apply_gemma3_scanned_blocks, which caused loss of attention context.
  2. KV Cache Shape Mismatches: MaxTextVllmRollout ignored rollout_vllm_additional_config, falling back to base configs that conflicted with HuggingFace config head dimensions.
  3. Native NNX Keys & Scalar Bugs: The weight sync preprocessor failed to unstack nested scanned weights because it mismatched on native NNX root keys, and a bug in the scan_length calculation for scalar RNG states caused it to drop all decoder parameters, resulting in random noise generation.
• Implementation details:
  • Explicitly passed kv_caches and attention_metadata in decoder blocks.
  • Updated the custom preprocessor in MaxTextVllmSampler to use robust string matching for NNX keys, safely ignore scalars when calculating scan_length, and correctly broadcast scalar values during unstacking.

Tests

inference vllm_decode

Run script : here
Run output : log

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.