lmdeploy/tests/pytorch/spec_decode/test_reject_sample.py

import torch

from lmdeploy.pytorch.engine.logits_process import SamplingInputs
from lmdeploy.pytorch.spec_decode.reject_sampler import (
    PLACEHOLDER_TOKEN_ID,
    _extract_outputs,
    rejection_greedy_sample_kernel,
    rejection_sample,
    sample_recovered_tokens_kernel,
    torch_greedy_rejection_sample,
)

device = 'cuda' if torch.cuda.is_available() else 'cpu'


def _make_peaked_logits(token_ids_2d, vocab_size):
    """Build logits where argmax(dim=-1) == token_ids_2d.

    token_ids_2d: list[list[int]] or Tensor [batch, num_spec]
    """
    if isinstance(token_ids_2d, torch.Tensor):
        token_ids_2d = token_ids_2d.tolist()
    batch_size = len(token_ids_2d)
    num_spec = len(token_ids_2d[0])
    logits = torch.full((batch_size, num_spec, vocab_size), -100.0, device=device)
    for b in range(batch_size):
        for s in range(num_spec):
            logits[b, s, token_ids_2d[b][s]] = 100.0
    return logits


def _extract_valid(output_token_ids):
    """Return list-of-lists of non-placeholder token ids."""
    out = []
    for row in output_token_ids:
        out.append(row[row != PLACEHOLDER_TOKEN_ID].tolist())
    return out


class TestRejectSample:
    """Tests for rejection_sample and torch_greedy_rejection_sample.

    Uses hand-crafted token ids so the expected output is obvious.
    """

    # ----- greedy: rejection_sample -----

    def test_greedy_all_match(self):
        """All draft == target -> accept all + bonus."""
        target_logits = _make_peaked_logits([[10, 20, 30], [10, 20, 30]], 64)
        draft = torch.tensor([[10, 20, 30], [10, 20, 30]], dtype=torch.long, device=device)
        bonus = torch.tensor([99, 88], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert out.shape == (2, 4)
        assert (rej == 0).all()
        assert _extract_valid(out) == [[10, 20, 30, 99], [10, 20, 30, 88]]
        assert last.tolist() == [99, 88]

    def test_greedy_first_mismatch(self):
        """Mismatch at position 0 -> keep only target[0]."""
        target_logits = _make_peaked_logits([[10, 20, 30], [10, 20, 30]], 64)
        draft = torch.tensor([[11, 20, 30], [11, 20, 30]], dtype=torch.long, device=device)
        bonus = torch.tensor([99, 88], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert (rej == 3).all()
        assert _extract_valid(out) == [[10], [10]]
        assert last.tolist() == [10, 10]

    def test_greedy_middle_mismatch(self):
        """Mismatch at position 2 -> keep target[0:3], no bonus."""
        target_logits = _make_peaked_logits([[10, 20, 30], [10, 20, 30]], 64)
        draft = torch.tensor([[10, 20, 31], [10, 20, 31]], dtype=torch.long, device=device)
        bonus = torch.tensor([99, 88], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert (rej == 1).all()
        assert _extract_valid(out) == [[10, 20, 30], [10, 20, 30]]
        assert last.tolist() == [30, 30]

    def test_greedy_mixed_mismatch_positions(self):
        """Each row has a different first-mismatch position.

        Row 0: all match       -> [10, 20, 30, 99]
        Row 1: mismatch at 0   -> [10]
        Row 2: mismatch at 1   -> [10, 20]
        Row 3: mismatch at 2   -> [10, 20, 30]
        """
        target_logits = _make_peaked_logits([[10, 20, 30]] * 4, 64)
        draft = torch.tensor([
            [10, 20, 30],
            [11, 20, 30],
            [10, 21, 30],
            [10, 20, 31],
        ],
                             dtype=torch.long,
                             device=device)
        bonus = torch.tensor([99, 88, 77, 66], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert rej.tolist() == [0, 3, 2, 1]
        assert _extract_valid(out) == [
            [10, 20, 30, 99],
            [10],
            [10, 20],
            [10, 20, 30],
        ]
        assert last.tolist() == [99, 10, 20, 30]

    def test_greedy_all_rejected(self):
        """Every draft wrong -> only first target token kept."""
        target_logits = _make_peaked_logits([[10, 20, 30], [10, 20, 30]], 64)
        draft = torch.tensor([[11, 21, 31], [11, 21, 31]], dtype=torch.long, device=device)
        bonus = torch.tensor([99, 88], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert (rej == 3).all()
        assert _extract_valid(out) == [[10], [10]]
        assert last.tolist() == [10, 10]

    # ----- greedy: compare rejection_sample vs torch_greedy_rejection_sample -----

    def test_greedy_triton_vs_torch_all_match(self):
        """Both implementations agree when all drafts match."""
        target_logits = _make_peaked_logits([[10, 20, 30]] * 4, 64)
        draft = torch.tensor([[10, 20, 30]] * 4, dtype=torch.long, device=device)
        bonus = torch.tensor([99, 88, 77, 66], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        t_out, t_rej, t_last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)
        p_out, p_rej, p_last = torch_greedy_rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert torch.equal(t_out, p_out)
        assert torch.equal(t_rej, p_rej)
        assert torch.equal(t_last, p_last)

    def test_greedy_triton_vs_torch_mixed(self):
        """Both implementations agree with mixed mismatch positions."""
        target_logits = _make_peaked_logits([[10, 20, 30]] * 4, 64)
        draft = torch.tensor(
            [
                [10, 20, 30],  # all match
                [11, 20, 30],  # pos 0
                [10, 21, 30],  # pos 1
                [10, 20, 31],  # pos 2
            ],
            dtype=torch.long,
            device=device)
        bonus = torch.tensor([99, 88, 77, 66], dtype=torch.long, device=device)
        si = SamplingInputs(max_top_k=1)

        t_out, t_rej, t_last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)
        p_out, p_rej, p_last = torch_greedy_rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert torch.equal(t_out, p_out), f'output mismatch: triton={t_out}, torch={p_out}'
        assert torch.equal(t_rej, p_rej), f'num_rejected: triton={t_rej}, torch={p_rej}'
        assert torch.equal(t_last, p_last), f'last_token: triton={t_last}, torch={p_last}'

    def test_greedy_triton_vs_torch_large_batch(self):
        """Both implementations agree with large batch + random mismatches."""
        batch_size, num_spec, vocab = 32, 6, 128
        target_logits = torch.randn(batch_size, num_spec, vocab, device=device)
        target_argmax = target_logits.argmax(dim=-1)
        draft = target_argmax.clone()

        torch.manual_seed(42)
        flip = torch.rand(batch_size, num_spec) < 0.3
        for i in range(batch_size):
            for j in range(num_spec):
                if flip[i, j]:
                    draft[i, j] = (target_argmax[i, j] + 1) % vocab

        bonus = torch.randint(0, vocab, (batch_size, ), device=device)
        si = SamplingInputs(max_top_k=1)

        t_out, t_rej, t_last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)
        p_out, p_rej, p_last = torch_greedy_rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        assert torch.equal(t_out, p_out)
        assert torch.equal(t_rej, p_rej)
        assert torch.equal(t_last, p_last)

    # ----- mixed greedy/random -----

    def test_mixed_batch(self):
        """Mixed batch: greedy (top_k=1) and random sequences."""
        target_logits = _make_peaked_logits([[5, 6, 7], [5, 6, 7], [5, 6, 7], [5, 6, 7]], 32)
        draft = target_logits.argmax(dim=-1)
        bonus = torch.tensor([99, 88, 77, 66], dtype=torch.long, device=device)

        top_k = torch.tensor([1, 10, 1, 10], device=device)
        si = SamplingInputs(max_top_k=10, top_k=top_k)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si)

        # Greedy rows (0, 2) must accept all since draft == target
        assert rej[0] == 0
        assert rej[2] == 0
        assert last[0] == 99
        assert last[2] == 77

    # ----- random: no draft_probs (MTP/n-gram) -----

    def test_random_no_draft_probs(self):
        """Random sampling with draft_probs=None (MTP/n-gram mode)."""
        target_logits = _make_peaked_logits([[5, 6, 7], [5, 6, 7]], 32)
        draft = torch.tensor([[5, 6, 7], [5, 6, 7]], dtype=torch.long, device=device)
        bonus = torch.tensor([99, 88], dtype=torch.long, device=device)

        top_k = torch.tensor([10, 10], device=device)
        si = SamplingInputs(max_top_k=10, top_k=top_k)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si, draft_probs=None)

        assert out.shape == (2, 4)
        assert (last >= 0).all()

    # ----- random: all reject (peaked target, wrong draft) -----

    def test_random_all_reject(self):
        """Random: draft picks wrong token, peaked target -> reject at pos 0."""
        vocab = 32
        target_logits = _make_peaked_logits([[5, 6, 7], [5, 6, 7]], vocab)
        draft = torch.tensor([[1, 1, 1], [1, 1, 1]], dtype=torch.long, device=device)
        draft_probs = torch.zeros(2, 3, vocab, device=device)
        draft_probs[:, :, 1] = 1.0
        bonus = torch.tensor([99, 88], dtype=torch.long, device=device)

        top_k = torch.tensor([10, 10], device=device)
        si = SamplingInputs(max_top_k=10, top_k=top_k)

        out, rej, last = rejection_sample(target_logits, draft, bonus, sampling_inputs=si, draft_probs=draft_probs)

        # Reject at pos 0 -> 1 recovered token per row
        valid = _extract_valid(out)
        for i in range(2):
            assert len(valid[i]) == 1, f'row {i}: expected 1 token, got {valid[i]}'

    # ----- _extract_outputs -----

    def test_extract_outputs(self):
        """Test _extract_outputs helper function."""
        output_token_ids = torch.tensor([
            [5, 6, 7, 8],  # all valid
            [10, 11, -1, -1],  # 2 valid
        ])
        out, rej, last = _extract_outputs(output_token_ids, num_spec_tokens=3)
        assert rej.tolist() == [0, 2]
        assert last.tolist() == [8, 11]


# ---------------------------------------------------------------------------
# Triton kernel tests (direct kernel invocations)
# ---------------------------------------------------------------------------


class TestTritonKernels:

    def test_greedy_kernel_mixed(self):
        """rejection_greedy_sample_kernel with mixed match/mismatch rows.

        Row 0: all match       -> [10, 20, 30, 40, 50]
        Row 1: mismatch at 2   -> [10, 20, 30, -1, -1]
        Row 2: mismatch at 0   -> [10, -1, -1, -1, -1]
        """
        batch_size, num_spec = 3, 4
        target_argmax = torch.tensor([[10, 20, 30, 40]] * 3, dtype=torch.long, device=device)
        draft = torch.tensor([
            [10, 20, 30, 40],
            [10, 20, 99, 40],
            [99, 20, 30, 40],
        ], dtype=torch.long, device=device)
        bonus = torch.tensor([50, 51, 52], dtype=torch.long, device=device)
        out = torch.full((batch_size, num_spec + 1), PLACEHOLDER_TOKEN_ID, dtype=torch.long, device=device)

        rejection_greedy_sample_kernel[(batch_size, )](out, draft, target_argmax, bonus, None, num_spec)

        assert out[0].tolist() == [10, 20, 30, 40, 50]
        assert out[1].tolist() == [10, 20, 30, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID]
        assert out[2].tolist() == [
            10, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID, PLACEHOLDER_TOKEN_ID
        ]

    def test_greedy_kernel_with_is_greedy_mask(self):
        """rejection_greedy_sample_kernel skips non-greedy rows.

        Row 0: greedy, all match -> [10, 20, 30, 50]
        Row 1: not greedy        -> all placeholder (untouched)
        """
        batch_size, num_spec = 2, 3
        target_argmax = torch.tensor([[10, 20, 30]] * 2, dtype=torch.long, device=device)
        draft = target_argmax.clone()
        bonus = torch.tensor([50, 51], dtype=torch.long, device=device)
        out = torch.full((batch_size, num_spec + 1), PLACEHOLDER_TOKEN_ID, dtype=torch.long, device=device)
        is_greedy = torch.tensor([True, False], dtype=torch.bool, device=device)

        rejection_greedy_sample_kernel[(batch_size, )](out, draft, target_argmax, bonus, is_greedy, num_spec)

        assert out[0].tolist() == [10, 20, 30, 50]
        assert (out[1] == PLACEHOLDER_TOKEN_ID).all()

    def test_recovered_tokens_with_draft_probs(self):
        """sample_recovered_tokens_kernel: recovered token = argmax of
        max(0, target - draft) * inv_q.

        target = [0, 0, 0, 0.9, 0, 0.1, 0, 0]
        draft  = [0, 0, 0, 0.9, 0, 0,   0, 0]
        diff   = [0, 0, 0, 0,   0, 0.1, 0, 0]  -> token 5
        """
        batch_size, num_spec, vocab = 2, 2, 8
        target_probs = torch.zeros(batch_size, num_spec, vocab, device=device, dtype=torch.float32)
        target_probs[:, :, 3] = 0.9
        target_probs[:, :, 5] = 0.1

        draft_probs = torch.zeros_like(target_probs)
        draft_probs[:, :, 3] = 0.9

        draft_ids = torch.full((batch_size, num_spec), 3, dtype=torch.long, device=device)
        inv_q = torch.ones(batch_size, vocab, device=device, dtype=torch.float32)
        recovered = torch.empty(batch_size, num_spec, dtype=torch.long, device=device)

        sample_recovered_tokens_kernel[(batch_size, num_spec)](recovered,
                                                               draft_ids,
                                                               draft_probs,
                                                               target_probs,
                                                               inv_q,
                                                               num_spec,
                                                               vocab,
                                                               8,
                                                               NO_DRAFT_PROBS=False)

        assert (recovered == 5).all()

    def test_recovered_tokens_no_draft_probs(self):
        """sample_recovered_tokens_kernel with NO_DRAFT_PROBS: draft token
        is masked out, recovered = argmax of remaining target probs.

        target   = [0, 0, 0.6, 0, 0, 0.3, 0, 0.1]
        draft_id = 2 (masked out)
        remaining= [0, 0, 0,   0, 0, 0.3, 0, 0.1]  -> token 5
        """
        batch_size, num_spec, vocab = 2, 2, 8
        target_probs = torch.zeros(batch_size, num_spec, vocab, device=device, dtype=torch.float32)
        target_probs[:, :, 2] = 0.6
        target_probs[:, :, 5] = 0.3
        target_probs[:, :, 7] = 0.1

        draft_ids = torch.full((batch_size, num_spec), 2, dtype=torch.long, device=device)
        inv_q = torch.ones(batch_size, vocab, device=device, dtype=torch.float32)
        recovered = torch.empty(batch_size, num_spec, dtype=torch.long, device=device)

        sample_recovered_tokens_kernel[(batch_size, num_spec)](recovered,
                                                               draft_ids,
                                                               None,
                                                               target_probs,
                                                               inv_q,
                                                               num_spec,
                                                               vocab,
                                                               8,
                                                               NO_DRAFT_PROBS=True)

        assert (recovered == 5).all()