Previously the vanilla reference was a Gemma4ForCausalLM shell wrapping
the language_model (the same construction used inside
gemma4_flex_inference.main for LoRA / state-dict hashing convenience).
That is not plain HF: `Gemma4Model.forward` uses
`create_masks_for_generate(..., mm_token_type_ids, pixel_values)` to
build attention masks, while the shell calls `Gemma4TextModel.forward`
directly, which builds its own per-regime masks via `create_causal_mask`
+ `create_sliding_window_causal_mask`. Both are correct for text-only
input but their mask-bias precision differs enough to produce a
measurable drift.
The script now keeps both references alive and reports three diffs:
shell vs raw, flex vs raw, flex vs shell. On unsloth/gemma-4-E2B-it
bf16 with a 6-token prompt:
shell vs raw max 6.9e-01 mean 3.0e-01 argmax=yes top-10=10/10
flex vs raw max 6.3e-01 mean 2.2e-01 argmax=yes top-10=10/10
flex vs shell max 3.8e-01 mean 8.0e-02 argmax=yes top-10=10/10
Flex is actually closer to raw HF than the shell is. About 0.30 mean of
the flex-vs-shell-and-vs-raw gap comes from the shell's mask
construction alone, not the flex kernel. Either way the bf16 drift
band matches Qwen3 (0.3 / 0.09) and Llama-3.2 (0.13 / 0.02), and
semantic top-1 + top-10 are exact.
One-shot correctness checks for the flex_attention + paged KV path
against a vanilla HF `model(input_ids, use_cache=False)` forward on
the same prompt. Report max / mean abs diff on last-position logits
plus argmax match and top-10 overlap, so numerical drift and semantic
equivalence are both visible.
Shared approach: load the base model, deep-copy it for the flex path
so the attention patching does not mutate the vanilla comparison, run
both on the same tokenized prompt, report diffs.
`verify_gemma4_numerics.py` mirrors `gemma4_flex_inference.py`'s
text-only loader (Gemma4ForConditionalGeneration -> drop vision / audio
towers -> language_model into a Gemma4ForCausalLM shell) before
deep-copying. The softcap is NOT re-applied on the vanilla logits
since `Gemma4ForCausalLM.forward` already applies
`final_logit_softcapping`.
`verify_qwen3_numerics.py` runs `qwen3_flex_inference.FlexInference`
against either Qwen3 or Llama-3.2 via `--model_name`. `fa4_prefill` is
disabled because short prompts hit a CuteDSL sm_100 shape mismatch in
`handle_block_sparse_empty_tile_correction_sm100`.
Results on B200 bf16, 6 to 7 token prompt:
| Model | max abs | mean abs | argmax | top-10 |
|---------------------------------|---------|----------|--------|--------|
| unsloth/Qwen3-4B-Base | 0.313 | 0.088 | yes | 10/10 |
| unsloth/Llama-3.2-3B-Instruct | 0.125 | 0.021 | yes | 10/10 |
| unsloth/gemma-4-E2B-it | 0.375 | 0.080 | yes | 10/10 |
All three land in the same bf16 Triton-flex vs eager-matmul drift band;
Gemma-4's extra per-layer-input path and layer_scalar do not widen the
gap despite 20 of 35 layers going through the shared-KV link.
