perf(qwen3-omni): project only final prefill logits

[CjhHa1]

Summary

Avoid materializing prompt-wide vocabulary logits in the Qwen3-Omni trainside prefill when generation samples only the final row.

The current Transformers 5.6 Qwen3-Omni Thinker forward does not expose logits_to_keep. A scoped helper selects hidden_states[:, -1:, :] at the lm_head boundary without replacing the module or changing its state dict. The hook exists only around the first prefill forward; decode and replay remain unchanged.

Related Issue

N/A

Test Plan

• SKIP=no-commit-to-branch pre-commit run --all-files --show-diff-on-failure — pass
• focused construction check: scoped LM-head selection returns [B, 1, V], matches the full projection's final row, and removes its hook after the forward
• H20 LM-head A/B microbenchmark (B=4, L=512, H=2048, V=152064, bf16, task selective-logits-head-bench-0913): peak allocation 0.580 GiB -> 1.160 MiB; median 9.09 ms -> 0.202 ms; selected logits exact
• full Qwen3-Omni Thinker first-prefill A/B on 8x H20, PyTorch 2.10, Transformers 5.6, bf16, Accelerate balanced device map, input [4, 512] (task pr446-full-prefill-ab-8h20-0913-v3): max per-GPU forward allocation 0.596 -> 0.307 GiB (-0.289 GiB, -48%); median 259.14 -> 254.48 ms (+1.8%); selected logits close and sampled KV cache exact
• production FSDP2 full-shard first-prefill A/B on 8x H20, one rank/GPU, production LoRA/wrap settings, input [4, 512] (task pr446-fsdp-qwen-8h20-0914): max-rank forward peak 3.958 -> 3.958 GiB (unchanged because FSDP all-gather dominates); median max-rank latency 363.03 -> 344.13 ms (+5.5%); selected logits close and sampled KV cache exact on every rank
• full token-generation/training loop: not run; the changed first-prefill forward is isolated above with the real checkpoint and production FSDP wrapping

Compatibility / Risk

No config, checkpoint, state-dict, or public model API changes. Prompts are left-padded before generation, so the final hidden row is each sample's last real token. The forward pre-hook is removed in finally; subsequent decode steps and replay do not carry it.

Under production FSDP the optimization improves prefill latency, but does not reduce the observed peak because the larger FSDP layer all-gather determines that peak.

Reviewer Notes

The native Qwen logits_to_keep changes live in #442. This PR covers Qwen3-Omni because its current forward lacks that argument.

HunyuanImage3 was benchmarked and removed from this PR: under its production FSDP topology, peak changed only 9.474 -> 9.471 GiB and latency regressed 2672.10 -> 2707.63 ms, so the scoped projection did not justify its added complexity.

Checklist

• I reviewed the changed code and removed unrelated/generated artifacts.
• I updated tests, docs, and configs where needed, or explained why not.