Skip to content

bugfix: redirect draft-extend placeholder KV write to padding block.#1944

Open
weizhehuang0827 wants to merge 1 commit into
xLLM-AI:mainfrom
weizhehuang0827:fix/draft-extend-kv-overwrite
Open

bugfix: redirect draft-extend placeholder KV write to padding block.#1944
weizhehuang0827 wants to merge 1 commit into
xLLM-AI:mainfrom
weizhehuang0827:fix/draft-extend-kv-overwrite

Conversation

@weizhehuang0827

@weizhehuang0827 weizhehuang0827 commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Description

Summary

  • When dp_enabled=true or enable_atb_spec_kernel=true, draft extend always adds a prev_token row for DP token-count alignment
  • If prev_token_id < 0 (first decode after prefill or all-draft-rejected), the placeholder row overwrites a valid KV cache entry at position-1, permanently degrading speculative acceptance rate
  • Fix: detect placeholder rows and redirect their new_cache_slots to block 0 (reserved padding block), preserving correct KV cache content

Test Results

Verified on DeepSeek-V3.2-w8a8 full model (DP=2, num_speculative_tokens=3, enable_atb_spec_kernel=false):

Case WITH fix WITHOUT fix Delta
seed1 1req 41.97% (175/417) 39.86% (171/429) +2.11%
seed2 1req 16.67% (2/12) 16.67% (2/12) 0 (short request)
seed3 1req 57.02% (130/228) 54.70% (128/234) +2.32%
4req seed1 48.47% (333/687) 48.33% (332/687) +0.14% (drafted identical)

Related Issues

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

Pull Request Checklist

Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.

PR Title and Commit Messages

  • The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Allowed types: feat, bugfix, docs, test, refactor, chore, style, revert, perf, model, build, release.
The subject should use clear English, start with a verb, include at least 4 words, and end with ..

Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit or an equivalent command.
  • I have installed the hooks with pre-commit install.
  • I have run pre-commit run --all-files and fixed any reported issues.

If you are unsure how to set up pre-commit, see the pre-commit documentation.

Self Review

  • I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
  • I have rebased this PR onto the latest main branch.

Build and Test Coverage

  • Tests have been added or updated as needed.
  • CUDA: python setup.py build test has passed on a CUDA machine.
  • NPU: python setup.py build test has passed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

When dp_enabled=true or use_chunked_prefill=true, draft extend always
adds a prev_token row. If prev_token_id < 0 (first decode after prefill
or all-draft-rejected), the placeholder row overwrites a valid KV cache
entry at position-1, permanently degrading speculative acceptance rate.

Fix: detect placeholder rows and redirect their new_cache_slots to
block 0 (reserved padding block), preserving the correct KV cache
content while maintaining DP token-count alignment.

Verified on DeepSeek-V3.2-w8a8 (DP=2, num_speculative_tokens=3):
- seed1 1req: 41.97% vs 39.86% (nofix), +2.11%
- seed3 1req: 16.67% vs 16.67% (identical, short request)
- seed4 1req: 57.02% vs 54.70% (nofix), +2.32%
- 4req seed1:  48.47% vs 48.33% (nofix), drafted count identical (687)
ustcfy
ustcfy approved these changes Jul 15, 2026
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.

2 participants