[Spec Decode] Support Muse Glimmer DFlash checkpoints

[lming2001]

Description

Add generic TPU DFlash compatibility for the released Muse Glimmer assistant checkpoint.

The existing JAX DFlash path already provides proposal and verification infrastructure. This PR makes it compatible with Muse's public checkpoint layout by:

• registering MuseGlimmerAssistantModel and DFlashMuseGlimmerAssistantModel;
• reading DFlash fields from either nested dflash_config or top-level checkpoint fields;
• honoring use_sliding_window=true and sliding_window=2048 in causal and non-causal draft attention;
• mapping encoder.fc and encoder.output_norm_enc weights;
• sizing draft KV cache entries from explicit head_dim=128; and
• resolving target LM-head parameter leaves safely.

PR #3527 supplies Muse Glimmer's configured auxiliary hidden states. This PR itself is independent and can merge in either order.

FIXES: #3526

Tests

JAX_PLATFORMS=cpu python -m pytest -q \
  tests/spec_decode/test_dflash.py::test_dflash_sliding_window_requires_explicit_enablement \
  tests/spec_decode/test_dflash.py::test_registry_supports_muse_glimmer_assistant_architectures \
  tests/runner/test_kv_cache_manager.py::TestKVCacheManager::test_get_kv_cache_spec_with_dflash_uses_explicit_head_dim

pre-commit run --files \
  tests/runner/test_kv_cache_manager.py \
  tests/spec_decode/test_dflash.py \
  tpu_inference/models/common/model_loader.py \
  tpu_inference/models/jax/dflash.py \
  tpu_inference/runner/kv_cache_manager.py \
  tpu_inference/spec_decode/jax/dflash.py

Results:

• New CPU regression tests: 3 passed.
• Combined target+DFlash focused suite on TPU v7: 57 passed.
• All applicable pre-commit hooks passed. The license hook was verified manually because the development host could not resolve proxy.golang.org; this PR adds no new source files.
• Current-upstream TPU v7 smoke: target and assistant loaded, reasoning/tool probes passed, and two repeated runs each completed 8/8 long-context requests. DFlash accepted 1,953 / 3,759 draft tokens (52.0%) and 1,958 / 3,744 (52.3%). These runs used on-demand compilation and are intentionally reported as correctness validation, not steady-state throughput benchmarks.

The assistant's trained sliding window is correctness-critical: prior TPU validation measured aggregate draft-token acceptance rising from 11.1% to 47.6%, with per-position acceptance of 70.5%, 44.8%, and 27.6% for three proposals. A K+1 runtime-width hypothesis reduced acceptance to 7.1%, so this change intentionally preserves the checkpoint's block_size=16.

AI assistance was used to port the internally validated implementation and run tests. The human submitter remains responsible for reviewing and defending every changed line.

Checklist

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• Unit tests cover architecture aliases, top-level config fields, sliding-window selection, and explicit KV head size.
• I have received at least 1 readability approval and 1 correctness approval.