From 4c9f1321a33b89d4d0f117aeaf62f2de953e10db Mon Sep 17 00:00:00 2001 From: At1a8 Date: Sun, 13 Sep 2026 01:41:43 +0000 Subject: [PATCH 1/3] fix simple acc issue --- python/sglang/srt/speculative/eagle_utils.py | 28 +++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 616f808d7ba6..8df33e31b0fd 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -21,6 +21,7 @@ ) from sglang.srt.runtime_context import get_parallel, get_spec from sglang.srt.utils import ( + get_bool_env_var, is_cpu, is_cuda, is_hip, @@ -42,6 +43,7 @@ _is_cuda = is_cuda() _is_hip = is_hip() +_SPEC_MATCH_SAMPLE = get_bool_env_var("SGLANG_SPEC_MATCH_SAMPLE") and _is_hip _is_npu = is_npu() _is_musa = is_musa() _is_xpu = is_xpu() @@ -374,6 +376,28 @@ def verify_tree_greedy_triton( ) +def select_target_predict( + next_token_logits: torch.Tensor, + sampling_info: SamplingBatchInfo, + draft_token_num: int, +) -> torch.Tensor: + if not _SPEC_MATCH_SAMPLE or sampling_info.is_all_greedy: + return torch.argmax(next_token_logits, dim=-1) + + if sampling_info.need_top_p_sampling or sampling_info.need_top_k_sampling: + logger.warning_once( + "SGLANG_SPEC_MATCH_SAMPLE: top_p/top_k truncation has no ROCm kernel; " + "falling back to greedy verify." + ) + return torch.argmax(next_token_logits, dim=-1) + + temperatures = torch.repeat_interleave( + sampling_info.temperatures, draft_token_num, dim=0 + ) + probs = torch.softmax(next_token_logits / temperatures, dim=-1) + return torch.multinomial(probs, num_samples=1).flatten() + + def verify_tree_greedy_func( predicts: torch.Tensor, accept_index: torch.Tensor, @@ -782,7 +806,9 @@ def eagle_sample( # Sample tokens target_predict = None if sampling_info.is_all_greedy or _is_cpu or _is_hip or _is_xpu: - target_predict = torch.argmax(next_token_logits, dim=-1) + target_predict = select_target_predict( + next_token_logits, sampling_info, verify_input.draft_token_num + ) target_predict = target_predict.reshape(bs, verify_input.draft_token_num) predict, accept_index, num_correct_drafts = verify_tree_greedy_func( predicts=predict, # mutable From 77b84983a2a3e0656bbfc37c7f465d69882c33dd Mon Sep 17 00:00:00 2001 From: At1a8 Date: Mon, 14 Sep 2026 07:07:53 +0000 Subject: [PATCH 2/3] rename env --- python/sglang/srt/speculative/eagle_utils.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 8df33e31b0fd..f6014dcf2b16 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -43,7 +43,9 @@ _is_cuda = is_cuda() _is_hip = is_hip() -_SPEC_MATCH_SAMPLE = get_bool_env_var("SGLANG_SPEC_MATCH_SAMPLE") and _is_hip +_SPEC_TEMPERATURE_SAMPLING_TARGET_VERIFY = ( + get_bool_env_var("SGLANG_SPEC_TEMPERATURE_SAMPLING_TARGET_VERIFY") and _is_hip +) _is_npu = is_npu() _is_musa = is_musa() _is_xpu = is_xpu() @@ -381,13 +383,13 @@ def select_target_predict( sampling_info: SamplingBatchInfo, draft_token_num: int, ) -> torch.Tensor: - if not _SPEC_MATCH_SAMPLE or sampling_info.is_all_greedy: + if not _SPEC_TEMPERATURE_SAMPLING_TARGET_VERIFY or sampling_info.is_all_greedy: return torch.argmax(next_token_logits, dim=-1) if sampling_info.need_top_p_sampling or sampling_info.need_top_k_sampling: logger.warning_once( - "SGLANG_SPEC_MATCH_SAMPLE: top_p/top_k truncation has no ROCm kernel; " - "falling back to greedy verify." + "SGLANG_SPEC_TEMPERATURE_SAMPLING_TARGET_VERIFY does not implement " + "top_p/top_k truncation yet; falling back to greedy verify." ) return torch.argmax(next_token_logits, dim=-1) From edfbbed504ceba7aa9154602b4e4fe76299945fe Mon Sep 17 00:00:00 2001 From: At1a8 Date: Mon, 14 Sep 2026 07:33:55 +0000 Subject: [PATCH 3/3] format --- python/sglang/srt/speculative/eagle_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index f6014dcf2b16..ce8a136f7743 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -383,7 +383,7 @@ def select_target_predict( sampling_info: SamplingBatchInfo, draft_token_num: int, ) -> torch.Tensor: - if not _SPEC_TEMPERATURE_SAMPLING_TARGET_VERIFY or sampling_info.is_all_greedy: + if not _SPEC_TEMPERATURE_SAMPLING_TARGET_VERIFY: return torch.argmax(next_token_logits, dim=-1) if sampling_info.need_top_p_sampling or sampling_info.need_top_k_sampling: @@ -807,10 +807,14 @@ def eagle_sample( # Sample tokens target_predict = None - if sampling_info.is_all_greedy or _is_cpu or _is_hip or _is_xpu: + if sampling_info.is_all_greedy or _is_cpu or _is_xpu: + target_predict = torch.argmax(next_token_logits, dim=-1) + elif _is_hip: target_predict = select_target_predict( next_token_logits, sampling_info, verify_input.draft_token_num ) + + if target_predict is not None: target_predict = target_predict.reshape(bs, verify_input.draft_token_num) predict, accept_index, num_correct_drafts = verify_tree_greedy_func( predicts=predict, # mutable