fix(xla): update Gemma 3n emitter to the shared-layer query offset #2466

Problem / Background

PR #2457 fixed the Gemma 3n shared-KV layer query offset in the MLX path (src/models/gemma3n.rs, the offset comes from rope_offset at line 439, matching mlx-lm's cache[0].offset). The mlxcel-xla emitter was deliberately left mirroring the old offset, so MLX-vs-XLA Gemma 3n parity now diverges. Related to epic #566.

Current Behavior

src/lib/mlxcel-xla/src/emitter/gemma3n_emit_ops.rs:790-800 builds query_positions for shared-KV layers (lw.wk.is_none()) as positions + real_len, with a comment saying it reuses the concrete layer's advanced cache offset. That is the pre-#2457 behavior. Separately, the xla-diagnostics call sites in src/models/gemma3n.rs were updated by hand and never compiled, because the xla feature does not build on the GB10 host.

Proposed Solution

Update the emitter so shared-layer query RoPE positions equal what the fixed MLX path computes. Read the post-#2457 gemma3n.rs logic first and port exactly that, then fix the comment. Compile the diagnostics call sites on a host that builds the xla feature and fix any errors.

Scope

In scope: gemma3n_emit_ops.rs, the xla-diagnostics call sites in src/models/gemma3n.rs, the Gemma 3n parity test.

Out of scope: the rest of epic #566.

Implementation Notes

• Reuse: the existing MLX-vs-XLA Gemma 3n parity check.
• Constraints: needs a host where --features xla builds (not GB10). If none is available, say so in the PR and ask the maintainer to run it.
• Edge cases: prefill versus decode steps, sliding versus full layers.
• Error handling: none new.

Acceptance Criteria

• The emitter computes shared-layer query positions as the fixed MLX path does.
• The existing MLX-vs-XLA Gemma 3n parity check passes on a host that builds xla, and failed there before the change.
• The xla-diagnostics call sites in gemma3n.rs compile with --features xla.
• The XLA Gemma 3n serve path uses the corrected emitter.

Verification

cargo build --features xla
cargo test --features xla gemma3n -- --test-threads=1
cargo fmt --all -- --check
cargo clippy --features xla -- -D warnings

A pass is a clean build with the diagnostics sites compiled and the parity check green.