Skip no-op EAGLE sampling renormalization (#31294)
Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
This commit is contained in:
@@ -683,6 +683,7 @@ def eagle_sample(
|
|||||||
next_token_logits / expanded_temperature, dim=-1
|
next_token_logits / expanded_temperature, dim=-1
|
||||||
) # (bs * num_draft_tokens, vocab_size)
|
) # (bs * num_draft_tokens, vocab_size)
|
||||||
maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
|
maybe_detect_nan(target_probs, "v2 verify: target_probs after softmax")
|
||||||
|
if sampling_info.need_top_k_sampling:
|
||||||
target_probs = top_k_renorm_prob(
|
target_probs = top_k_renorm_prob(
|
||||||
target_probs,
|
target_probs,
|
||||||
torch.repeat_interleave(
|
torch.repeat_interleave(
|
||||||
@@ -690,6 +691,7 @@ def eagle_sample(
|
|||||||
),
|
),
|
||||||
) # (bs * num_draft_tokens, vocab_size)
|
) # (bs * num_draft_tokens, vocab_size)
|
||||||
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
|
maybe_detect_nan(target_probs, "v2 verify: target_probs after top_k_renorm")
|
||||||
|
if sampling_info.need_top_p_sampling:
|
||||||
target_probs = top_p_renorm_prob(
|
target_probs = top_p_renorm_prob(
|
||||||
target_probs,
|
target_probs,
|
||||||
torch.repeat_interleave(
|
torch.repeat_interleave(
|
||||||
|
|||||||
Reference in New Issue
Block a user