spec_generate in dflash.py reads only 6 of 7 noise rows

#1
by zxytim - opened

Hi, I think there's a small but costly bug in the reference spec_generate in dflash.py.

The draft logits are sliced as:

draft_logits = target.lm_head(
    self(... )[:, -block_size + 1 :, :]
)

so only rows 1..6 of the 7-row noise block produce proposals (6 draft tokens, verify width 7+bonus). But the model card says "block size 7 draft tokens (verify width 8, including the target bonus token)" β€” i.e. the anchor row (block position 0) should also predict the next token, giving 7 draft tokens per iteration. All 7 noise rows should be read (row j predicts position start+j+1).

With the 6-of-7 readout, proposals come out shifted by one position and measured acceptance drops far below the card's numbers (roughly ~1.7 vs ~4.5 on GSM8K-style prompts in my tests). Switching to all-7-rows readout brings acceptance right into the card's range.

Thank you for catching this. The checkpoint and current SGLang runtime use block_size=gamma=7 proposed draft tokens and an 8-token verify window.

I opened PR #2 to align the reference spec_generate with that contract. The patch:

  • keeps the 7-row noise block;
  • reads all 7 draft logits;
  • verifies anchor + 7 drafts;
  • reserves the complete target-bonus slot at the output-buffer boundary.

A targeted H100 contract probe observed 7 noise -> 6 logits -> verify width 7 before the patch and 7 noise -> 7 logits -> verify width 8 after it. A live SGLang load of this checkpoint independently resolved gamma=7 and verify_num_draft_tokens=8.

Sign up or log in to comment