Kimi-K3-speculator.dspark

RedHatAI/Kimi-K3-speculator.dspark

This is a DSpark speculator model for moonshotai/Kimi-K3.

Training Details

This model was trained using the Speculators library. Training compute for this model was generously provided by Verda, a leading cloud platform for AI training and inference. Verda runs its own data centers and offers on-demand NVIDIA GPU clusters with NVLink and InfiniBand interconnect, backed by an in-house AI Lab that works directly with open-source projects on large-scale training and inference optimization.

Kimi-K3 does not fit on a single GB300 node, so training used a multi-node setup with hidden-state extraction and draft training on disjoint node groups, streaming hidden states between them through a Mooncake store. Two separate virtual environments are required:

• Extraction: vllm, hs_connectors, and mooncake-transfer-engine
• Training: speculators, hs_connectors, and mooncake-transfer-engine

<details> <summary>Commands</summary>

Prepare data

  # In virtual environment with speculators installed
  python scripts/prepare_data.py \
    --model moonshotai/Kimi-K3 \
    --trust-remote-code \
    --data ./data.jsonl \
    --output ./runs/kimi_k3_dspark/data \
    --seq-length 8192 \
    --num-preprocessing-workers 48 \
    --minimum-valid-tokens 16

Launch the Mooncake master

One per run, anywhere both node groups can reach.

  mooncake_master --rpc_port 50051 --metrics_port 9003 \
    --rpc_thread_num 8 --enable_disk_eviction=false --logtostderr=true

Launch extraction (both nodes of the TP8 pair)

NODE_RANK=0 on the head node, 1 on the other; rank 1 runs --headless.

  # In (separate) virtual environment with vllm installed
  vllm_venv/bin/python scripts/launch_vllm.py moonshotai/Kimi-K3 \
    --hidden-states-backend mooncake \
    --mooncake-master <host>:50051 \
    --mooncake-metadata-server P2PHANDSHAKE \
    --mooncake-protocol tcp \
    --mooncake-global-segment-gib 32 \
    --mooncake-local-buffer-gib 4 \
    --mooncake-writer-threads 4 \
    --target-layer-ids 24 48 72 88 92 \
    --trust-remote-code \
    -- \
    --served-model-name moonshotai/Kimi-K3 \
    --trust-remote-code \
    --load-format fastsafetensors \
    --moe-backend auto \
    --all2all-backend flashinfer_nvlink_one_sided \
    --enable-expert-parallel \
    --gpu-memory-utilization 0.95 \
    --compilation-config '{"pass_config":{"fuse_allreduce_rms":false}}' \
    --tensor-parallel-size 8 \
    --nnodes 2 \
    --node-rank <0|1> \
    --master-addr <head-fabric-ip> \
    --port 8000 \
    --max-model-len 8193 \
    --max-num-seqs 64 \
    --max-num-batched-tokens 32768 \
    --kv-cache-dtype auto \
    --attention-config '{"mla_prefill_backend":"TRTLLM_RAGGED","use_prefill_query_quantization":false}' \
    --no-enable-prefix-caching \
    --language-model-only

Launch training

Must be run once the extractor reports healthy on /health.

  # In virtual environment with speculators installed
  torchrun --standalone --nproc-per-node 4 \
    scripts/train.py \
    --verifier-name-or-path moonshotai/Kimi-K3 \
    --trust-remote-code \
    --draft-config k3_draft_layer_config.json \
    --data-path ./runs/kimi_k3_dspark/data \
    --save-path ./runs/kimi_k3_dspark/checkpoints \
    --draft-vocab-size 163840 \
    --mask-token-id 163837 \
    --epochs 1 \
    --checkpoint-freq 0.1 \
    --total-seq-len 8192 \
    --train-data-ratio 0.999 \
    --speculator-type dspark \
    --target-layer-ids 24 48 72 88 92 \
    --block-size 8 \
    --max-anchors 1024 \
    --dflash-decay-gamma 4.0 \
    --markov-rank 256 \
    --markov-head-type vanilla \
    --enable-confidence-head \
    --confidence-head-with-markov \
    --confidence-head-alpha 1.0 \
    --loss-fn '{"ce":0.1,"tv":0.9}' \
    --optimizer muon \
    --lr 1e-4 \
    --scheduler-type cosine \
    --scheduler-warmup-ratio 0.03 \
    --hidden-states-backend mooncake \
    --mooncake-master <host>:50051 \
    --mooncake-metadata-server P2PHANDSHAKE \
    --mooncake-protocol tcp \
    --mooncake-global-segment-gib 0 \
    --mooncake-local-buffer-gib 4 \
    --mooncake-writer-threads 4 \
    --vllm-endpoint http://<head-fabric-ip>:8000/v1 \
    --on-missing generate \
    --on-generate delete \
    --request-timeout 900 \
    --max-retries 5 \
    --generation-validation-retries 2 \
    --max-consecutive-generation-failures 20 \
    --num-workers 2 \
    --prefetch-factor 2 \
    --log-freq 20

All four components can be run together on a 3-node Slurm allocation with run.sbatch from the example directory.

</details>

Model Specifications

Deployment

  docker run --gpus all \
  --privileged --ipc=host -p 8000:8000 \
  -v ~/.cache/huggingface:/root/.cache/huggingface \
  -e GLOO_SOCKET_IFNAME=$IFACE_NAME \
  -e NCCL_SOCKET_IFNAME=$IFACE_NAME \
  -e VLLM_ALLREDUCE_USE_FLASHINFER=1 \
  -e VLLM_ENGINE_READY_TIMEOUT_S=3600 \
  -e VLLM_USE_V2_MODEL_RUNNER=1 \
  -e VLLM_USE_RUST_FRONTEND=1 \
  vllm/vllm-openai:latest moonshotai/Kimi-K3 \
  --trust-remote-code \
  --gpu-memory-utilization 0.95 \
  --tensor-parallel-size 16 \
  --nnodes 2 \
  --node-rank 0 \
  --master-addr $HEAD_IP \
  --load-format fastsafetensors \
  --no-enable-flashinfer-autotune \
  --max-model-len 1048576 \
  --kv-cache-dtype fp8 \
  --attention-config '{"use_prefill_query_quantization":true,"mla_prefill_backend":"TOKENSPEED_MLA"}' \
  --enable-prefix-caching \
  --attention-backend TOKENSPEED_MLA \
  --prefix-match-unit 128 \
  --reasoning-parser kimi_k3 \
  --language-model-only \
  --speculative-config '{"model":"RedHatAI/Kimi-K3-speculator.dspark", "num_speculative_tokens":8, "method": "dspark", "draft_sample_method": "probabilistic", "rejection_sample_method": "block"}'  

Acceptance Rates

Per-position token acceptance rates across datasets, at 8 draft tokens per step:

Long-context acceptance

Acceptance rates on LongBench dataset by sub-domain and input context length, from 2K to 20K tokens.

<details> <summary>Full breakdown</summary> </details>

Performance Eval