fix(models): declare and honor logits_to_keep on Mistral3 VLM forward

[kabirvashisht4-glitch]

What does this PR do ?

Makes Mistral3ForConditionalGeneration declare and honour logits_to_keep, so
a configured fused loss is no longer silently replaced.

forward absorbed the kwarg into **kwargs and never read it.
_supports_logits_to_keep inspects the forward signature
(components/utils/model_utils.py L76-77), so the check returned False and
the recipe swapped the configured loss for a fresh MaskedCrossEntropy
(recipes/llm/train_ft.py L159, recipes/vlm/finetune.py L554,
_transformers/infrastructure.py L640). The class is registered
(_transformers/registry.py L194-197), so this is user-reachable.

Two consequences, both from #3510:

1. The full logits tensor is materialised. Mistral4Config.vocab_size
   defaults to 131072, so the head output is [tokens, 131072].
2. The configured loss_fn's settings are dropped with it — the replacement
   is a fresh MaskedCrossEntropy(), so fp32_upcast: false is lost.

Same defect and same fix as #3510 (MiniMax-M3). Evidence it was an oversight:
its own text backbone Mistral4ForCausalLM declares it (same file, L417), and
its sibling Mistral3FP8VLMForConditionalGeneration declares it
(mistral3_vlm/model.py L206).

Contract. The recipe calls model(logits_to_keep=1, **batch) and then
requires "hidden_states" in out; FusedLinearCrossEntropy.forward takes
hidden_states and applies the lm_head itself, fused with the loss. So the
fused path returns the hidden states rather than sliced logits — matching what
#3510 did for M3. The default path (logits_to_keep=0) is untouched and still
returns the logits tensor.

Changelog

• Mistral3ForConditionalGeneration.forward declares
  logits_to_keep: Union[int, torch.Tensor] = 0 and returns
  {"hidden_states": ...} when it is set; the default path still returns
  logits, unchanged.
• Hoisted the THD check into a single is_thd local (it was computed twice from
  kwargs, once before and once after squeeze_input_for_thd mutates it) and
  applied the same batch-dim restoration on the hidden-states path.
• New tests/unit_tests/models/mistral4/test_mistral4_logits_to_keep.py
  covering the recipe's _supports_logits_to_keep probe, the fused path
  returning hidden states without materialising logits, and the unchanged
  default path.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

On the test file location. These live in a new file rather than in
test_mistral4_model.py because that module carries a module-level
pytest.mark.skipif(not torch.cuda.is_available()), so the contract would not
be checked on the CPU tier. The new file is CPU-runnable, mirroring
tests/unit_tests/models/minimax_m3_vl/test_minimax_m3_logits_to_keep.py.

Verification (CPU, no GPU or checkpoint needed):

• pytest tests/unit_tests/models/mistral4/test_mistral4_logits_to_keep.py → 3 passed.
• Reverting only the model change makes test_recipe_detects_logits_to_keep_support
  and test_logits_to_keep_returns_hidden_states_and_skips_lm_head fail;
  test_without_logits_to_keep_still_returns_logits passes either way by design,
  guarding the path this must not change.
• pytest tests/unit_tests/models/mistral4/ tests/unit_tests/models/mistral3/ tests/unit_tests/models/mistral3_vlm/ → 147 passed, 70 skipped.
• Wider tests/unit_tests/models/ tests/unit_tests/_transformers/ run: 84 failures
  before and after this change (gemma4 / glm5_next / retrieval), all pre-existing
  on main in my environment and unrelated — the only delta is the 2 tests above
  going green. diffusion_gemma and gemma4_unified were excluded locally: they
  fail to import transformers.models.diffusion_gemma on main too.
• ruff format / ruff check clean.

Additional Information

• Closes Mistral3ForConditionalGeneration: forward does not declare logits_to_keep, so fused losses are silently replaced #3778

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.