Kimi-K3 support for speculative-decoding draft training

modelopt/torch/speculative/plugins/modeling_final_norm.py

@@ -90,6 +90,10 @@ def extra_repr(self):
     "deepseek_v3": "rmsnorm",
     "kimi_k2": "rmsnorm",  # Kimi-K2 / K2-Thinking (DeepSeek-V3 arch) report model_type "kimi_k2"
     "kimi_k25": "rmsnorm",  # Kimi-K2.5 / K2.6 / K2.7 all report model_type "kimi_k25"
+    # Kimi-K3 is a VLM: its OUTER config reports model_type "kimi_k3", but the text
+    # backbone this table is keyed on reports "kimi_linear". Keying on "kimi_k3" would
+    # silently miss and drop the final norm.
+    "kimi_linear": "rmsnorm",
     # M3's final norm is always gemma-style; map it here too so a config that lost its
     # use_gemma_norm flag still gets the correct flavor instead of silently dropping the +1.
     "minimax_m3_vl_text": "gemma_rmsnorm",

modelopt/torch/utils/loss_mask.py

@@ -79,20 +79,34 @@ def get_loss_mask_recovery(tokenizer) -> LossMaskRecovery | None:
 
 _KIMI_ROLE_MARKERS = ("<|im_user|>", "<|im_assistant|>", "<|im_system|>")
 
+# Kimi-K3's XTML structural tags (see the Kimi-K3 section below). Declared here
+# because ``_kimi_detect`` uses them to hand K3 tokenizers off to the K3 recovery.
+_K3_MARKERS = ("<|open|>", "<|close|>", "<|sep|>", "<|end_of_msg|>")
 
-def _kimi_detect(tokenizer) -> bool:
-    """Whether ``tokenizer`` defines Kimi's chat role markers as real tokens."""
+
+def _has_all_tokens(tokenizer, tokens) -> bool:
+    """Whether ``tokenizer`` maps every one of ``tokens`` to a real (non-unk) id."""
     unk = getattr(tokenizer, "unk_token_id", None)
     try:
-        ids = [
-            tokenizer.convert_tokens_to_ids(t)
-            for t in (*_KIMI_ROLE_MARKERS, "<|im_middle|>", "<|im_end|>")
-        ]
+        ids = [tokenizer.convert_tokens_to_ids(t) for t in tokens]
     except Exception:
         return False
     return all(i is not None and i != unk for i in ids)
 
 
+def _kimi_detect(tokenizer) -> bool:
+    """Whether ``tokenizer`` defines Kimi's chat role markers as real tokens.
+
+    K3 keeps the K2 ``<|im_*|>`` markers for back-compat but *also* defines the XTML
+    structural tags that classic Kimi lacks. When those are present the tokenizer is
+    K3, so defer to the ``kimi_k3`` recovery whose ``compute`` understands the XTML
+    turn layout; matching here would silently produce an empty mask.
+    """
+    if not _has_all_tokens(tokenizer, (*_KIMI_ROLE_MARKERS, "<|im_middle|>", "<|im_end|>")):
+        return False
+    return not _has_all_tokens(tokenizer, _K3_MARKERS)
+
+
 def _kimi_compute(tokenizer, input_ids) -> torch.Tensor:
     """Recover the assistant-content mask from already-tokenized Kimi chat ids.
 
@@ -136,3 +150,77 @@ def _kimi_compute(tokenizer, input_ids) -> torch.Tensor:
 register_loss_mask_recovery(
     LossMaskRecovery(name="kimi", detect=_kimi_detect, compute=_kimi_compute)
 )
+
+
+# ---------------------------------------------------------------------------
+# Kimi-K3
+#
+# K3 replaces the K2/K2.5 <|im_*|> turn markers with an XTML tag format:
+#   <|open|> {tag} {attrs..} <|sep|> {content} <|close|> {tag} <|sep|>  [<|end_of_msg|>]
+# Tag names and attribute values (including the role) are ORDINARY text tokens;
+# only open/close/sep/end_of_msg are special tokens. An assistant turn reads
+#   <|open|> message role assistant <|sep|> {content} <|close|> message <|sep|>
+# and its content may nest further tags (``think``, ``response``), so the content
+# span is found by tracking open/close depth rather than by scanning for the next
+# marker.
+# ---------------------------------------------------------------------------
+
+
+def _k3_detect(tokenizer) -> bool:
+    """Whether ``tokenizer`` defines K3's XTML structural markers as real tokens."""
+    return _has_all_tokens(tokenizer, _K3_MARKERS)
+
+
+def _k3_compute(tokenizer, input_ids) -> torch.Tensor:
+    """Recover the assistant-content mask from already-tokenized K3 chat ids.
+
+    Marks the content of each ``message role=assistant`` turn -- from the token after
+    the message-open ``<|sep|>`` up to (excluding) the matching ``<|close|>`` -- which
+    includes the nested ``think``/``response`` sub-tags the model generates. Role
+    headers, other roles, and the turn-closing tokens stay unmasked, matching the
+    ``{% generation %}`` span a fast tokenizer would report.
+    """
+    ids = input_ids.tolist() if hasattr(input_ids, "tolist") else list(input_ids)
+    open_id = tokenizer.convert_tokens_to_ids("<|open|>")
+    close_id = tokenizer.convert_tokens_to_ids("<|close|>")
+    sep_id = tokenizer.convert_tokens_to_ids("<|sep|>")
+
+    n = len(ids)
+    mask = [0] * n
+    i = 0
+    while i < n:
+        if ids[i] != open_id:
+            i += 1
+            continue
+        # The header tokens sit between this <|open|> and its <|sep|>.
+        j = i + 1
+        while j < n and ids[j] != sep_id:
+            j += 1
+        if j >= n:
+            break
+        header = tokenizer.decode(ids[i + 1 : j]).lower()
+        if "message" not in header or "assistant" not in header:
+            i = j + 1
+            continue
+        # Content = [after this <|sep|>, matching <|close|>), skipping nested tags.
+        start = j + 1
+        depth = 1
+        k = start
+        while k < n:
+            if ids[k] == open_id:
+                depth += 1
+            elif ids[k] == close_id:
+                depth -= 1
+                if depth == 0:
+                    break
+            k += 1
+        for t in range(start, k):
+            mask[t] = 1
+        i = k + 1
+
+    return torch.tensor(mask, dtype=torch.long)
+
+
+register_loss_mask_recovery(
+    LossMaskRecovery(name="kimi_k3", detect=_k3_detect, compute=_k3_compute)
+)

tests/unit/torch/speculative/plugins/test_modeling_final_norm.py

@@ -37,6 +37,11 @@
         ("deepseek_v3", "rmsnorm"),
         ("kimi_k2", "rmsnorm"),
         ("kimi_k25", "rmsnorm"),
+        # Kimi-K3's text backbone reports "kimi_linear"; the outer VLM config's
+        # "kimi_k3" is deliberately NOT a key, since this table is keyed on the
+        # resolved text config.
+        ("kimi_linear", "rmsnorm"),
+        ("kimi_k3", None),
         # gpt_oss is intentionally DISABLED: its RMSNorm forward/weight dtype differs from the
         # Llama variant we reuse, so it must not resolve to a norm type until a matching class
         # is added.

tests/unit/torch/utils/test_loss_mask.py

@@ -142,6 +142,106 @@ def test_kimi_mask_accepts_list_input():
     assert mask.tolist() == [0, 0, 0, 1, 1, 0]
 
 
+# ---------------------------------------------------------------------------
+# Kimi-K3 (XTML chat format)
+# ---------------------------------------------------------------------------
+
+_K3_MARKER_IDS = {"<|open|>": 10, "<|close|>": 11, "<|sep|>": 12, "<|end_of_msg|>": 13}
+# Header/content tokens are ORDINARY text tokens in K3, so the recovery decodes them.
+_K3_WORDS = {100: "message", 101: "role", 102: "assistant", 103: "user", 104: "think"}
+
+
+class FakeK3Tokenizer:
+    """K3 tokenizer: XTML structural markers, plain-text tag names."""
+
+    is_fast = False
+    unk_token_id = 999
+
+    def convert_tokens_to_ids(self, token):
+        return _K3_MARKER_IDS.get(token, self.unk_token_id)
+
+    def decode(self, ids):
+        return " ".join(_K3_WORDS.get(i, "") for i in ids)
+
+
+class FakeK3WithLegacyMarkersTokenizer(FakeK3Tokenizer):
+    """Real K3 keeps the K2 ``<|im_*|>`` markers too; K3 must still win."""
+
+    def convert_tokens_to_ids(self, token):
+        if token in _K3_MARKER_IDS:
+            return _K3_MARKER_IDS[token]
+        return _MARKER_IDS.get(token, self.unk_token_id)
+
+
+def _k3_turn(role_id, content_ids):
+    # <|open|> message role {role} <|sep|> {content} <|close|> message <|sep|>
+    return [
+        _K3_MARKER_IDS["<|open|>"],
+        100,  # "message"
+        101,  # "role"
+        role_id,
+        _K3_MARKER_IDS["<|sep|>"],
+        *content_ids,
+        _K3_MARKER_IDS["<|close|>"],
+        100,
+        _K3_MARKER_IDS["<|sep|>"],
+    ]
+
+
+def test_k3_recovery_is_registered():
+    recovery = get_loss_mask_recovery(FakeK3Tokenizer())
+    assert recovery is not None
+    assert recovery.name == "kimi_k3"
+
+
+def test_k3_tokenizer_does_not_match_the_k2_recovery():
+    """K3 defines both marker sets; the K2 recovery must defer or the mask is empty."""
+    recovery = get_loss_mask_recovery(FakeK3WithLegacyMarkersTokenizer())
+    assert recovery is not None
+    assert recovery.name == "kimi_k3"
+
+
+def test_k3_mask_marks_only_assistant_content():
+    tok = FakeK3Tokenizer()
+    ids = _k3_turn(103, [200]) + _k3_turn(102, [300, 301])  # user turn, assistant turn
+    mask = get_loss_mask_recovery(tok).compute(tok, torch.tensor(ids))
+
+    marked = {i for i, v in enumerate(mask.tolist()) if v == 1}
+    assert marked == {i for i, v in enumerate(ids) if v in (300, 301)}
+
+
+def test_k3_mask_includes_nested_tags():
+    """Assistant content nests ``think``/``response`` sub-tags; they belong to the span."""
+    tok = FakeK3Tokenizer()
+    nested = [
+        200,
+        _K3_MARKER_IDS["<|open|>"],
+        104,  # "think"
+        _K3_MARKER_IDS["<|sep|>"],
+        201,
+        _K3_MARKER_IDS["<|close|>"],
+        104,
+        _K3_MARKER_IDS["<|sep|>"],
+        202,
+    ]
+    ids = _k3_turn(102, nested)
+    mask = get_loss_mask_recovery(tok).compute(tok, torch.tensor(ids)).tolist()
+
+    # Content spans everything between the header <|sep|> and the matching <|close|>.
+    start = 5
+    end = len(ids) - 3
+    assert mask[:start] == [0] * start
+    assert mask[start:end] == [1] * (end - start)
+    assert mask[end:] == [0] * (len(ids) - end)
+
+
+def test_k3_mask_accepts_list_input():
+    tok = FakeK3Tokenizer()
+    ids = _k3_turn(102, [300, 301])
+    mask = get_loss_mask_recovery(tok).compute(tok, ids)  # plain list, not a tensor
+    assert mask.tolist() == [0, 0, 0, 0, 0, 1, 1, 0, 0, 0]
+
+
 def test_register_and_lookup_custom_recovery(restore_registry):
     sentinel = object()
 

tools/launcher/examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml

@@ -0,0 +1,137 @@
+# DSpark streaming speculative-decoding training for Kimi-K3 (multi-node).
+# DSpark = the DFlash backbone + a lightweight Markov head + a confidence head,
+# generating a causal block semi-autoregressively; see dspark.yaml for the head
+# and loss config. Runs the shared streaming pipeline
+# (common/eagle3/train_eagle_streaming.sh) with the K3-specific base, draft dims,
+# mask token and sliding window; trained from scratch. A starting point for
+# reproduction — tune node counts, batch, steps and serve limits for your cluster.
+#
+# Kimi-K3 specifics this yaml encodes (each was a silent failure mode):
+#   * Loss mask. K3 replaces the K2/K2.5 <|im_*|> turn markers with an XTML tag
+#     format (<|open|> message role assistant <|sep|> ... <|close|>) and ships a
+#     slow tiktoken tokenizer, so answer_only_loss depends on the `kimi_k3`
+#     loss-mask recovery. Without it the legacy `kimi` recovery matches (K3 keeps
+#     the old markers for back-compat) and returns an ALL-ZERO mask.
+#   * Final norm. K3 is a VLM whose outer config reports model_type "kimi_k3"
+#     while its text backbone reports "kimi_linear"; the final-norm table is keyed
+#     on the text config, so the `kimi_linear` entry is what stops FakeBaseModel
+#     from silently building no norm and reconstructing target logits from an
+#     un-normed hidden.
+#   * EAGLE_CAPTURE_IDS ends at 93 == num_hidden_layers, the TRUE final hidden.
+#     K3 has a block-residual backbone (attn_res_block_size=12) whose output is
+#     produced after the layer loop by attn_res, so capturing that id requires
+#     vllm#50815. Without it the "final" capture is prefix_sum + hidden_states,
+#     which is missing the block-residual combination and skews the KD teacher.
+#   * The DSpark draft does NOT inherit the base dims, so they are set explicitly
+#     below. Note these are draft-side choices, not mirrors of the backbone: K3's
+#     text config carries no rope_theta at all (so the draft's value is a free
+#     hyperparameter), and its FFN is 33792 dense / 3072 MoE, neither of which the
+#     draft copies.
+#   * Serving K3 needs TP8, and this pipeline runs one serve replica per node, so
+#     the serve nodes must have >= 8 GPUs.
+#
+# Run ON the cluster login node (paramiko can't reach it through the login proxy):
+#   export SLURM_HOST=localhost SLURM_ACCOUNT=<your_account> \
+#          SLURM_PARTITION=<multi_node_partition> \
+#          SLURM_HF_LOCAL=<hf_models_dir> \
+#          SLURM_JOB_DIR=<experiments_dir> \
+#          NEMORUN_HOME=$PWD
+#   uv run launch.py --yaml examples/moonshotai/Kimi-K3/hf_streaming_dspark_multi_node.yaml \
+#          identity=$HOME/.ssh/id_ecdsa detach=True --yes
+#
+# The export lands in /scratchspace/export.
+
+job_name: Kimi-K3_DSpark_streaming_multi_node
+pipeline:
+  allow_to_fail: false
+  skip: false
+  note:
+
+  global_vars:
+    hf_model: /hf-local/moonshotai/Kimi-K3
+
+  # Build /scratchspace/data/train.jsonl. Point data.data_path at the full
+  # Spec-Decoding-Dataset-v2 corpus to reproduce.
+  task_0:
+    script: common/eagle3/make_dataset.sh
+    args:
+      - -f modules/Model-Optimizer/examples/dataset/example_data_config.yaml
+      - --full-conversations
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 1
+      ntasks_per_node: 1
+      gpus_per_node: 8
+      container: nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc10
+
+  task_1:
+    script: common/eagle3/train_eagle_streaming.sh
+    args:
+      - --config modules/Model-Optimizer/modelopt_recipes/general/speculative_decoding/dspark.yaml
+      - model.model_name_or_path=<<global_vars.hf_model>>
+      - model.use_fake_base_for_offline=true
+      - model.trust_remote_code=true
+      - data.mode=streaming
+      - data.data_path=/scratchspace/data/train.jsonl
+      - training.output_dir=/scratchspace/dspark
+      - training.training_seq_len=4096
+      - training.disable_tqdm=true
+      - training.ar_validate_steps=500000
+      - training.num_train_epochs=1
+      - training.per_device_train_batch_size=4
+      - training.gradient_accumulation_steps=1
+      - training.save_steps=1000
+      - training.logging_steps=20
+      - training.learning_rate=1.0e-4
+      - training.warmup_steps=2000
+      # K3's slow tokenizer can't emit assistant masks; the `kimi_k3` loss-mask
+      # recovery reconstructs them from token ids (see header).
+      - training.answer_only_loss=true
+      # The vLLM serve container has no tensorboard -> trainer init crash.
+      - training.report_to=none
+      # Draft dims are NOT inherited from the base — set them explicitly or the
+      # draft is silently built with the Qwen3Config defaults.
+      - dflash.dflash_architecture_config.num_hidden_layers=6
+      - dflash.dflash_architecture_config.num_key_value_heads=8
+      - dflash.dflash_architecture_config.intermediate_size=12288
+      # K3's text config exposes no rope_theta, so this is a draft-side choice
+      # rather than a value inherited from the backbone.
+      - dflash.dflash_architecture_config.rope_theta=10000
+      # Semi-AR generation block (dspark.yaml ships 16; the Kimi backbone uses 8).
+      - dflash.dflash_block_size=8
+      # K3 has no dedicated mask token; 163606 is a reserved slot.
+      - dflash.dflash_mask_token_id=163606
+      # Sliding-window attention over the context; the block itself stays
+      # bidirectional. Must be >= dflash_block_size.
+      - dflash.dflash_swa_window_size=1024
+    environment:
+      - HF_MODEL_CKPT: <<global_vars.hf_model>>
+      # 6 evenly-spaced aux capture ids + the true final hidden (93 == the base's
+      # num_hidden_layers). The final id requires vllm#50815 on a block-residual
+      # backbone; see header.
+      - EAGLE_CAPTURE_IDS: "[2,19,36,52,69,86,93]"
+      - SERVE_NODES: "8"
+      - SERVE_TP: "8"
+      - STREAMING_NUM_WORKERS: "4"
+      # K3's custom-modeling base needs trust_remote_code at export and serve.
+      - EXPORT_EXTRA_ARGS: "--trust_remote_code"
+      - SERVE_EXTRA_ARGS: "--trust-remote-code"
+      - SERVE_MAX_MODEL_LEN: "4160"
+      - SERVE_MAX_NUM_SEQS: "16"
+      - SERVE_GPU_MEM_UTIL: "0.9"
+      # K3 is large; allow a long load before the trainer gives up on the serve.
+      - SERVE_READY_TIMEOUT: "6000"
+      - VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: "1200"
+      - VLLM_ENGINE_ITERATION_TIMEOUT_S: "1200"
+      # RDMA transport is UCX (InfiniBand) by default. On AWS EFA, uncomment:
+      # - NIXL_BACKENDS: "LIBFABRIC"
+      # - FI_PROVIDER: "efa"
+      # - NCCL_IB_DISABLE: "1"
+    slurm_config:
+      _factory_: "slurm_factory"
+      nodes: 9
+      ntasks_per_node: 1
+      gpus_per_node: 8
+      # vLLM build with native Kimi-K3 support (vllm/models/kimi_k3) and the
+      # final-layer aux-capture fix (vllm#50815).
+      container: <vllm-image-with-native-kimi-k3>