Kimi-K3.1-Pro-DSpark

DSpark — draft-модель от LithosAI для lossless speculative decoding с Kimi-K3.1-Pro. Это голова на 2,25 млрд параметров (BF16), которая читает уже посчитанные hidden states целевой модели, параллельно предлагает блок токенов, а target проверяет его за один forward pass. Сохранённые токены берутся из распределения target, поэтому ответы точно такие же, как у Kimi-K3.1-Pro; меняется только скорость.

🦾 Ключевые особенности

• 💥 Блок из 8 токенов: target проверяет 8 токенов за шаг — текущий плюс 7 предложенных.
• 🦾 5 dense decoder-слоёв, hidden size 7168, GQA с 64 query- и 16 KV-головами; внутри draft-блока внимание двунаправленное.
• 👁 Читает hidden states целевых слоёв 7, 23, 67, 83 и 91 из 93 и сливает их одной линейной проекцией.
• ♟ Низкоранговый Markov-голова (rank 256) и confidence-голова для оценки принятия по позициям.
• ⬇ Общие с target token embedding и LM head; отдельно в репозитории не хранятся.
• ✨ YaRN: factor 16 поверх 65 536, сконфигурировано до 1 048 576 позиций.

😎 Пример

python -m sglang.launch_server \
  --model-path moonshotai/Kimi-K3.1-Pro \
  --tp-size 8 \
  --kv-cache-dtype fp8_e4m3 \
  --speculative-algorithm DSPARK \
  --speculative-draft-model-path LithosAI/Kimi-K3.1-Pro-DSpark \
  --speculative-dspark-block-size 7 \
  --speculative-num-draft-tokens 8

Подробнее

Средняя длина принятия — 3,448 на 10 бенчмарках при block size 8; на AA-LCR с промптами ~95K токенов — 3,281.
Запуск в карточке: Kimi-K3.1-Pro на одном узле 8× B300, tensor parallel 8, FP8 KV cache.
Draft-модель не самостоятельная: нужен target Kimi-K3.1-Pro и его hidden states; token embedding и LM head общие и не входят в репозиторий.
Обучение: многоходовые agentic coding-сессии от Kimi-K3.1-Pro, 5 drafted tokens на блок, дистилляция в final hidden state target, 2 эпохи, 8× B300.