feat: pipeline-parallel hidden-state export, a patched vLLM image, and the K3 DSpark MLA draft

configs/draft_models/kimi_k3_dspark_mla.json

@@ -0,0 +1,54 @@
+{
+  "architectures": [
+    "K3DSparkModel"
+  ],
+  "model_type": "k3_dspark",
+  "hidden_size": 7168,
+  "intermediate_size": 14336,
+  "num_hidden_layers": 5,
+  "num_attention_heads": 64,
+  "num_key_value_heads": 64,
+  "q_lora_rank": 1536,
+  "kv_lora_rank": 512,
+  "qk_nope_head_dim": 128,
+  "qk_rope_head_dim": 64,
+  "v_head_dim": 128,
+  "mla_use_output_gate": false,
+  "vocab_size": 163840,
+  "rms_norm_eps": 1e-05,
+  "max_position_embeddings": 1048576,
+  "rope_theta": 50000.0,
+  "num_target_layers": 5,
+  "target_hidden_size": 7168,
+  "target_num_hidden_layers": 93,
+  "target_layer_ids": [
+    7,
+    31,
+    47,
+    63,
+    87
+  ],
+  "fc_norm": true,
+  "mask_token_id": 163837,
+  "bos_token_id": 163584,
+  "eos_token_id": 163586,
+  "pad_token_id": 163839,
+  "markov_rank": 256,
+  "markov_head_type": "vanilla",
+  "enable_confidence_head": true,
+  "confidence_head_with_markov": true,
+  "tie_word_embeddings": false,
+  "draft_vocab_size": 163840,
+  "_torchspec_version": "0.1.0",
+  "torch_dtype": "bfloat16",
+  "rope_parameters": {
+    "rope_type": "yarn",
+    "factor": 32.0,
+    "original_max_position_embeddings": 32768,
+    "rope_theta": 50000.0,
+    "beta_fast": 32,
+    "beta_slow": 1,
+    "mscale": 1.0,
+    "mscale_all_dim": 1.0
+  }
+}

configs/vllm_kimi_k3_dspark_stage1.yaml

@@ -0,0 +1,137 @@
+# Kimi-K3 DSpark draft training — 8-node setup
+#
+# Node layout (40 GPUs across 8 nodes):
+#   - 4 nodes (32 GPUs) for inference: two vLLM engines, each TP=8 across 2
+#     pipeline stages
+#   - 2 nodes (8 GPUs) for training (FSDP FULL_SHARD)
+#
+# Pipeline-parallel inference requires a vLLM built from
+# patches/vllm/<image-tag>/, because every pipeline stage has to return the aux
+# hidden states it owns; released vLLM captures on the last stage only. See
+# docker/vllm/<image-tag>/Dockerfile.
+#
+# Usage:
+#   python -m torchspec.train_entry --config configs/vllm_kimi_k3_dspark_stage1.yaml
+
+model:
+  target_model_path: ???  # set to your Kimi-K3 checkpoint (local dir or HF hub id)
+  draft_model_config: configs/draft_models/kimi_k3_dspark_mla.json
+  target_model_backend: vllm
+  trust_remote_code: true
+  embedding_key: language_model.model.embed_tokens.weight
+  lm_head_key: language_model.lm_head.weight
+  norm_key: language_model.model.norm.weight
+
+dataset:
+  train_data_path: ???  # set to your dataset path (HF hub id, local dir, or jsonl)
+  eval_data_path: ???  # set to your eval dataset path, or null to skip eval
+  eval_interval: 50
+  eval_prompt_key: conversations
+  prompt_key: conversations
+  renderer: kimi-k3
+  chat_template: null
+  defer_tokenization: false
+  last_turn_loss_only: true
+  shuffle_dataset: true
+  length_group_size: 1
+  # DSpark needs at least two complete seven-token diffusion blocks.
+  min_loss_tokens: 14
+
+training:
+  max_seq_length: 32768
+  seed: 42
+  micro_batch_size: 1
+  # Eight training ranks x one sequence x 16 microbatches = 128.
+  draft_accumulation_steps: 16
+  learning_rate: 5.0e-5
+  warmup_ratio: 0.04
+  lr_decay_style: cosine
+  min_lr: 0.0
+  weight_decay: 0.0
+  max_grad_norm: 1.0
+  ttt_length: 7
+  num_epochs: 5
+  save_interval: 1000
+  save_per_epoch: true
+  max_concurrent_batches: 1
+  gradient_checkpointing: true
+  attention_backend: flex_attention
+  training_num_nodes: 2
+  training_num_gpus_per_node: 4
+  fsdp_strategy: FULL_SHARD
+  distributed_timeout_minutes: 60
+  # Set the socket interfaces to your cluster's NICs. The MNNVL/IMEX entries are
+  # only needed on NVLink-switched multi-node systems.
+  # train_env_vars: '{"NCCL_SOCKET_IFNAME":"<iface>","GLOO_SOCKET_IFNAME":"<iface>","NCCL_MNNVL_ENABLE":"1","NCCL_CUMEM_ENABLE":"1","NVIDIA_IMEX_CHANNELS":"all"}'
+
+  dflash_block_size: 7
+  dspark_num_anchors: 512
+  dspark_num_target_layers: 5
+  dspark_loss_decay_gamma: 4.0
+  dspark_ce_loss_alpha: 0.1
+  dspark_l1_loss_alpha: 0.9
+  dspark_confidence_head_alpha: 1.0
+
+  # Keeping both checkpoint inputs disabled prevents an accidental warm start.
+  load_path: null
+  continual_training: false
+
+inference:
+  inference_engine_type: vllm
+  inference_num_gpus: 32
+  inference_num_gpus_per_engine: 16
+  inference_num_gpus_per_node: 4
+  inference_batch_size: 4
+  inference_buffer_threshold: 8
+  inference_fetch_batch: 8
+  max_sample_pool_size: 32
+  store_last_hidden_states: true
+  # Must match target_layer_ids in the draft config above.
+  aux_hidden_states_layers: [7, 31, 47, 63, 87]
+  last_hidden_states_prenorm: true
+
+  vllm:
+    # 16 GPUs per engine at pp_size=2 gives tp_size=8 with two pipeline stages.
+    tp_size: 8
+    pp_size: 2
+    nnodes: 4
+    init_timeout: 3600
+    dist_timeout: 1800
+    mem_fraction_static: 0.85
+    extra_args:
+      load_format: fastsafetensors
+      max_num_batched_tokens: 32768
+      enable_chunked_prefill: true
+      async_scheduling: true
+      enforce_eager: true
+      kv_cache_memory_bytes: 34359738368
+
+mooncake:
+  master_server_address: null
+  metadata_server: null
+  protocol: rdma  # tcp works but is much slower for hidden-state transfer
+  # device_name: mlx5_0,mlx5_1  # Please change this to correct network devices.
+  enable_gpu_direct: false
+  hidden_dim: 7168
+  enable_hard_pin: true
+  kv_lease_ttl_s: 180.0
+  get_retry_max_wait_seconds: 180.0
+  global_segment_size: 512GB
+  local_buffer_size: 16GB
+
+output_dir: ./outputs/kimi-k3-dspark-stage1
+cache_dir: ./cache/kimi-k3-dspark-stage1
+model_download_dir: null
+
+debug:
+  save_debug_train_data: null
+  enable_perf_metrics: true
+
+logging:
+  report_to: none
+  wandb_key: null
+  wandb_project: train_kimi_k3_dspark
+  wandb_team: null
+  wandb_group: null
+  wandb_mode: online
+  wandb_random_suffix: true

configs/vllm_qwen3_8b_pp2.yaml

@@ -0,0 +1,71 @@
+# Configuration for train_entry.py with a pipeline-parallel vLLM target
+#
+# GPU allocation:
+#   - 2 GPUs for inference: one engine split across 2 pipeline stages
+#   - 2 GPUs for training (DP/FSDP: model sharded across 2 GPUs)
+#   - Total: 4 GPUs
+#
+# Requires a vLLM built from patches/vllm/<image-tag>/, because every pipeline
+# stage has to return the aux hidden states it owns; released vLLM captures on
+# the last stage only. See docker/vllm/<image-tag>/Dockerfile.
+#
+# Usage:
+#   python -m torchspec.train_entry --config configs/vllm_qwen3_8b_pp2.yaml
+
+model:
+  target_model_path: Qwen/Qwen3-8B
+  trust_remote_code: true
+
+dataset:
+  train_data_path: ../examples/data/sample_conversations.jsonl
+  eval_data_path: ../examples/data/eval_conversations.jsonl
+  eval_interval: 100
+  chat_template: qwen
+  prompt_key: conversations
+
+training:
+  attention_backend: flex_attention
+  micro_batch_size: 1
+  draft_accumulation_steps: 1
+  learning_rate: 1e-4
+  max_concurrent_batches: 1
+  max_grad_norm: 0.5
+  max_seq_length: 16384
+  num_epochs: 1
+  seed: 42
+  training_num_gpus_per_node: 2
+  training_num_nodes: 1
+  ttt_length: 7
+  save_per_epoch: true
+  warmup_ratio: 0.015
+
+inference:
+  inference_engine_type: vllm
+  inference_num_gpus: 2
+  inference_num_gpus_per_engine: 2
+  inference_num_gpus_per_node: 4
+  max_sample_pool_size: 64
+  inference_buffer_threshold: 32
+  inference_batch_size: 8
+  vllm:
+    # The engine's TP degree is derived as
+    # inference_num_gpus_per_engine / pp_size, so 2 GPUs at pp_size=2 run
+    # tp_size=1 with two pipeline stages.
+    pp_size: 2
+    mem_fraction_static: 0.7
+    extra_args:
+      max_num_batched_tokens: 8192
+
+mooncake:
+  master_server_address: null
+  metadata_server: null
+  protocol: tcp
+  global_segment_size: 32GB
+  local_buffer_size: 4GB
+
+output_dir: ./outputs/qwen3-8b-pp2
+cache_dir: ./cache
+model_download_dir: null
+
+debug:
+  save_debug_train_data: null

docker/justfile

@@ -1,3 +1,23 @@
+# Each backend version is a directory under docker/<backend>/<version>/ holding a
+# Dockerfile, paired with patches/<backend>/<version>/ if that version needs
+# patching. Select one with BACKEND plus the matching *_VERSION:
+#
+#   BACKEND=sglang just build
+#   BACKEND=vllm VLLM_VERSION=v0.22.1 just build
+#
+# Since every backend builds to the same IMAGE_REPO:version tag, use
+# ARG_TAG_POSTFIX to keep two backends' images apart in one registry.
+#
+# Pipeline-parallel hidden-state export needs the patched nightly rather than a
+# released vLLM, because every pipeline stage has to return the aux hidden
+# states it owns and released vLLM captures on the last stage only:
+#
+#   BACKEND=vllm VLLM_VERSION=nightly-7794b1e08bf505ff28664515ffaaeeec955ab796 \
+#       ARG_TAG_POSTFIX=-vllm-pp just build
+#
+# That Dockerfile applies patches/vllm/<version>/series in order rather than
+# globbing, so a patched image is reproducible from this repo alone; see the
+# series file for what each patch is for and which ones are upstream candidates.
 BACKEND := env("BACKEND", "sglang")
 SGLANG_VERSION := env("SGLANG_VERSION", "v0.5.14")
 VLLM_VERSION := env("VLLM_VERSION", "v0.22.1")

docker/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/Dockerfile

@@ -0,0 +1,43 @@
+FROM vllm/vllm-openai:nightly-7794b1e08bf505ff28664515ffaaeeec955ab796
+
+WORKDIR /root/
+
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends nvtop rsync dnsutils && \
+    rm -rf /var/lib/apt/lists/*
+
+# Applied in the order given by `series` rather than by shell glob order, so the
+# stack stays explicit as patches are added or renamed. See that file for what
+# each patch is and whether it is an upstream candidate.
+COPY patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/series \
+     patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/*.patch \
+     /tmp/patches/
+RUN cd /usr/local/lib/python3.12/dist-packages && \
+    while read -r p; do \
+      case "$p" in ''|'#'*) continue;; esac; \
+      echo "applying $p"; \
+      patch -p1 --batch --forward < "/tmp/patches/$p" || exit 1; \
+    done < /tmp/patches/series && \
+    rm -rf /tmp/patches
+
+# NOTE: the base image ships its own CUDA-13-linked mooncake-transfer-engine
+# (0.3.10.post2 as of the v0.22.1 image), which satisfies torchspec's
+# >=0.3.10.post1 floor. Do NOT upgrade or reinstall mooncake from PyPI here:
+# the generic wheel links libcudart.so.12 (broken on this CUDA 13 image) and
+# the cuda13 variant's aarch64 wheel needs glibc >= 2.39 (this image is
+# Ubuntu 22.04 / 2.35).
+COPY . /root/torchspec
+# The `fa` extra is deliberately not installed here, unlike in the v0.22.1
+# image. Its flash-attn-4 pins nvidia-cutlass-dsl==4.6.0.dev0 against this
+# image's 4.6.0 and pulls apache-tvm-ffi past the version vLLM's flashinfer was
+# built against, after which every worker aborts while profiling with
+# `tvm::ffi::Error: TypeAttr __ffi_repr__ is already registered for type index
+# 132`. Pinning the base versions instead does not help: the two are mutually
+# exclusive, and the image is not pip-reproducible anyway (its torch pins
+# nvidia-nccl-cu13==2.29.7 while the image ships 2.30.7, so any install of a
+# torch dependant downgrades nccl). The extra only enables the optional
+# flash_attention draft backend, whose import is already guarded and which
+# training.attention_backend does not select by default.
+RUN cd /root/torchspec && pip install --no-cache-dir -e .
+
+WORKDIR /root/torchspec

patches/vllm/nightly-7794b1e08bf505ff28664515ffaaeeec955ab796/series

@@ -0,0 +1,26 @@
+# Patch stack for vllm/vllm-openai:nightly-7794b1e08bf505ff28664515ffaaeeec955ab796
+# (vLLM 0.26.1rc1.dev353+g7794b1e08, main @ 2026-08-05).
+#
+# Applied top to bottom; blank lines and #-comments are skipped. Files under
+# tests/ are deliberately not listed — the installed dist-packages tree has no
+# tests directory.
+#
+# The order below is the order the patches were authored in: vllm_k3.patch and
+# vllm_pp_hidden_states.patch both edit vllm/models/kimi_k3/nvidia/model.py, and
+# the latter is generated against the tree the former produces. Either order
+# applies cleanly today — patch absorbs the ~80-line shift as an offset — so
+# listing them here is about not depending on that, and on glob order.
+#
+# Upstream intent differs per patch. Keep them separate so the upstream-candidate
+# ones can be format-patch'd out, and so a rebase onto a newer nightly shows
+# which hunks upstream has absorbed (a shrinking patch is the progress signal).
+
+# Ours indefinitely: Kimi-K3 model + env flag, not upstreamable as-is.
+vllm_k3.patch
+
+# Upstream candidate: per-PP-stage aux hidden-state capture, the prefill-only
+# extract path, and the PP completion barrier. Forward-ported onto this nightly
+# from an earlier working tree based on vLLM PR #50000 (head 0498dc7ea9204f18).
+# Adds SupportsPP to ExtractHiddenStatesModel, which is what lets
+# pipeline_parallel_size > 1 past ModelConfig.verify_with_parallel_config.
+vllm_pp_hidden_states.patch