Accept or Reject Draft Tokens
Hard · Solved in PyTorch · Machine learning coding practice
Problem
Implement the acceptance step from Accelerating Large Language Model Decoding with Speculative Sampling (Chen et al., 2023).
Target rows target_probs p of shape (K+1, V) and draft rows draft_probs q of shape (K, V) are float64, nonnegative, and sum to one per row. draft_tokens t of shape (K,) is int64 with q_{i, t_i} > 0. The draws uniforms u of shape (K,) and the Python float sample_uniform s all lie in [0, 1); generate no randomness of your own.
Topics: inference, vllm, pytorch
The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Accept or Reject Draft Tokens in the browser and run it against the tests.