Speculative Acceptance and Correction

Problem

Verify a speculative draft block of N token proposals. For proposal x at index i, accept exactly when acceptance_uniforms[i] < min(1, target_probs[i][x]/draft_probs[i][x]). At the first rejection, sample one correction from normalized max(target_probs[i]−draft_probs[i],0), append it after accepted tokens, and stop.

If all proposals are accepted, append one bonus sampled from target_probs[N]. Sampling is deterministic inverse CDF in ascending token-ID order: choose the first cumulative probability strictly greater than sample_uniform. Return {"tokens": accepted prefix plus correction/bonus, "accepted": accepted count, "rejectedAt": rejection index or −1}. There is no EOS handling in this block-level exercise.

def speculative_block(proposals, target_probs, draft_probs, acceptance_uniforms, sample_uniform):

Constraints

• 0 ≤ N ≤ 16; 1 ≤ V ≤ 16. target_probs is (N+1) × V; draft_probs is N × V.
• Each row is a finite nonnegative normalized distribution; every proposed token has positive draft probability.
• acceptance_uniforms has N values and sample_uniform is a scalar, all in [0,1).
• Inputs satisfy these constraints; invalid-input handling is not graded. Use only Python standard library and NumPy.

Examples

Arguments: [
  [
    0
  ],
  [
    [
      0.5,
      0.5
    ],
    [
      0.2,
      0.8
    ]
  ],
  [
    [
      0.5,
      0.5
    ]
  ],
  [
    0.9
  ],
  0.5
]
Expected result: {
  "tokens": [
    0,
    1
  ],
  "accepted": 1,
  "rejectedAt": -1
}

Arguments: [
  [
    0
  ],
  [
    [
      0.2,
      0.8
    ],
    [
      0.5,
      0.5
    ]
  ],
  [
    [
      0.8,
      0.2
    ]
  ],
  [
    0.5
  ],
  0.2
]
Expected result: {
  "tokens": [
    1
  ],
  "accepted": 0,
  "rejectedAt": 0
}