#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

from vllm.triton_utils import tl, triton

from vllm_ascend.ops.triton.triton_utils import get_element, get_vectorcore_num


def cal_grid_and_block_size(batch_size: int):
    vectorcore_num = get_vectorcore_num()
    if batch_size <= vectorcore_num:
        grid = batch_size
        block_size = 1
    else:
        grid = vectorcore_num
        block_size = triton.next_power_of_2(triton.cdiv(batch_size, grid))
    return grid, block_size


@triton.jit(do_not_specialize=["max_spec_len"])
def rejection_greedy_sample_spec_len_1_triton(
    output_token_ids_ptr,  # [batch_size, 2]
    draft_token_ids_ptr,  # [num_tokens]
    target_argmax_ptr,  # [num_tokens]
    bonus_token_ids_ptr,
    vec_len,
    BLOCK_SIZE: tl.constexpr,
):
    block_idx = tl.program_id(0)
    offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offset < vec_len

    draft_token_id = tl.load(draft_token_ids_ptr + offset, mask)
    target_argmax_id = tl.load(target_argmax_ptr + offset, mask)
    bonus_token_id = tl.load(bonus_token_ids_ptr + offset, mask)

    tl.store(output_token_ids_ptr + offset * 2, target_argmax_id, mask)

    accept_mask = (draft_token_id == target_argmax_id) & mask
    tl.store(output_token_ids_ptr + offset * 2 + 1, bonus_token_id, accept_mask)


@triton.jit(do_not_specialize=["max_spec_len"])
def bonus_renew(
    bonus_token_ids_ptr,
    position,
    output_token_ids_ptr,
    max_spec_len,
    num_tokens1,
):
    bonus_token_id = tl.load(bonus_token_ids_ptr + position)
    tl.store(output_token_ids_ptr + position * (max_spec_len + 1) + num_tokens1, bonus_token_id)


@triton.jit(do_not_specialize=["vec_len", "max_spec_len"])
def rejection_greedy_sample_triton(
    output_token_ids_ptr,  # [batch_size, max_spec_len + 1]
    cu_num_draft_tokens_ptr,  # [batch_size]
    draft_token_ids_ptr,  # [num_tokens]
    target_argmax_ptr,  # [num_tokens]
    bonus_token_ids_ptr,  # [batch_size]
    is_greedy_ptr,  # [batch_size] or None
    vec_len,
    max_spec_len,
    BLOCK_SIZE: tl.constexpr,
):
    block_idx = tl.program_id(0)
    offset = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offset < vec_len

    if is_greedy_ptr is None:
        is_greedy_mask = mask
    else:
        is_greedy = tl.load(is_greedy_ptr + offset, mask=mask, other=0)
        is_greedy_mask = mask & (is_greedy != 0)

    start_idx = tl.where(offset == 0, 0, tl.load(cu_num_draft_tokens_ptr + offset - 1, is_greedy_mask))
    end_idx = tl.load(cu_num_draft_tokens_ptr + offset, is_greedy_mask)
    num_draft_tokens = end_idx - start_idx

    for pos in tl.range(0, BLOCK_SIZE):
        num_tokens1 = get_element(num_draft_tokens, (pos,))
        rejected = False
        start_idx1 = get_element(start_idx, (pos,))
        is_greedy_mask1 = get_element(is_greedy_mask, (pos,))
        position = block_idx * BLOCK_SIZE + pos
        for i in range(num_tokens1):
            if not rejected:
                draft_token_id = tl.load(draft_token_ids_ptr + start_idx1 + i)
                target_argmax_id = tl.load(target_argmax_ptr + start_idx1 + i)
                tl.store(
                    output_token_ids_ptr + position * (max_spec_len + 1) + i,
                    target_argmax_id,
                )
                if draft_token_id != target_argmax_id:
                    # Reject.
                    rejected = True

        if not rejected and is_greedy_mask1:
            bonus_renew(
                bonus_token_ids_ptr,
                position,
                output_token_ids_ptr,
                max_spec_len,
                num_tokens1,
            )


@triton.jit(do_not_specialize=["max_spec_len"])
def rejection_random_sample_kernel(
    output_token_ids_ptr,  # [batch_size, max_spec_len + 1]
    cu_num_draft_tokens_ptr,  # [batch_size]
    draft_token_ids_ptr,  # [num_tokens]
    draft_probs_ptr,  # [num_tokens, vocab_size] or None
    target_probs_ptr,  # [num_tokens, vocab_size] or [num_tokens, selected_vocab_size] if ENABLE_REDUCE_SAMPLING
    target_indices_ptr,  # [num_tokens, selected_vocab_size] global vocab indices, only used if ENABLE_REDUCE_SAMPLING
    bonus_token_ids_ptr,  # [batch_size]
    recovered_token_ids_ptr,  # [num_tokens]
    uniform_probs_ptr,  # [num_tokens]
    is_greedy_ptr,  # [batch_size]
    max_spec_len,
    vocab_size,  # vocab_size or selected_vocab_size if ENABLE_REDUCE_SAMPLING
    global_vocab_size,  # global vocab size for draft_probs indexing (only used if ENABLE_REDUCE_SAMPLING)
    vec_len,
    ori_target_probs_ptr,  # [num_tokens, ori_vocab_size] original probs for entropy
    NO_ORI_TARGET_PROBS: tl.constexpr,
    NO_DRAFT_PROBS: tl.constexpr,
    ENABLE_REDUCE_SAMPLING: tl.constexpr,  # Whether using reduce sampling
    ENTROPY_VERIFY: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
    VOCAB_BLOCK_SIZE: tl.constexpr = 512,
    POSTERIOR_THRESHOLD: tl.constexpr = 0.95,
    POSTERIOR_ALPHA: tl.constexpr = 0.4,
    SUB_BLOCK: tl.constexpr = 4096,
    EPSILON: tl.constexpr = 1e-10,
):
    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))
    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:
            rejected = False
            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):
                if not rejected:
                    if ENABLE_REDUCE_SAMPLING:
                        token_idx = start_idx + pos
                        draft_token_id = tl.load(draft_token_ids_ptr + token_idx)

                        if draft_token_id == -1:
                            rejected = True
                            token_id = tl.load(recovered_token_ids_ptr + token_idx)
                        else:
                            target_prob = 0.0
                            found = False

                            for v_offset in range(0, vocab_size, VOCAB_BLOCK_SIZE):
                                if not found:
                                    vocab_offsets = v_offset + tl.arange(0, VOCAB_BLOCK_SIZE)
                                    vocab_mask = vocab_offsets < vocab_size

                                    candidate_indices = tl.load(
                                        target_indices_ptr + token_idx * vocab_size + vocab_offsets,
                                        mask=vocab_mask,
                                        other=-1,
                                    )

                                    match_mask = candidate_indices == draft_token_id

                                    candidate_probs = tl.load(
                                        target_probs_ptr + token_idx * vocab_size + vocab_offsets,
                                        mask=vocab_mask,
                                        other=0.0,
                                    )

                                    current_match_prob = tl.sum(candidate_probs * match_mask, axis=0)
                                    if current_match_prob > 0.0:
                                        target_prob = current_match_prob
                                        found = True

                            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)

                            # 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)
                    else:
                        token_idx = start_idx + pos
                        draft_token_id = tl.load(draft_token_ids_ptr + token_idx)
                        if draft_token_id == -1:
                            rejected = True
                            token_id = tl.load(recovered_token_ids_ptr + token_idx)
                        else:
                            target_prob = tl.load(target_probs_ptr + token_idx * global_vocab_size + 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)

                            if ENTROPY_VERIFY:
                                loop = (vocab_size + SUB_BLOCK - 1) // SUB_BLOCK
                                entropy = 0.0
                                for loop_i in range(loop):
                                    vocab_start = loop_i * SUB_BLOCK
                                    vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK)
                                    vocab_mask = vocab_offset < vocab_size
                                    if NO_ORI_TARGET_PROBS:
                                        probs = tl.load(
                                            target_probs_ptr + token_idx * vocab_size + vocab_offset,
                                            vocab_mask,
                                            other=0,
                                        )
                                    else:
                                        probs = tl.load(
                                            ori_target_probs_ptr + token_idx * vocab_size + vocab_offset,
                                            vocab_mask,
                                            other=0,
                                        )
                                    log_probs = tl.log(probs + EPSILON)
                                    entropy_contrib = -probs * log_probs
                                    entropy += tl.sum(entropy_contrib)

                                exp_neg_entropy = tl.exp(-entropy * POSTERIOR_ALPHA)
                                threshold_by_entropy = exp_neg_entropy
                                threshold = tl.minimum(threshold_by_entropy, POSTERIOR_THRESHOLD)
                                _uniform_prob = threshold * uniform_prob
                            else:
                                _uniform_prob = uniform_prob
                            # 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 + token_idx)
                        tl.store(output_token_ids_ptr + req_idx * (max_spec_len + 1) + pos, token_id)

            if not rejected:
                # 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,
                    bonus_token_id,
                )


@triton.jit(do_not_specialize=["replace_from", "replace_to", "vec_len"])
def expand_kernel(
    output_ptr,  # [num_tokens]
    input_ptr,  # [batch_size]
    cu_num_tokens_ptr,  # [batch_size]
    replace_from,
    replace_to,
    vec_len,
    MAX_NUM_TOKENS: tl.constexpr,
    BLOCK_SIZE: tl.constexpr,
):
    req_idx = tl.program_id(0)
    offset = req_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    len_mask = offset < vec_len

    start_idx = tl.where(offset == 0, 0, tl.load(cu_num_tokens_ptr + offset - 1, len_mask))
    end_idx = tl.load(cu_num_tokens_ptr + offset, len_mask)
    num_tokens = end_idx - start_idx

    src_val = tl.load(input_ptr + offset, len_mask)
    src_val = tl.where(src_val == replace_from, replace_to, src_val)

    for i in tl.range(0, BLOCK_SIZE):
        num_tokens1 = get_element(num_tokens, (i,))
        start_idx1 = get_element(start_idx, (i,))
        src_val1 = get_element(src_val, (i,))
        offset1 = tl.arange(0, MAX_NUM_TOKENS)
        tl.store(output_ptr + start_idx1 + offset1, src_val1, mask=offset1 < num_tokens1)


@triton.jit
def sample_recovered_tokens_kernel(
    output_token_ids_ptr,
    cu_num_draft_tokens_ptr,
    draft_token_ids_ptr,
    draft_probs_ptr,
    target_probs_ptr,
    target_indices_ptr,
    q_ptr,
    vocab_size,
    global_vocab_size,
    NO_DRAFT_PROBS: tl.constexpr,
    ENABLE_REDUCE_SAMPLING: tl.constexpr,
    SUB_BLOCK: tl.constexpr,
    VOCAB_BLOCK_SIZE: tl.constexpr = 512,
):
    req_idx = tl.program_id(0)
    pos = tl.program_id(1)

    # Compute token index
    start_idx = tl.where(req_idx == 0, 0, tl.load(cu_num_draft_tokens_ptr + req_idx - 1))
    end_idx = tl.load(cu_num_draft_tokens_ptr + req_idx)
    num_draft_tokens = end_idx - start_idx

    if pos >= num_draft_tokens:
        return

    token_idx = start_idx + pos

    if ENABLE_REDUCE_SAMPLING:
        C = vocab_size
        n_loop = tl.cdiv(C, VOCAB_BLOCK_SIZE)

        global_max_p = tl.full((), -float("inf"), tl.float32)
        global_recovered_id = tl.full((), -1, tl.int64)
        draft_token_id = tl.load(draft_token_ids_ptr + token_idx).to(tl.int64)

        for li in range(n_loop):
            c_start = li * VOCAB_BLOCK_SIZE
            offs = c_start + tl.arange(0, VOCAB_BLOCK_SIZE)
            mask = offs < C

            # Load target prob and global index
            tprob = tl.load(target_probs_ptr + token_idx * C + offs, mask=mask, other=0.0).to(tl.float32)

            gidx = tl.load(target_indices_ptr + token_idx * C + offs, mask=mask, other=0).to(tl.int64)

            if NO_DRAFT_PROBS:
                is_draft = (gidx == draft_token_id) & mask
                prob = tl.where(is_draft, 0.0, tprob)
            else:
                valid = (gidx >= 0) & (gidx < global_vocab_size) & mask
                dprob = tl.load(draft_probs_ptr + token_idx * global_vocab_size + gidx, mask=valid, other=0.0).to(
                    tl.float32
                )
                prob = tl.maximum(tprob - dprob, 0.0)

            qv = tl.load(q_ptr + req_idx * C + offs, mask=mask, other=1.0).to(tl.float32)

            bad_q = (qv <= 0) | (qv != qv) | (qv == float("inf")) | (qv == -float("inf"))
            score = tl.where(bad_q, float("-inf"), prob / qv)
            score = tl.where(mask, score, float("-inf"))

            block_best_score = tl.max(score, axis=0)
            block_best_idx = tl.argmax(score, axis=0).to(tl.int64)
            block_best_global_id = tl.load(target_indices_ptr + token_idx * C + (c_start + block_best_idx)).to(tl.int64)

            better = block_best_score > global_max_p
            global_max_p = tl.where(better, block_best_score, global_max_p)
            global_recovered_id = tl.where(better, block_best_global_id, global_recovered_id)

        tl.store(output_token_ids_ptr + token_idx, global_recovered_id)
    else:
        vocab_size = global_vocab_size
        loop = (vocab_size + SUB_BLOCK - 1) // SUB_BLOCK
        global_recovered_id = -1
        global_max_p = -1.0
        if NO_DRAFT_PROBS:
            draft_token_id = tl.load(draft_token_ids_ptr + start_idx + pos)
            for loop_i in range(loop):
                vocab_start = loop_i * SUB_BLOCK
                vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK)
                prob = tl.load(
                    target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset,
                    mask=vocab_offset < vocab_size,
                    other=0,
                )
                prob = tl.where(vocab_offset == draft_token_id, 0.0, prob)
                q = tl.load(
                    q_ptr + req_idx * vocab_size + vocab_offset, mask=vocab_offset < vocab_size, other=float("-inf")
                )
                new_p = prob / q
                recovered_id = tl.argmax(new_p, axis=-1)
                max_p = get_element(new_p, (recovered_id,))
                if max_p > global_max_p:
                    global_max_p = max_p
                    global_recovered_id = vocab_start + recovered_id
        else:
            for loop_i in range(loop):
                vocab_start = loop_i * SUB_BLOCK
                vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK)
                draft_prob = tl.load(
                    draft_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset,
                    mask=vocab_offset < vocab_size,
                    other=0,
                )
                target_prob = tl.load(
                    target_probs_ptr + (start_idx + pos) * vocab_size + vocab_offset,
                    mask=vocab_offset < vocab_size,
                    other=0,
                )
                prob = tl.maximum(target_prob - draft_prob, 0)
                # NOTE(woosuk): We don't need `prob = prob / tl.sum(prob)` here because
                # `tl.argmax` will select the maximum value.

                q = tl.load(
                    q_ptr + req_idx * vocab_size + vocab_offset, mask=vocab_offset < vocab_size, other=float("-inf")
                )
                new_p = prob / q
                recovered_id = tl.argmax(new_p, axis=-1)
                max_p = get_element(new_p, (recovered_id,))
                if max_p > global_max_p:
                    global_max_p = max_p
                    global_recovered_id = vocab_start + recovered_id

        tl.store(output_token_ids_ptr + start_idx + pos, global_recovered_id)


def rejection_greedy_sample_with_triton(
    output_token_ids,
    num_draft_tokens,
    cu_num_draft_tokens,
    draft_token_ids,
    target_argmax,
    bonus_token_ids,
    is_greedy,
    max_spec_len,
    grid,
    block_size,
):
    vec_len = output_token_ids.shape[0]

    if min(num_draft_tokens) == 1 and max(num_draft_tokens) == 1 and is_greedy is None:
        rejection_greedy_sample_spec_len_1_triton[(grid,)](
            output_token_ids,
            draft_token_ids,
            target_argmax,
            bonus_token_ids,
            vec_len,
            BLOCK_SIZE=block_size,
        )
    else:
        rejection_greedy_sample_triton[(grid,)](
            output_token_ids,
            cu_num_draft_tokens,
            draft_token_ids,
            target_argmax,
            bonus_token_ids,
            is_greedy,
            vec_len,
            max_spec_len,
            BLOCK_SIZE=block_size,
        )


def expand_triton(batch_size, expanded_x, x, cu_num_tokens, replace_from, replace_to, max_num_tokens):
    vec_len = batch_size
    grid, block_size = cal_grid_and_block_size(batch_size)

    expand_kernel[(grid,)](
        expanded_x,
        x,
        cu_num_tokens,
        replace_from,
        replace_to,
        vec_len,
        MAX_NUM_TOKENS=max_num_tokens,  # To avoid recompilation.
        BLOCK_SIZE=block_size,
    )


@triton.jit(do_not_specialize=["max_spec_len"])
def rejection_random_sample_block_verify_kernel(
    output_token_ids_ptr,  # [batch_size, max_spec_len + 1]
    cu_num_draft_tokens_ptr,  # [batch_size]
    draft_token_ids_ptr,  # [num_tokens]
    draft_probs_ptr,  # [num_tokens, vocab_size] or None
    target_probs_ptr,  # [num_tokens, vocab_size] or [num_tokens, selected_vocab_size] if ENABLE_REDUCE_SAMPLING
    target_indices_ptr,  # [num_tokens, selected_vocab_size] global vocab indices, only used if ENABLE_REDUCE_SAMPLING
    bonus_token_ids_ptr,  # [batch_size]
    recovered_token_ids_ptr,  # [num_tokens]
    uniform_probs_ptr,  # [num_tokens]
    is_greedy_ptr,  # [batch_size]
    max_spec_len,
    vocab_size,  # vocab_size or selected_vocab_size if ENABLE_REDUCE_SAMPLING
    global_vocab_size,  # global vocab size for draft_probs indexing (only used if ENABLE_REDUCE_SAMPLING)
    vec_len,
    ori_target_probs_ptr,  # [num_tokens, ori_vocab_size] original probs for entropy
    NO_ORI_TARGET_PROBS: tl.constexpr,
    NO_DRAFT_PROBS: tl.constexpr,
    ENABLE_REDUCE_SAMPLING: tl.constexpr,  # Whether using reduce_sampling
    BLOCK_SIZE: tl.constexpr,
    ENTROPY_VERIFY: tl.constexpr,
    VOCAB_BLOCK_SIZE: tl.constexpr = 512,
    POSTERIOR_THRESHOLD: tl.constexpr = 0.95,
    POSTERIOR_ALPHA: tl.constexpr = 0.4,
    SUB_BLOCK: tl.constexpr = 4096,
    EPSILON: tl.constexpr = 1e-10,
):
    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
    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

    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 draft_token_id == -1:
                        pi = 0.0
                    else:
                        target_prob = 0.0
                        found = False

                        for v_offset in range(0, vocab_size, VOCAB_BLOCK_SIZE):
                            if not found:
                                vocab_offsets = v_offset + tl.arange(0, VOCAB_BLOCK_SIZE)
                                vocab_mask = vocab_offsets < vocab_size

                                candidate_indices = tl.load(
                                    target_indices_ptr + token_idx * vocab_size + vocab_offsets,
                                    mask=vocab_mask,
                                    other=-1,
                                )

                                match_mask = candidate_indices == draft_token_id

                                candidate_probs = tl.load(
                                    target_probs_ptr + token_idx * vocab_size + vocab_offsets,
                                    mask=vocab_mask,
                                    other=0.0,
                                )

                                current_match_prob = tl.sum(candidate_probs * match_mask, axis=0)

                                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:
        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 draft_token_id == -1:
                        pi = 0.0
                    else:
                        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:
                            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 ENTROPY_VERIFY:
                            loop = (vocab_size + SUB_BLOCK - 1) // SUB_BLOCK
                            entropy = 0.0
                            for loop_i in range(loop):
                                vocab_start = loop_i * SUB_BLOCK
                                vocab_offset = vocab_start + tl.arange(0, SUB_BLOCK)
                                vocab_mask = vocab_offset < vocab_size
                                if NO_ORI_TARGET_PROBS:
                                    probs = tl.load(
                                        target_probs_ptr + token_idx * vocab_size + vocab_offset,
                                        vocab_mask,
                                        other=0,
                                    )
                                else:
                                    probs = tl.load(
                                        ori_target_probs_ptr + token_idx * vocab_size + vocab_offset,
                                        vocab_mask,
                                        other=0,
                                    )
                                log_probs = tl.log(probs + EPSILON)
                                entropy_contrib = -probs * log_probs
                                entropy += tl.sum(entropy_contrib)

                            exp_neg_entropy = tl.exp(-entropy * POSTERIOR_ALPHA)
                            threshold_by_entropy = exp_neg_entropy
                            threshold = tl.minimum(threshold_by_entropy, POSTERIOR_THRESHOLD)
                            _uniform_prob = threshold * uniform_prob
                        else:
                            _uniform_prob = uniform_prob

                        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)