[BugFix][v0.20.2rc]Reduce sampling is reconstructed to eliminate all patch behaviors and support DFlash and MTP (#9946)

<!--  Thanks for sending a pull request!

BEFORE SUBMITTING, PLEASE READ
https://docs.vllm.ai/en/latest/contributing/overview.html

-->
### What this PR does / why we need it?
<!--
- Please clarify what changes you are proposing. The purpose of this
section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster
reviews in your PR.

- Please clarify why the changes are needed. For instance, the use case
and bug description.

- Fixes #
-->
The Reduce_sampling optimization is reconstructed to eliminate all patch
behaviors and support the DFlash and MTP.
When sampling optimization is enabled, if speculative decoding is Eagle3
or DFlash, sampling can be optimized for both the main model and the MTP
layer. If the speculative decoding method is MTP, some models can
optimize sampling for both the main model and the MTP layer, while
others can only optimize sampling for the main model.
### Does this PR introduce _any_ user-facing change?
<!--
Note that it means *any* user-facing change including all aspects such
as API, interface or other behavior changes.
Documentation-only updates are not considered user-facing changes.
-->

### How was this patch tested?
<!--
CI passed with new added/existing test.
If it was tested in a way different from regular unit tests, please
clarify how you tested step by step, ideally copy and paste-able, so
that other reviewers can test and check, and descendants can verify in
the future.
If tests were not added, please describe why they were not added and/or
why it was difficult to add.
-->

---------

Signed-off-by: hzx55906 <513464215@qq.com>

Author: hzx55906

tests/ut/spec_decode/test_eagle_proposer.py

@@ -2161,9 +2161,8 @@ def test_intermediate_variables_precision(self):
 class MockDraftModel:
     """Draft model that records prepared forward inputs."""
 
-    def __init__(self, returns_tuple=True, enable_reduce_sample=False, vocab_size=200000):
+    def __init__(self, returns_tuple=True, vocab_size=200000):
         self.returns_tuple = returns_tuple
-        self.enable_reduce_sample = enable_reduce_sample
         self.vocab_size = vocab_size
         self.calls = []
         self.logit_inputs = []
@@ -2187,13 +2186,11 @@ def __call__(self, **kwargs):
             return last_hidden_states, hidden_states
         return last_hidden_states
 
-    def compute_logits(self, sample_hidden_states, enable_reduce_sample=False):
+    def compute_logits(self, sample_hidden_states):
         self.logit_inputs.append(sample_hidden_states.clone())
         token_ids = sample_hidden_states[:, 0].to(torch.long)
         logits = torch.full((sample_hidden_states.shape[0], self.vocab_size), -1000.0)
         logits[torch.arange(sample_hidden_states.shape[0]), token_ids] = 1000.0
-        if self.enable_reduce_sample:
-            logits = logits.argmax(dim=-1)
         return logits
 
     def embed_input_ids(self, input_ids):
@@ -2417,7 +2414,7 @@ def check_mock(self):
         assert sig_name == ["self", "input_ids", "positions", "hidden_states", "inputs_embeds"]
         sig = inspect.signature(RunnerCls.compute_logits)
         sig_name = self.get_param_names(sig)
-        assert sig_name == ["self", "hidden_states", "enable_reduce_sample"]
+        assert sig_name == ["self", "hidden_states"]
 
         import vllm_ascend.ascend_forward_context
 
@@ -2513,7 +2510,17 @@ def get_param_names(self, sig):
         return [p.name for p in sig.parameters.values()]
 
     def test_run_merged_draft_eagle3_decode_prepares_each_forward_input(self):
-        self.proposer.model = MockDraftModel(returns_tuple=True, enable_reduce_sample= True)
+        self.proposer.model = MockDraftModel(returns_tuple=True)
+
+        def compute_draft_token_ids(sample_hidden_states):
+            self.proposer.model.logit_inputs.append(sample_hidden_states.clone())
+            token_ids = sample_hidden_states[:, 0].to(torch.long)
+            logits = torch.full((sample_hidden_states.shape[0], self.proposer.model.vocab_size), -1000.0)
+            logits[torch.arange(sample_hidden_states.shape[0]), token_ids] = 1000.0
+            logits = logits.argmax(dim=-1)
+            return logits
+
+        self.proposer.compute_draft_token_ids = compute_draft_token_ids
         self.proposer.supports_mm_inputs = True
         initial_input_ids = torch.tensor(
             [279, 1196, 374, 8014, 151667, 198, 32313, 11, 151667, 198, 32313, 11],

vllm_ascend/ops/triton/bincount.py

@@ -24,6 +24,7 @@
 from vllm.distributed.parallel_state import get_tp_group
 from vllm.triton_utils import tl, triton
 
+from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ops.triton.triton_utils import get_vectorcore_num
 
 
@@ -80,10 +81,10 @@ def token_bin_counts_and_mask_kernel(
     )
 
     local_token = token - vocab_start_idx
-
     token_in_range = pos_mask & (token >= vocab_start_idx) & (local_token < vocab_size)
 
-    count_ptr = batch_counts_start + local_token * counts_vocab_stride
+    safe_local_token = tl.where(token_in_range, local_token, 0)
+    count_ptr = batch_counts_start + safe_local_token * counts_vocab_stride
     tl.atomic_add(count_ptr, 1, mask=token_in_range)
 
 
@@ -129,8 +130,11 @@ def get_token_bin_counts_and_mask_triton(
     progs = min(core_num, total_blocks)
     grid = (progs, triton.cdiv(total_blocks, progs))
 
-    tp_group = get_tp_group()
-    tp_rank = tp_group.rank_in_group
+    if get_ascend_config().enable_reduce_sample:
+        tp_group = get_tp_group()
+        tp_rank = tp_group.rank_in_group
+    else:
+        tp_rank = 0
     token_bin_counts_and_mask_kernel[grid](
         tokens,
         tokens.stride(0),

vllm_ascend/ops/triton/reject_sample.py

@@ -179,10 +179,10 @@ def rejection_random_sample_kernel(
 
             for pos in range(num_draft_tokens):
                 if not rejected:
-                    token_idx = start_idx + pos
-                    draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
-
                     if ENABLE_REDUCE_SAMPLING:
+                        token_idx = start_idx + pos
+                        draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
+
                         target_prob = 0.0
                         found = False
 
@@ -209,30 +209,47 @@ def rejection_random_sample_kernel(
                                 if current_match_prob > 0.0:
                                     target_prob = current_match_prob
                                     found = True
-                    else:
-                        target_prob = tl.load(target_probs_ptr + token_idx * vocab_size + draft_token_id)
 
-                    if NO_DRAFT_PROBS:
-                        draft_prob = 1.0
-                    else:
-                        vocab_for_draft = global_vocab_size if ENABLE_REDUCE_SAMPLING else vocab_size
-                        draft_prob = tl.load(draft_probs_ptr + token_idx * vocab_for_draft + draft_token_id)
+                        if NO_DRAFT_PROBS:
+                            draft_prob = 1
+                        else:
+                            draft_prob = tl.load(draft_probs_ptr + token_idx * global_vocab_size + draft_token_id)
 
-                    uniform_prob = tl.load(uniform_probs_ptr + token_idx)
+                        uniform_prob = tl.load(uniform_probs_ptr + token_idx)
 
-                    # Acceptance condition
-                    if draft_prob > 0 and target_prob / draft_prob >= uniform_prob:
-                        # Accept
-                        token_id = draft_token_id
-                    else:
-                        # Reject - use recovered token
-                        rejected = True
-                        token_id = tl.load(recovered_token_ids_ptr + token_idx)
+                        # Acceptance condition
+                        if draft_prob > 0 and target_prob / draft_prob >= uniform_prob:
+                            # Accept
+                            token_id = draft_token_id
+                        else:
+                            # Reject - use recovered token
+                            rejected = True
+                            token_id = tl.load(recovered_token_ids_ptr + token_idx)
 
-                    tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id)
+                        tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id)
+                    else:
+                        draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
+                        target_prob = tl.load(target_probs_ptr + (start_idx + pos) * global_vocab_size + draft_token_id)
+                        if NO_DRAFT_PROBS:
+                            draft_prob = 1
+                        else:
+                            draft_prob = tl.load(
+                                draft_probs_ptr + (start_idx + pos) * global_vocab_size + draft_token_id
+                            )
+                        uniform_prob = tl.load(uniform_probs_ptr + start_idx + pos)
+                        # NOTE(woosuk): While the draft probability should never be 0,
+                        # we check it to avoid NaNs. If it happens to be 0, we reject.
+                        if draft_prob > 0 and target_prob / draft_prob >= uniform_prob:
+                            # Accept.
+                            token_id = draft_token_id
+                        else:
+                            # Reject. Use recovered token.
+                            rejected = True
+                            token_id = tl.load(recovered_token_ids_ptr + start_idx + pos)
+                        tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id)
 
             if not rejected:
-                # All tokens accepted - append bonus token
+                # If all tokens are accepted, append the bonus token.
                 bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
                 tl.store(
                     output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens,
@@ -470,38 +487,40 @@ def rejection_random_sample_block_verify_kernel(
     NO_DRAFT_PROBS: tl.constexpr,
     ENABLE_REDUCE_SAMPLING: tl.constexpr,  # Whether using reduce_sampling
     BLOCK_SIZE: tl.constexpr,
-    VOCAB_BLOCK_SIZE: tl.constexpr = 512,
+    SUB_BLOCK: tl.constexpr = 512,
 ):
     block_idx = tl.program_id(0)
     offsets = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
     mask = offsets < vec_len
     is_greedy = tl.load(is_greedy_ptr + offsets, mask, other=1)
     not_greedy_mask = is_greedy == 0
-    start_idxs = tl.where(offsets == 0, 0, tl.load(cu_num_draft_tokens_ptr + offsets - 1, not_greedy_mask))
+    prev_mask = not_greedy_mask & (offsets > 0)
+    prev_end_idxs = tl.load(cu_num_draft_tokens_ptr + offsets - 1, prev_mask, other=0)
+    start_idxs = tl.where(offsets == 0, 0, prev_end_idxs)
     end_idxs = tl.load(cu_num_draft_tokens_ptr + offsets, not_greedy_mask)
     n_num_draft_tokens = end_idxs - start_idxs
 
-    for req_i in range(BLOCK_SIZE):
-        not_greedy = get_element(not_greedy_mask, (req_i,))
-        if not_greedy:
-            pi = 1.0
-            uniform_prob = 1.0
-            last_accepted_token_pos = -1
-            start_idx = get_element(start_idxs, (req_i,))
-            req_idx = block_idx * BLOCK_SIZE + req_i
-            num_draft_tokens = get_element(n_num_draft_tokens, (req_i,))
-
-            for pos in range(num_draft_tokens):
-                token_idx = start_idx + pos
-                draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
+    if ENABLE_REDUCE_SAMPLING:
+        for req_i in range(BLOCK_SIZE):
+            not_greedy = get_element(not_greedy_mask, (req_i,))
+            if not_greedy:
+                pi = 1.0
+                uniform_prob = 1.0
+                last_accepted_token_pos = -1
+                start_idx = get_element(start_idxs, (req_i,))
+                req_idx = block_idx * BLOCK_SIZE + req_i
+                num_draft_tokens = get_element(n_num_draft_tokens, (req_i,))
+
+                for pos in range(num_draft_tokens):
+                    token_idx = start_idx + pos
+                    draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
 
-                if ENABLE_REDUCE_SAMPLING:
                     target_prob = 0.0
                     found = False
 
-                    for v_offset in range(0, vocab_size, VOCAB_BLOCK_SIZE):
+                    for v_offset in range(0, vocab_size, SUB_BLOCK):
                         if not found:
-                            vocab_offsets = v_offset + tl.arange(0, VOCAB_BLOCK_SIZE)
+                            vocab_offsets = v_offset + tl.arange(0, SUB_BLOCK)
                             vocab_mask = vocab_offsets < vocab_size
 
                             candidate_indices = tl.load(
@@ -519,37 +538,115 @@ def rejection_random_sample_block_verify_kernel(
                             if current_match_prob > 0.0:
                                 target_prob = current_match_prob
                                 found = True
+
+                    tmp_uniform_prob = tl.load(uniform_probs_ptr + token_idx)
+                    uniform_prob = uniform_prob * tmp_uniform_prob
+
+                    if NO_DRAFT_PROBS:
+                        draft_prob = 1.0
+                    else:
+                        draft_prob = tl.load(draft_probs_ptr + token_idx * global_vocab_size + draft_token_id)
+
+                    pi = min(pi * target_prob / draft_prob, 1.0)
+                    if draft_prob > 0 and pi >= uniform_prob:
+                        last_accepted_token_pos = pos
+
+                # Store accepted tokens
+                if last_accepted_token_pos > -1:
+                    for pos in range(last_accepted_token_pos + 1):
+                        token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
+                        tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id)
+
+                # Store recovered or bonus token
+                if last_accepted_token_pos + 1 < num_draft_tokens:
+                    # Rejected - store recovered token
+                    recovered_token_id = tl.load(recovered_token_ids_ptr + start_idx + last_accepted_token_pos + 1)
+                    tl.store(
+                        output_token_ids_ptr + req_idx * (max_spec_len + 1) + last_accepted_token_pos + 1,
+                        recovered_token_id,
+                    )
                 else:
+                    # All accepted - store bonus token
+                    bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
+                    tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, bonus_token_id)
+    else:
+        vocab_size = global_vocab_size
+        loop = (vocab_size + SUB_BLOCK - 1) // SUB_BLOCK
+        for req_i in range(BLOCK_SIZE):
+            not_greedy = get_element(not_greedy_mask, (req_i,))
+            if not_greedy:
+                start_idx = get_element(start_idxs, (req_i,))
+                req_idx = block_idx * BLOCK_SIZE + req_i
+                num_draft_tokens = get_element(n_num_draft_tokens, (req_i,))
+                if num_draft_tokens == 0:
+                    bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
+                    tl.store(
+                        output_token_ids_ptr + req_idx * (max_spec_len + 1),
+                        bonus_token_id,
+                    )
+                    continue
+
+                accepted_len = 0
+                prefix_prob = 1.0
+                for pos in range(num_draft_tokens):
+                    token_idx = start_idx + pos
+                    draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
                     target_prob = tl.load(target_probs_ptr + token_idx * vocab_size + draft_token_id)
 
-                tmp_uniform_prob = tl.load(uniform_probs_ptr + token_idx)
-                uniform_prob = uniform_prob * tmp_uniform_prob
+                    if NO_DRAFT_PROBS:
+                        draft_prob = 1.0
+                    else:
+                        draft_prob = tl.load(draft_probs_ptr + token_idx * vocab_size + draft_token_id)
 
-                if NO_DRAFT_PROBS:
-                    draft_prob = 1.0
-                else:
-                    vocab_for_draft = global_vocab_size if ENABLE_REDUCE_SAMPLING else vocab_size
-                    draft_prob = tl.load(draft_probs_ptr + token_idx * vocab_for_draft + draft_token_id)
+                    if draft_prob > 0:
+                        prefix_prob = min(prefix_prob * target_prob / draft_prob, 1.0)
+                    else:
+                        prefix_prob = 0.0
 
-                pi = min(pi * target_prob / draft_prob, 1.0)
-                if draft_prob > 0 and pi >= uniform_prob:
-                    last_accepted_token_pos = pos
+                    if pos == num_draft_tokens - 1:
+                        h_block = prefix_prob
+                    else:
+                        next_token_idx = token_idx + 1
+                        if NO_DRAFT_PROBS:
+                            next_draft_token_id = tl.load(draft_token_ids_ptr + next_token_idx)
+                            next_target_prob = tl.load(
+                                target_probs_ptr + next_token_idx * vocab_size + next_draft_token_id
+                            )
+                            residual_mass = prefix_prob * (1.0 - next_target_prob)
+                        else:
+                            residual_mass = 0.0
+                            for loop_i in range(loop):
+                                vocab_start = loop_i * SUB_BLOCK
+                                vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK)
+                                next_draft_prob = tl.load(
+                                    draft_probs_ptr + next_token_idx * vocab_size + vocab_offset,
+                                    mask=vocab_offset < vocab_size,
+                                    other=0,
+                                )
+                                next_target_prob = tl.load(
+                                    target_probs_ptr + next_token_idx * vocab_size + vocab_offset,
+                                    mask=vocab_offset < vocab_size,
+                                    other=0,
+                                )
+                                residual_prob = tl.maximum(prefix_prob * next_target_prob - next_draft_prob, 0.0)
+                                residual_mass += tl.sum(residual_prob, axis=0)
+                        denom = residual_mass + 1.0 - prefix_prob
+                        h_block = residual_mass / denom if denom > 0 else 0.0
 
-            # Store accepted tokens
-            if last_accepted_token_pos > -1:
-                for pos in range(last_accepted_token_pos + 1):
+                    uniform_prob = tl.load(uniform_probs_ptr + token_idx)
+                    if uniform_prob <= h_block:
+                        accepted_len = pos + 1
+
+                for pos in range(accepted_len):
                     token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
                     tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id)
 
-            # Store recovered or bonus token
-            if last_accepted_token_pos + 1 < num_draft_tokens:
-                # Rejected - store recovered token
-                recovered_token_id = tl.load(recovered_token_ids_ptr + start_idx + last_accepted_token_pos + 1)
-                tl.store(
-                    output_token_ids_ptr + req_idx * (max_spec_len + 1) + last_accepted_token_pos + 1,
-                    recovered_token_id,
-                )
-            else:
-                # All accepted - store bonus token
-                bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
-                tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, bonus_token_id)
+                if accepted_len == num_draft_tokens:
+                    bonus_token_id = tl.load(bonus_token_ids_ptr + req_idx)
+                    tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + num_draft_tokens, bonus_token_id)
+                else:
+                    recovered_token_id = tl.load(recovered_token_ids_ptr + start_idx + accepted_len)
+                    tl.store(
+                        output_token_ids_ptr + req_idx * (max_spec_len + 1) + accepted_len,
+                        recovered_token_id,
+                    )

vllm_ascend/patch/worker/__init__.py

@@ -45,9 +45,6 @@
 else:
     import vllm_ascend.patch.worker.patch_idex_310  # noqa
 import vllm_ascend.patch.worker.patch_rejection_sampler  # noqa
-
-# TODO(MengqingCao): remove after the upstream community is modified
-import vllm_ascend.patch.worker.patch_llama_eagle3  # noqa
 import vllm_ascend.patch.worker.patch_npugraph_ex_triton  # noqa
 import vllm_ascend.patch.worker.patch_kimi_k25  # noqa
 import vllm_ascend.patch.worker.patch_draft_quarot  # noqa