Work on part of speculative decoding in TKW

[tpopp]

This is incomplete, but I tried to leave TODO's for the obvious incomplete portions. It has not been run yet as it was trying to get all the pieces outlined.

Handles the following:

last_offset = cur_prob_offset
q = target_probs[bx, last_offset]
p = draft_probs[bx, last_offset] if num_accepted_tokens !=
num_speculative_tokens - 1 else torch.zeros_like(q)
relu_diff = torch.relu(q - p)
sum_relu = relu_diff.sum()

u = coin * sum_relu
sampled_id = d - 1
aggregate = 0.0
for i in range(d):
    val = relu_diff[i]
    if val <= 0:
	continue
    aggregate += val
    if aggregate > u:
	sampled_id = i
	break