Add input unpacking for tuple inputs in decoder layers

[NicoGrande]

Description

This PR finishes the work started by @gagika in #2753. Credits to @gagika for uncovering this regression!

This PR fixes a regression introduced in a recent commit where decoder layers began returning a tuple (hidden_states, kv_cache) even when scan_layers=False. This caused failures in SequentialBlockDecoderLayers (used in pipeline parallelism) because the subsequent layer would receive a tuple as input instead of a tensor, leading to a ValueError in normalization layers.

Changes

• Added a check at the beginning of the __call__ method for all Decoder Layer implementations to handle tuple inputs gracefully:

  # Unpack inputs if it's a tuple (e.g. from a previous layer returning (hidden_states, kv_cache))
  if isinstance(inputs, tuple):
    inputs = inputs[0]
    

Changes

• Updated Qwen3NextScannableBlock to correctly unpack the (output, kv_cache) tuple inside its scan loop.

If the change fixes a bug:
FIXES: b/462751017

Affected Models

This fix covers the following model architectures:

• DeepSeek (Dense & MoE)
• Gemma, Gemma 2, Gemma 3
• GPT-3, GPT-OSS
• Llama 2, Llama 4
• Mistral, Mixtral
• Qwen 3 (Dense, MoE, and Next)
• Simple Layer (Testing)

Tests

internal google3 tests.

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.

[NuojCheng]

Thank you Nico! LGTM

[gagika]

Thanks Nico