Skip to content

Fix: Gemma 3 & 4 Base Model and Flax Linen Decoding #4066

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jul 2, 2026
Merged

Fix: Gemma 3 & 4 Base Model and Flax Linen Decoding #4066

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/maxtext/configs/inference/vllm.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,9 @@ logical_axis_rules: [
['mlp', ['attn_dp', 'model']],
['embed', []],
['norm', []],
['layers', []],
['dense_layers', []],
['moe_layers', []],
# ==========================================
# Inference(Prefill, Decode, Cache)
# ==========================================
Expand Down
143 changes: 110 additions & 33 deletions src/maxtext/layers/decoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1056,6 +1056,8 @@ def __call__(
bidirectional_mask_value,
previous_chunk,
slot,
kv_caches=kv_caches,
attention_metadata=attention_metadata,
)
elif cfg.decoder_block == DecoderBlockType.GEMMA4:
bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None
Expand All @@ -1068,6 +1070,8 @@ def __call__(
bidirectional_mask_value,
previous_chunk,
slot,
kv_caches=kv_caches,
attention_metadata=attention_metadata,
)
elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4:
y = self._apply_deepseek4_scanned_blocks(
Expand Down Expand Up @@ -1303,6 +1307,8 @@ def _apply_gemma3_scanned_blocks(
bidirectional_mask,
previous_chunk,
slot,
kv_caches=None,
attention_metadata=None,
):
"""Applies Gemma3 scanned decoder blocks, handling main scan and remainders."""

Expand All @@ -1316,27 +1322,43 @@ def _apply_gemma3_scanned_blocks(
policy = self.get_remat_policy()
RemattedGemma3Block = self.set_remat_policy([gemma3.Gemma3ScannableBlockToLinen], policy)[0]

layer_call_kwargs = {"bidirectional_mask": bidirectional_mask}
layer_kwargs = {"num_of_layers": attention_pattern_length}

# Apply the main scan over the full blocks
if scan_length > 0:
broadcast_args = (
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
kv_caches, scan_length, attention_pattern_length, stack=True
)
y, _ = self.scan_decoder_layers(

broadcast_args_spec = [
(decoder_segment_ids, nn.broadcast),
(decoder_positions, nn.broadcast),
(deterministic, nn.broadcast),
(model_mode, nn.broadcast),
(slot, nn.broadcast),
(None, nn.broadcast), # page_state
(previous_chunk, nn.broadcast),
(bidirectional_mask, nn.broadcast),
(kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast),
(attention_metadata, nn.broadcast),
]
broadcast_args = tuple(arg for arg, _ in broadcast_args_spec)
in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec)

y, returned_kv_cache = self.scan_decoder_layers(
cfg,
RemattedGemma3Block,
scan_length,
"layers",
mesh,
in_axes_tuple=(nn.broadcast,) * len(broadcast_args),
in_axes_tuple=in_axes_tuple,
model_mode=self.model_mode,
**layer_kwargs,
)(y, *broadcast_args, **layer_call_kwargs)
)(y, *broadcast_args)

maxtext_utils.update_kv_caches_after_scan(
kv_caches, returned_kv_cache, scan_length, attention_pattern_length, stacked=True
)

# Apply any remaining layers that did not fit into a full scanned block
num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length
Expand All @@ -1346,16 +1368,37 @@ def _apply_gemma3_scanned_blocks(
layer = RemattedGemma3Block(
config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, name="layers_remainder", **rem_layer_kwargs
) # pytype: disable=wrong-keyword-args
y, _ = layer(

remainder_kv = None
if kv_caches is not None:
start_idx = scan_length * attention_pattern_length
remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers])

y_and_kv = layer(
y,
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
previous_chunk=previous_chunk,
slot=slot,
**layer_call_kwargs,
bidirectional_mask=bidirectional_mask,
kv_cache=remainder_kv,
attention_metadata=attention_metadata,
)

if isinstance(y_and_kv, tuple):
y = y_and_kv[0]
updated_remainder_kv = y_and_kv[1]
else:
y = y_and_kv
updated_remainder_kv = None

if kv_caches is not None and updated_remainder_kv is not None:
start_idx = scan_length * attention_pattern_length
for offset, updated_item in enumerate(updated_remainder_kv):
kv_caches[start_idx + offset] = updated_item

return y

def _apply_gemma4_scanned_blocks(
Expand All @@ -1368,6 +1411,8 @@ def _apply_gemma4_scanned_blocks(
bidirectional_mask,
previous_chunk,
slot,
kv_caches=None,
attention_metadata=None,
):
"""Applies Gemma4 scanned decoder blocks, handling main scan and remainders."""

Expand All @@ -1379,28 +1424,32 @@ def _apply_gemma4_scanned_blocks(
num_full_blocks = cfg.num_decoder_layers // block_pattern_len
remainder_layers = cfg.num_decoder_layers % block_pattern_len

broadcast_args = (
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
)

# Pass slot/previous_chunk/bidirectional_mask by keyword only (via layer_call_kwargs),
# never positionally in broadcast_args: Gemma4DecoderLayer and Gemma4ScannableBlock
# declare slot and previous_chunk in swapped order, so positional passing misroutes them.
layer_call_kwargs = {
"slot": slot,
"previous_chunk": previous_chunk,
"bidirectional_mask": bidirectional_mask,
}

if num_full_blocks > 0:
ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen
policy = self.get_remat_policy()
RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0]

kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan(
kv_caches, num_full_blocks, block_pattern_len, stack=True
)

broadcast_args_spec = [
(decoder_segment_ids, nn.broadcast),
(decoder_positions, nn.broadcast),
(deterministic, nn.broadcast),
(model_mode, nn.broadcast),
(slot, nn.broadcast),
(None, nn.broadcast), # page_state
(previous_chunk, nn.broadcast),
(bidirectional_mask, nn.broadcast),
(kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast),
(attention_metadata, nn.broadcast),
]
broadcast_args = tuple(arg for arg, _ in broadcast_args_spec)
in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec)

# For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks
y, _ = nn.scan(
y, returned_kv_cache = nn.scan(
RemattedGemma4Block,
variable_axes={
"params": cfg.param_scan_axis,
Expand All @@ -1410,7 +1459,7 @@ def _apply_gemma4_scanned_blocks(
"_overwrite_with_gradient": 0,
},
split_rngs={"params": True, "dropout": cfg.enable_dropout},
in_axes=(nn.broadcast,) * len(broadcast_args),
in_axes=in_axes_tuple,
length=num_full_blocks,
metadata_params={
nn.PARTITION_NAME: "layers",
Expand All @@ -1424,7 +1473,11 @@ def _apply_gemma4_scanned_blocks(
num_of_layers=block_pattern_len,
name="scanned_blocks",
)(
y, *broadcast_args, **layer_call_kwargs
y, *broadcast_args
)

maxtext_utils.update_kv_caches_after_scan(
kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True
)

# Process any remaining layers that don't fit into a full scanned block
Expand All @@ -1438,9 +1491,33 @@ def _apply_gemma4_scanned_blocks(
attention_type=attention_type,
layer_idx=layer_id,
)
y = layer(y, *broadcast_args, **layer_call_kwargs)
if cfg.scan_layers:
y = y[0]
kv_cache = kv_caches[layer_id] if kv_caches is not None else None

# inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode,
# previous_chunk, page_state, slot, bidirectional_mask, kv_cache, attention_metadata
remainder_args = (
decoder_segment_ids,
decoder_positions,
deterministic,
model_mode,
previous_chunk,
None, # page_state
slot,
bidirectional_mask,
kv_cache,
attention_metadata,
)

y_and_kv = layer(y, *remainder_args)
if isinstance(y_and_kv, tuple):
y = y_and_kv[0]
new_kv = y_and_kv[1]
else:
y = y_and_kv
new_kv = None

if kv_caches is not None and new_kv is not None:
kv_caches[layer_id] = new_kv

return y

Expand Down
Loading
Loading