Function test_top_k_top_p

tests/test_core.py:105–138

Verify that top-k and top-p filtering restricts samples to the expected tokens.

(provider, vocab_size, n_hidden_states)

@pytest.mark.parametrize("vocab_size", [100, 200, 256])
@pytest.mark.parametrize("provider", ["naive-pt", "naive-compiled", "fused-topk"])
def test_top_k_top_p(provider, vocab_size, n_hidden_states):
    """Verify that top-k and top-p filtering restricts samples to the expected tokens."""
    inputs = make_synthetic_inputs(vocab_size=vocab_size, n_hidden_states=n_hidden_states)
    temperature = torch.tensor(1.0, device=device)
    top_k = 10
    top_p = 0.9
    num_samples = 5_000

    sampler = get_sampler(provider, weights=inputs.weights)
    sampler.prepare()
    samples = sampler.sample(
        weights=inputs.weights,
        hidden_states=inputs.hidden_states,
        num_samples=num_samples,
        temperature=temperature,
        top_k=top_k,
        top_p=top_p,
    )

    # Use the same bf16 matmul logits the sampler sees, not the exact float32 logits.
    ref_logits = inputs.hidden_states @ inputs.weights.T
    for seq_idx in range(n_hidden_states):
        allowed_tokens = reference_top_k_top_p(
            logits=ref_logits[seq_idx], temperature=temperature, top_k=top_k, top_p=top_p
        )
        unique_sampled = torch.unique(samples[seq_idx])
        allowed_set = set(allowed_tokens.cpu().tolist())
        sampled_set = set(unique_sampled.cpu().tolist())
        assert sampled_set <= allowed_set, (
            f"seq {seq_idx}: sampled tokens not in allowed set. Extra: {sampled_set - allowed_set}"
        )
        assert sampled_set == allowed_set, (
            f"seq {seq_idx}: not all allowed tokens sampled. Missing: {allowed_set - sampled_set}"
        )


@pytest.mark.parametrize("provider", ["naive-pt", "fused-topk"])