nemotron_h: load in-checkpoint MTP head for self-speculation

[pierre427]

nemotron_h: load in-checkpoint MTP head for self-speculation

Stacked on #1456 (MTP self-speculation machinery + qwen3_5 MTP-head
template). Please review/merge #1456 first.

What

Nemotron 3 Super checkpoints ship a DeepSeek-style multi-token-prediction
(MTP) head alongside the main backbone. This wires that head into
nemotron_h.py so it loads from the checkpoint and drives self-speculative
decoding through the same interface #1456 established for qwen3_5
(logits / make_mtp_cache / mtp_step, plus sanitize that drops the
module when a checkpoint carries no mtp.* tensors).

Because the backbone is a hybrid Mamba-2/attention stack, the recurrent
(SSM) layers cannot be trimmed the way a KV cache can. Rolling back a
rejected speculative block therefore needs to replay the accepted prefix
through the SSM update rather than truncate it. This PR adds that support:

• ssm_sink — an optional list threaded through the Mamba mixer / block
  / model forwards. When present, each Mamba layer appends the exact inputs
  (hidden_states, B, C, dt, per-layer A_log/D/dt_bias, the
  pre-update state, mask, and padded conv input) needed to replay its update.
  When None (the default, i.e. normal decoding) there is zero overhead.
• Model.rollback_speculative_cache(caches, ssm_states, keep, block_size)
  — trims the KV caches normally and rebuilds each Mamba ArraysCache by
  re-running ssm_update on the kept prefix (or restoring the captured
  pre-verify state verbatim when keep == 0). Single-sequence (B=1,
  unpadded) only; raises on padded/batched caches.
• NemotronHMTPBlock / NemotronHMTP — the head itself. The first MTP
  layer carries the embed/hidden fusion (eh_proj / enorm / hnorm), the
  last carries final_layernorm; each block is otherwise a standard
  NemotronHBlock. The shipped head is transformer (attention) blocks, so
  mtp_step builds the attention mask.
• ModelArgs gains num_nextn_predict_layers, mtp_layers_block_type,
  and mtp_hybrid_override_pattern (the block-type list is normalized to
  single-char pattern codes in __post_init__, mirroring the backbone's
  hybrid_override_pattern). The MTP module is built only when the config
  declares MTP layers and dropped again in sanitize if the checkpoint has
  no mtp.* tensors, so strict loading stays consistent for plain
  (non-MTP) Nemotron-H checkpoints.

How it wires to the self-spec machinery

Identical head interface to #1456's qwen3_5 template:

• logits(hidden) projects a hidden through lm_head.
• make_mtp_cache() builds the head's own per-layer cache (KV for
  attention layers).
• mtp_step(hidden, tokens, mtp_cache) runs one MTP forward: fuses the
  backbone hidden with the next token's embedding via
  enorm/hnorm/eh_proj, runs the head layers, applies final_layernorm,
  and returns (logits, post_hidden) so draft depth can be chained.
• sanitize uses the same "drop the module when the checkpoint has no
  mtp.* tensors" logic (and only stacks experts for prefixes actually
  present, covering both backbone and MTP layers).

The one Nemotron-specific addition over the qwen template is the SSM
rollback path (ssm_sink + rollback_speculative_cache). Nemotron's
recurrent state does not use the base's ArraysCache.record_rollback
hook — it is self-contained: the verify forward captures replay inputs via
ssm_sink, and rollback_speculative_cache rebuilds the state directly. No
new cache hook was required; all helpers it needs (ssm_update,
create_attention_mask, ArraysCache, KVCache) already exist on the
#1456 base.

Test

tests/test_nemotron_h_mtp.py (mirrors tests/test_qwen3_5_mtp.py):

• test_mtp_step_shapes_and_chaining — tiny hybrid config with a
  2-layer attention MTP head; real backbone forward, then mtp_step
  produces a valid next-token prediction with correct shapes and cache
  offsets, and chaining on the head's own hidden advances the cache.
• test_mtp_module_dropped_without_weights — a checkpoint with no
  mtp.* tensors drops the module (model.mtp is None).
• test_rollback_speculative_cache_matches_fresh_forward — exercises
  ssm_sink capture + rollback_speculative_cache: after a block_size-4
  verify forward with keep-2 accepted, both the KV caches and the Mamba
  recurrent states match a fresh forward over just the committed prefix plus
  kept tokens.

All pass; test_qwen3_5_mtp.py and the existing nemotron_h model test are
unaffected. black + isort --profile black clean.