[Question] Speculative Decoding: Sampling from the draft probs vector

[keyboardAnt]

Description

The speculative sampling kernel introduced via pull request #3373 (sgl-kernel/src/sgl-kernel/csrc/speculative_sampling.cuh) appears to be incomplete. There are clear comments like // FIXME: leverage draft probs indicating that the implementation does not yet sample from the drafter's probability vector, instead defaulting to zeros as placeholder values.

This discrepancy directly undermines speculative decoding’s effectiveness. If not fixed, it significantly reduces the acceptance rate:

• Desired behavior (sampling):
  Acceptance rate = ∑ₓ min(p(x), q(x))
  where p(x) is the target model’s distribution and q(x) is the drafter’s.

• Current behavior (greedy/zeros):
  Acceptance rate = p(x)

---

References

• The kernel file with the FIXME comment:
  sgl-kernel/src/sgl-kernel/csrc/speculative_sampling.cuh as introduced in pull request support speculative decoding kernel in sgl-kernel #3373.
  PR link

• Slack discussion highlighting that draft_probs may be set to zeros:
  eagle_utils.py#L474

---

Questions & Action Items

1. Is the drafter's probability vector intended to be sampled from in the current kernel implementation?
2. If not, is there a timeline or PR planned to update the kernel to properly incorporate sampling from draft probabilities?

---

This issue blocks #9539, and may also block #8581 and #8391.

[keyboardAnt]

It seems that vLLM currently has a similar issue (vllm-project/vllm#16899) but already has a draft fix in progress (vllm-project/vllm#20459).

[keyboardAnt]

It seems that vLLM currently has a similar issue (vllm-project/vllm#16899) but already has a draft fix in progress (vllm-project/vllm#20459).

I benchmarked vLLM’s fix on H100 and saw a slight improvement in the mean TPOT - vllm-project/vllm#20459 (comment)

[hongyanz]

@zhyncs Can you please assign @hongdaxia to solve this issue? He is happy to work together with @keyboardAnt on this issue. I think solving this issue is very important for improving the performance when T > 0.

[Qiaolin-Yu]

@zhyncs Can you please assign @hongdaxia to solve this issue? He is happy to work together with @keyboardAnt on this issue. I think solving this issue is very important for improving the performance when T > 0.

@hongyanz @hongdaxia Sounds great! Thanks for your help!