Skip to content

[WIP] feat: add DFlash2 draft serving support - #1211

Draft
torchspec-bot wants to merge 6 commits into
mainfrom
feat/dflash2-official-node-cudagraph
Draft

[WIP] feat: add DFlash2 draft serving support#1211
torchspec-bot wants to merge 6 commits into
mainfrom
feat/dflash2-official-node-cudagraph

Conversation

@torchspec-bot

@torchspec-bot torchspec-bot commented Aug 23, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add DFlash2 draft-model support using the official grouped dynamic convolution and candidate-selector design
  • dispatch DFlash2 checkpoints through the existing DFLASH runtime
  • connect the draft q distribution to stochastic verification with min(1, p / q) acceptance and residual relu(p - q) sampling
  • support GQA and MLA DFlash2 checkpoints, including Kimi-K3 target-hidden capture and MLA context-cache writes
  • keep Kimi-K3 draft MLA KV in the unified full-retention cache group while carrying per-layer sliding-window metadata on the compute attention objects
  • keep CUDA Graph compatibility and fail fast for unsupported random-sampling backends
  • add focused tests and user-facing configuration documentation

Validation

  • DFlash2 and Kimi-K3 focused tests: 17 passed
  • Kimi-K3 BF16 target plus DFlash2 MLA draft loaded at TP8 and captured CUDA Graph buckets [1, 2, 3, 4]
  • GSM8K, 128 samples, concurrency 1, temperature 1.0, top-p 0.95, top-k 20: mean acceptance length 5.271, median 5.403, proposal acceptance 61.01%
  • GSM8K boxed-or-last-number accuracy: 114/128 (89.06%)
  • pre-commit run --all-files passed

WIP

  • broader topology qualification and performance tuning will follow

torchspec-bot and others added 6 commits August 23, 2026 03:28
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
Signed-off-by: torchspec-bot <torchspec-bot@users.noreply.github.com>
Signed-off-by: torchspec-bot <torchspec-bot@users.noreply.github.com>
Signed-off-by: torchspec-bot <torchspec-bot@users.noreply.github.com>
Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant