Function _create_share_inputs

tests/layers/test_speculative_sampler.py:108–149

(max_num_seqs, max_draft_token_num, max_model_len, vocab_size)

def _create_share_inputs(max_num_seqs, max_draft_token_num, max_model_len, vocab_size):
    share_inputs = {}
    share_inputs["seq_lens_this_time"] = paddle.full([max_num_seqs, 1], 2, dtype="int32")

    cu_seqlens_q_output = [0] + paddle.cumsum(share_inputs["seq_lens_this_time"]).numpy().tolist()
    share_inputs["cu_seqlens_q_output"] = paddle.to_tensor(cu_seqlens_q_output).cast("int32")
    share_inputs["batch_id_per_token_output"] = paddle.arange(max_num_seqs, dtype="int32") * 2

    share_inputs["accept_tokens"] = paddle.full(
        shape=[max_num_seqs, max_draft_token_num + 1], fill_value=0, dtype="int64"
    )
    share_inputs["accept_num"] = paddle.full(shape=[max_num_seqs], fill_value=1, dtype="int32")
    share_inputs["step_idx"] = paddle.full([max_num_seqs, 1], 1, dtype="int64")
    share_inputs["stop_flags"] = paddle.full([max_num_seqs, 1], False, dtype="bool")
    share_inputs["seq_lens_encoder"] = paddle.full([max_num_seqs, 1], 0, dtype="int32")
    share_inputs["seq_lens_decoder"] = paddle.full([max_num_seqs, 1], 2, dtype="int32")
    share_inputs["draft_tokens"] = paddle.full(
        shape=[max_num_seqs, max_draft_token_num + 1], fill_value=0, dtype="int64"
    )
    share_inputs["max_dec_len"] = paddle.full([max_num_seqs, 1], max_model_len, dtype="int64")
    share_inputs["is_block_step"] = paddle.full([max_num_seqs], False, dtype="bool")
    share_inputs["actual_draft_token_num"] = paddle.full(
        shape=[max_num_seqs], fill_value=max_draft_token_num, dtype="int32"
    )

    share_inputs["batch_token_num"] = paddle.where(
        share_inputs["seq_lens_encoder"] != 0,
        paddle.ones_like(share_inputs["seq_lens_encoder"]),
        share_inputs["seq_lens_this_time"],
    ).squeeze(1)
    share_inputs["next_token_num"] = paddle.full(shape=[max_num_seqs], fill_value=0, dtype="int32")
    share_inputs["cu_batch_token_offset"] = paddle.concat(
        [paddle.to_tensor([0]), paddle.cumsum(share_inputs["accept_num"])]
    ).astype("int32")
    share_inputs["cu_next_token_offset"] = paddle.full(shape=[max_num_seqs + 1], fill_value=0, dtype="int32")
    share_inputs["substep"] = 0
    share_inputs["draft_logits"] = paddle.full(
        [max_num_seqs * (max_draft_token_num + 1), vocab_size], -1, dtype="float32"
    )
    share_inputs["reasoning_status"] = paddle.zeros([max_num_seqs], dtype="int32")

    return share_inputs


def _create_padding_inputs():