Kimi-K3 support for speculative-decoding draft training

[h-guo18]

What does this PR do?

Type of change: Bug fix

Two model-specific gaps prevented speculative-decoding draft training on Kimi-K3 from being numerically correct. Both are additive — no behavior change for any other model_type or tokenizer.

1. Loss mask — K3's chat format is not recognised

K3 replaces the K2/K2.5 <|im_*|> turn markers with an XTML tag format:

<|open|> message role assistant <|sep|> {content} <|close|> message <|sep|>

where tag names and attribute values (including the role) are ordinary text tokens; only <|open|> / <|close|> / <|sep|> / <|end_of_msg|> are special. Like the other Kimi models it ships only a slow tiktoken tokenizer, so apply_chat_template cannot emit an assistant mask and the loss_mask recovery registry is the only route.

This adds a kimi_k3 recovery that walks the XTML turns and marks the assistant content span. It tracks <|open|>/<|close|> depth rather than scanning for the next marker, so the nested think / response sub-tags the model generates stay inside the masked span.

There is a second-order problem: K3 keeps the legacy <|im_*|> markers for back-compat, so the existing kimi recovery also matches a K3 tokenizer — and since get_loss_mask_recovery returns the first match, K3 would silently get an all-zero mask (no <|im_middle|> ever appears in a K3-rendered sample). _kimi_detect now defers when the XTML markers are present.

2. Final norm — K3 is a VLM and the table is keyed on the text config

_FINAL_NORM_TYPE_BY_MODEL_TYPE is keyed on the resolved text config's model_type. Kimi-K3's outer config reports kimi_k3, but its text backbone reports kimi_linear, so nothing matched, _select_final_norm_type returned None, and FakeBaseModel built no final norm at all. The offline/streaming producers then reconstructed target logits from an un-normed hidden state.

This adds the kimi_linear entry. Note the entry deliberately is not kimi_k3 — a test pins that, since keying on the outer VLM model_type is the exact mistake this fixes.

3. Example — Kimi-K3 DSpark streaming training

Adds tools/launcher/examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml, mirroring the existing MiniMax-M3 and Kimi-K2.6 DSpark streaming examples. Its header documents the K3 settings that are otherwise silent failure modes — the two fixes above, the final aux capture id (93 == num_hidden_layers, which needs vllm#50815 on K3's block-residual backbone), and the explicit draft dims.

One deliberate difference from the sibling examples: the draft dims are described as draft-side choices, not mirrors of the backbone. K3's text config carries no rope_theta at all, and its FFN is 33792 dense / 3072 MoE — the draft copies neither, so claiming otherwise (as would be natural by analogy with the M3/K2.6 examples) would be wrong.

Usage

No API change. Both paths are selected automatically from the tokenizer / base config:

from modelopt.torch.utils.loss_mask import get_loss_mask_recovery

# Kimi-K3 tokenizer -> the XTML-aware recovery (previously: the K2 one, all-zero mask)
recovery = get_loss_mask_recovery(tokenizer)
assert recovery.name == "kimi_k3"
loss_mask = recovery.compute(tokenizer, input_ids)

Testing

Extends the existing unit tests; 27 passed locally:

pytest tests/unit/torch/utils/test_loss_mask.py \
       tests/unit/torch/speculative/plugins/test_modeling_final_norm.py

New cases:

• test_k3_recovery_is_registered — a K3 tokenizer selects kimi_k3.
• test_k3_tokenizer_does_not_match_the_k2_recovery — a tokenizer carrying both marker sets (i.e. real K3) still routes to kimi_k3. This is the regression that produced empty masks.
• test_k3_mask_marks_only_assistant_content — user turns stay unmasked.
• test_k3_mask_includes_nested_tags — nested think sub-tag stays inside the assistant span.
• test_k3_mask_accepts_list_input — parity with the existing K2 test.
• test_select_final_norm_type — adds kimi_linear -> "rmsnorm" and kimi_k3 -> None.

Beyond unit tests, both changes are exercised by a real Kimi-K3 DSpark drafter training run (streaming hidden-state extraction from vLLM, drafter exported); the loss mask is what selects the trained tokens and the final norm is what makes the reconstructed teacher logits correct.

Before your PR is "Ready for review"

• Is this change backward compatible?: ✅ — purely additive. The only change to existing behavior is _kimi_detect returning False for tokenizers that define K3's XTML markers, which no pre-K3 Kimi tokenizer does.
• If you copied code from any other sources or added a new PIP dependency, did you follow guidance in CONTRIBUTING.md: ✅ — no new dependencies, no copied code.
• Did you write any new necessary tests?: ✅ — unit tests for both fixes; the launcher example is config only.
• Did you update Changelog?: ❌ — happy to add an entry if you would like one for a model-support bug fix.
• Did you get Claude approval on this PR?: N/A — not an NVIDIA org member on this account.

Additional Information

Marked as draft for one reason: the commit carries the DCO sign-off but is not yet cryptographically signed (git commit -s -S). I will force-push a signed version before marking ready.

[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.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 067d8bf6-b540-43bd-a743-1f0520f344aa

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

_{Comment @coderabbitai help to get the list of available commands.}

[codecov]

Codecov Report

❌ Patch coverage is 97.82609% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 66.96%. Comparing base (2f6e77f) to head (0e25818).

Files with missing lines	Patch %	Lines
modelopt/torch/utils/loss_mask.py	97.82%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2049      +/-   ##
==========================================
+ Coverage   66.94%   66.96%   +0.02%     
==========================================
  Files         519      519              
  Lines       59401    59445      +44     
==========================================
+ Hits        39767    39810      +43     
- Misses      19634    19635       +1     

Flag	Coverage Δ
unit	55.19% <97.82%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.