Fix DeepSeek decode running the prefill attention path under NNX

[ecnal-cienet]

Description

Decoding deepseek2-16b through MaxEngine on the NNX path emits a single repeated token. On the nightly checkpoint (v6e-8, maxtext.inference.decode, greedy, mla_naive_kvcache=false):

prompt: An attention function can be described as mapping a query and a set of key-value pairs to an output,
        where the query, keys, values, and outputs are all vectors. The output is

before:  a a a a a a a a a a a a a a a a a a a a a ...   (all 1024 tokens)
after:   a weighted sum of the values, where the weights are computed by a function of the query and the keys.

Root cause

DeepSeekGenericLayer.attention_op passed model_mode=self.model_mode to attention — the value captured at __init__ — instead of the model_mode its __call__ receives. Same in self_attention_with_norm_op for the MHC branch.

MaxEngine's NNX path builds the model once in PREFILL mode and passes the mode per call: _nnx_run_model always merged self.graphdef, which is the prefill graphdef. So on every autoregressive step each DeepSeek layer still read self.model_mode == "prefill" and took the prefill KV-cache branch. Prefill logits stay correct, which is why the first token is right and everything after it is wrong.

Only DeepSeek is affected — llama2 and the other families already forward the call-time model_mode. A sweep of every non-__init__ read of self.model_mode in src/maxtext found DeepSeek to be the only one feeding a stale mode into attention.

The bug was latent under Linen, which re-binds layers on each apply so self.model_mode is always current. It only surfaces once a model object is built once and reused across modes, which is exactly what pure NNX does.

Fix

File	Change
models/deepseek.py	attention_op and self_attention_with_norm_op take model_mode and forward the call-time value.
models/deepseek4.py	Updated call site (DeepSeek4DecoderLayer inherits both helpers).
inference/maxengine/maxengine.py	Cache a graphdef_ar and merge with the graphdef built for the requested mode.

Either change alone restores correct output; both are included deliberately, so a future layer repeating the mistake cannot silently corrupt decode and the two protections cannot regress together unnoticed.

Tests

Two new CPU-only files, neither needing a checkpoint. They close a real gap: golden-logit tests only exercise the forward pass, so nothing covered the prefill → autoregressive handoff.

tests/unit/decoder_layer_model_mode_test.py — 16 cases. Asserts at the attention boundary that a layer attends in the mode it was called with. Covers deepseek dense/moe, llama2, mistral, gemma, gemma2, qwen2 and qwen3, in both directions (built-prefill/called-AR and built-AR/called-prefill) so no layer can pass by hardcoding a mode.

tests/unit/deepseek_decode_consistency_test.py — 4 cases. Greedy decode through prefill/insert/generate must equal the argmax rollout of a plain forward pass over the same prefix. Parameterized over deepseek-moe, deepseek-dense and a generic model, so adding a family is one line. Includes a separate guard that generate merges an autoregressive graphdef, covering the engine side on its own.

Verification

All on wanglance-v6e-8 (TPU v6 lite, 8 chips). Rows marked unmodified main were run in a throwaway git worktree, so the branch was never reverted.

Item	Result	Log
Real-checkpoint decode, with fix	-> a weighted sum of the values, where the weights are computed by a function of the query and the keys.	log
Real-checkpoint decode, unmodified main	-> a a a a a a a a a ... for all 1024 tokens — the bug reproduced	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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[xibinliu]

Good catch of the model_mode difference for deepseek attention_op

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.