Skip to content

Fix positioned top-k sampling for speculative generation - #2029

Open
tjansn wants to merge 3 commits into
Blaizzy:mainfrom
tjansn:fix-positioned-top-k-sampling
Open

Fix positioned top-k sampling for speculative generation#2029
tjansn wants to merge 3 commits into
Blaizzy:mainfrom
tjansn:fix-positioned-top-k-sampling

Conversation

@tjansn

@tjansn tjansn commented Aug 26, 2026

Copy link
Copy Markdown

Summary

  • apply top_k inside _PositionedTargetSampler, before temperature scaling
  • keep seeded position-keyed sampling enabled when top_k is non-default
  • add regressions for positioned-sampler selection and top-k filtering

Problem

Seeded speculative generation already uses a position-keyed target sampler so proposal draws cannot perturb target draws. The selection gate excluded every request with non-default top_k, however, and the positioned sampler did not apply top-k itself. That sent seeded top-k generation back through the stateful sampler path, where serial and speculative grouping can consume target randomness differently.

The server sampling path applies top-k before temperature/top-p sampling. This patch ports that same ordering to the positioned target sampler and removes the now-unnecessary top_k == 0 selection restriction.

Validation

  • targeted selection/filtering assertions pass locally
  • Python compilation passes
  • a local Qwen MTP matrix covering two quantization widths, two seeds, opposite serial/speculative arm orders, and five sampling regimes passed 40/40 reports after the fix, including temperature=1, top_p=1, top_k=40
  • seeded completions repeated across opposite arm orders; proposal draws did not perturb target draws

I could not run the full upstream pytest suite in this environment because pytest is not installed in the pinned runtime.

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