Gemma3 Model Token Logits Mismatch with Different max_prefill_predict_length

[babyplutokurt]

Bug report

Gemma3 Model Token Logits Mismatch with Different max_prefill_predict_length

Description

When running prefill with the Gemma3 4B model, the output token logits differ depending on the max_prefill_predict_length value.
In contrast, other models (e.g., Qwen3, LLaMA 2, LLaMA 3) produce consistent logits regardless of the prefill length.

Steps to Reproduce

1. Load the Gemma3 4B model.
2. Use the same input prompt: The capital of France is
3. Run inference with:

• Case 1: max_prefill_predict_length = 1024, max_target_length = 2048
• Case 2: max_prefill_predict_length = 4096, max_target_length = 8192

Observed Behavior (Gemma3 4B)

Case 1:
[null, {"The": -1.68888}, {" capital": -20.57046}, {" of": -0.15805}, {" France": -3.24782}, {" is": -1.67389}]
Full output: The capital of France is Paris.\n\nParis is a global center for art

Case 2:
[null, {"The": -1.68888}, {" capital": -20.30916}, {" of": -0.16284}, {" France": -3.24021}, {" is": -1.64998}]
Full output: The capital of France is Paris�Ar¡ 3<start_of_image>--- 3

Note: The very first token logits matches! Remaining tokens start to divergence. Strongly suspect related to kv cache.

Note: The small cumulative logit differences lead to large divergence in generated output for longer sequences.

Expected Behavior

Token logits for the same input prompt should remain identical across different max_prefill_predict_length values, as observed with Qwen3, LLaMA 2, and LLaMA 3 models.

Additional Information

• Affected Model: Gemma3 4B
• Verified Models (working as expected): Qwen3, LLaMA 2, LLaMA 3

Logs/Output

No response

Environment Information

No response

Additional Context

No response

[gagika]

Could you share your setup and print output tokens?
I am wondering if there is an EOS token after The capital of France is Paris, and after EOS token it's just generating a garbage.

[babyplutokurt]

Could you share your setup and print output tokens? I am wondering if there is an EOS token after The capital of France is Paris, and after EOS token it's just generating a garbage.

Config:

    MaxText/configs/base.yml \
    model_name="gemma3-4b" \
    tokenizer_path="assets/tokenizer.gemma3" \
    load_parameters_path="gs://maxtext-gemma/unified/gemma3/4b/unscanned/2025-08-05-18-18/0/items" \
    per_device_batch_size=4 \
    ici_tensor_parallelism=4 \
    max_prefill_predict_length=4096 \
    max_target_length=8192 \
    async_checkpointing=false \
    scan_layers=false \
    attention="dot_product" \
    return_log_prob=True

Case 1: max_prefill_predict_length = 1024, max_target_length = 2048
Prompt + Output:"The capital of France is Paris.\n\nParis is a global center for art"
Tokenized token:["","The"," capital"," of"," France"," is"," Paris",".","\n\n","Paris"," is"," a"," global"," center"," for"," art"]
"Token_logprobs":[null,{"The":-1.6888837814331055},{" capital":-20.570457458496094},{" of":-0.15805160999298096},{" France":-3.2478187084198},{" is":-1.6738873720169067},{" Paris":-0.19891169667243958},{".":-0.10209408402442932},{"\n\n":-0.704276978969574},{"Paris":-0.5836753845214844},{" is":-0.0494668111205101},{" a":-1.1100436449050903},{" global":-0.8351733088493347},{" center":-0.09056396782398224},{" for":-0.005966233555227518},{" art":-0.0009556006407365203}]

Case 2: max_prefill_predict_length = 4096, max_target_length = 8192
Prompt + Output:"The capital of France is Paris�Ar¡ 3<start_of_image>--- 3"
Tokenized token: ["","The"," capital"," of"," France"," is"," Paris","�","Ar","¡"," ","3","<start_of_image>","---"," ","3"]
"Token_logprobs":[null,{"The":-1.6888837814331055},{" capital":-20.309162139892578},{" of":-0.16284018754959106},{" France":-3.240208387374878},{" is":-1.6499804258346558},{" Paris":-0.21631987392902374},{"�":-0.6105182766914368},{"Ar":-2.0301079750061035},{"¡":-0.417588472366333},{" ":-1.4758585691452026},{"3":-1.1177611351013184},{"<start_of_image>":-1.1808792352676392},{"---":-2.1718084812164307},{" ":-0.7211616635322571},{"3":-0.05767672508955002}]

First token is "" (BOS), predict next token "the", with identical logprob -1.6888837814331055, after the first token, all logprobs starts to diverge.

The logprobs should be identical to different max prefill length.

I suspect something related to QK cache normalization. Other models like llama3, qwen3, llama2 are verified matched logprobs in different prefill length