[Spec Decode][DSpark] Markov sampling in the Gumbel kernel, wider DFlash padding, top-p fix

tests/v1/sample/test_topk_topp_sampler.py

@@ -411,6 +411,27 @@ def test_extreme_p_values(self):
         p = torch.tensor([0.1, 0.5, 0.9, 1.0] * 4, dtype=torch.float32)
         self._compare_results(logits.clone(), k=None, p=p)
 
+    @pytest.mark.parametrize("k_val", [None, 50])
+    def test_tiny_p_keeps_max_token(self, k_val: int | None):
+        """Above 64 rows, a top-p below the fp32 resolution of the kept mass
+        must still keep a max token instead of masking the whole row."""
+        from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton
+
+        batch_size = 156
+        logits = (
+            3.0
+            * torch.randn(
+                batch_size, 32000, generator=self.generator, dtype=torch.float32
+            )
+        ).bfloat16()
+        logits[1::4, :2] = logits[1::4].amax(dim=-1, keepdim=True)  # tied max
+        logits = logits.float()
+        is_max = logits == logits.amax(dim=-1, keepdim=True)
+        k = None if k_val is None else torch.full((batch_size,), k_val)
+        p = torch.full((batch_size,), 1e-9, dtype=torch.float32)
+        kept = apply_top_k_top_p_triton(logits.clone(), k, p) != float("-inf")
+        assert (kept & is_max).any(dim=-1).all()
+
     def test_large_batch(self):
         """Test with a large batch size."""
         batch_size, vocab_size = 512, 32000

tests/v1/spec_decode/test_dspark_markov_sampling.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DSpark's Markov loop must give the draft ids and cached draft logits of
+adding the Markov bias before sampling, whether or not the Gumbel kernel adds
+the bias itself. The LM head, Markov embedding and Markov bias run the real
+layers."""
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    UnquantizedEmbeddingMethod,
+)
+from vllm.model_executor.models.qwen3_dspark import (
+    DSparkMarkovHead,
+    Qwen3DSparkForCausalLM,
+)
+from vllm.v1.worker.gpu.spec_decode.draft_logits import DraftLogitsCache
+from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator
+
+DEVICE = "cuda"
+HIDDEN = 128
+MARKOV_RANK = 256
+NUM_STEPS = 3
+MAX_NUM_REQS = 64
+DRAFT_VOCAB = 8000
+DTYPE = torch.bfloat16
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs GPU")
+
+
+class _LMHead(nn.Module):
+    def __init__(self, weight: torch.Tensor):
+        super().__init__()
+        self.weight = nn.Parameter(weight, requires_grad=False)
+        self.quant_method = UnquantizedEmbeddingMethod()
+
+
+class _DraftModel(nn.Module):
+    markov_embed = Qwen3DSparkForCausalLM.markov_embed
+    markov_bias = Qwen3DSparkForCausalLM.markov_bias
+    map_draft_to_target = Qwen3DSparkForCausalLM.map_draft_to_target
+
+    def __init__(self, lm_head, markov_head, d2t):
+        super().__init__()
+        self.lm_head = lm_head
+        self.model = SimpleNamespace(markov_head=markov_head)
+        self.logits_processor = LogitsProcessor(DRAFT_VOCAB)
+        self.draft_id_to_target_id = d2t
+
+    def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.logits_processor(self.lm_head, hidden_states, skip_gather=True)
+
+
+def _make_weights(reduced_vocab: bool):
+    g = torch.Generator(device=DEVICE).manual_seed(0)
+    lm_head = _LMHead(
+        torch.randn(DRAFT_VOCAB, HIDDEN, generator=g, device=DEVICE, dtype=DTYPE)
+    )
+    target_vocab = DRAFT_VOCAB + 500 if reduced_vocab else DRAFT_VOCAB
+    markov_head = DSparkMarkovHead(
+        target_vocab, DRAFT_VOCAB, MARKOV_RANK, prefix="markov_head"
+    ).to(device=DEVICE, dtype=DTYPE)
+    with torch.no_grad():
+        markov_head.markov_w1.weight.normal_(generator=g)
+        markov_head.markov_w2.weight.normal_(generator=g)
+    d2t = None
+    if reduced_vocab:
+        # Draft id i maps to target id i + d2t[i] (sorted, unique targets).
+        target_ids = torch.randperm(target_vocab, generator=g, device=DEVICE)
+        target_ids = target_ids[:DRAFT_VOCAB].sort().values
+        d2t = target_ids - torch.arange(DRAFT_VOCAB, device=DEVICE)
+    return lm_head, markov_head, d2t, target_vocab
+
+
+def _make_speculator(
+    lm_head, markov_head, d2t, target_vocab, probabilistic, inputs
+) -> DSparkSpeculator:
+    spec = DSparkSpeculator.__new__(DSparkSpeculator)
+    spec.model = _DraftModel(lm_head, markov_head, d2t)
+    spec.num_speculative_steps = NUM_STEPS
+    spec.max_num_reqs = MAX_NUM_REQS
+    spec.device = torch.device(DEVICE)
+    spec.dtype = DTYPE
+    spec.use_fp64_gumbel = False
+    spec.draft_watermarker = None
+    spec.acceptance_estimator = None
+    spec.use_confidence_head = False
+    spec._draft_topk = None
+    spec.draft_logits = (
+        DraftLogitsCache(
+            MAX_NUM_REQS,
+            NUM_STEPS,
+            target_vocab,
+            dtype=DTYPE,
+            fill=0.0,
+            device=DEVICE,
+            sharded=False,
+        )
+        if probabilistic
+        else None
+    )
+    spec._d2t_scatter_index = None
+    spec._draft_scatter_buf = None
+    if probabilistic and d2t is not None:
+        spec._d2t_scatter_index = torch.arange(DRAFT_VOCAB, device=DEVICE) + d2t
+        spec._draft_scatter_buf = torch.full(
+            (MAX_NUM_REQS, target_vocab), float("-inf"), dtype=DTYPE, device=DEVICE
+        )
+    spec._step_cols = torch.arange(NUM_STEPS, dtype=torch.int32, device=DEVICE)
+    spec._anchor_idx = (
+        torch.arange(MAX_NUM_REQS, dtype=torch.int64, device=DEVICE) * NUM_STEPS
+    )
+    spec.draft_tokens = torch.zeros(
+        MAX_NUM_REQS, NUM_STEPS, dtype=torch.int64, device=DEVICE
+    )
+    spec.temperature = inputs["temperature"].clone()
+    spec.seeds = inputs["seeds"].clone()
+    spec.sample_indices = inputs["sample_indices"].clone()
+    spec.sample_pos = inputs["sample_pos"].clone()
+    spec.sample_idx_mapping = inputs["sample_idx_mapping"].clone()
+    spec.input_buffers = SimpleNamespace(input_ids=inputs["input_ids"].clone())
+    return spec
+
+
+def _make_inputs(rng: np.random.Generator, slots: np.ndarray, target_vocab: int):
+    """Buffers as prepare_dflash_inputs leaves them for `slots` (the batch)."""
+    num_reqs = len(slots)
+    max_rows = MAX_NUM_REQS * NUM_STEPS
+    sample_idx_mapping = np.full(max_rows, -1, dtype=np.int32)
+    sample_idx_mapping[: num_reqs * NUM_STEPS] = np.repeat(slots, NUM_STEPS)
+    sample_indices = np.zeros(max_rows, dtype=np.int64)
+    sample_indices[: num_reqs * NUM_STEPS] = np.arange(num_reqs * NUM_STEPS)
+    sample_pos = np.zeros(max_rows, dtype=np.int64)
+    last_pos = rng.integers(10, 5000, size=num_reqs)
+    sample_pos[: num_reqs * NUM_STEPS] = (
+        last_pos[:, None] + 2 + np.arange(NUM_STEPS)
+    ).ravel()
+    temperature = np.zeros(MAX_NUM_REQS, dtype=np.float32)
+    # A mix of greedy and seeded random requests.
+    temperature[slots] = rng.choice([0.0, 0.6, 1.0, 1.3], size=num_reqs)
+    seeds = rng.integers(0, 2**62, size=MAX_NUM_REQS, dtype=np.int64)
+    input_ids = rng.integers(0, target_vocab, size=max_rows, dtype=np.int32)
+
+    def dev(a):
+        return torch.from_numpy(np.ascontiguousarray(a)).to(DEVICE)
+
+    return {
+        "sample_idx_mapping": dev(sample_idx_mapping),
+        "sample_indices": dev(sample_indices),
+        "sample_pos": dev(sample_pos),
+        "temperature": dev(temperature),
+        "seeds": dev(seeds),
+        "input_ids": dev(input_ids),
+    }
+
+
+def _sample_sequential_unfused(spec, num_reqs: int, head_hidden: torch.Tensor):
+    """The Markov loop with the bias added outside the Gumbel kernel."""
+    num_sample = num_reqs * NUM_STEPS
+    sample_hidden = head_hidden[spec.sample_indices[:num_sample]]
+    base_logits = spec.model.compute_draft_logits(sample_hidden)
+    base_logits = base_logits.view(num_reqs, NUM_STEPS, -1)
+    idx_map = spec.sample_idx_mapping[:num_sample].view(num_reqs, NUM_STEPS)
+    sample_pos = spec.sample_pos[:num_sample].view(num_reqs, NUM_STEPS)
+    prev = spec.input_buffers.input_ids[spec._anchor_idx[:num_reqs]]
+    for i in range(NUM_STEPS):
+        bias = spec.model.markov_bias(spec.model.markov_embed(prev))
+        prev = spec._sample_logits(
+            base_logits[:, i] + bias, idx_map[:, i], sample_pos[:, i], i
+        )
+        spec.draft_tokens[:num_reqs, i] = prev
+
+
+# (slots, padded batch size); padded requests have no slot.
+_BATCHES = [
+    ([5, 9, 13], 4),
+    ([0, 7, 2, 33, 41, 12, 63, 18, 22, 3, 50, 44, 29], 16),
+    (None, 40),  # 37 random slots
+]
+
+
+@pytest.mark.parametrize("probabilistic", [True, False])
+@pytest.mark.parametrize("reduced_vocab", [False, True])
+@pytest.mark.parametrize("batch", range(len(_BATCHES)))
+def test_markov_sampling_matches_unfused(
+    default_vllm_config, probabilistic: bool, reduced_vocab: bool, batch: int
+):
+    """In every mode, the Markov loop gives the draft ids and cached draft
+    logits of adding the Markov bias first. Probabilistic full-vocab drafting
+    adds it inside the Gumbel kernel."""
+    rng = np.random.default_rng(batch)
+    slots, num_reqs_padded = _BATCHES[batch]
+    if slots is None:
+        slots = rng.choice(MAX_NUM_REQS, size=37, replace=False)
+    lm_head, markov_head, d2t, target_vocab = _make_weights(reduced_vocab)
+    inputs = _make_inputs(rng, np.asarray(slots, dtype=np.int64), target_vocab)
+    head_hidden = torch.randn(
+        num_reqs_padded * NUM_STEPS, HIDDEN, device=DEVICE, dtype=DTYPE
+    )
+    specs = [
+        _make_speculator(lm_head, markov_head, d2t, target_vocab, probabilistic, inputs)
+        for _ in range(2)
+    ]
+    specs[0]._sample_sequential(num_reqs_padded, head_hidden)
+    _sample_sequential_unfused(specs[1], num_reqs_padded, head_hidden)
+    assert torch.equal(specs[0].draft_tokens, specs[1].draft_tokens)
+    if probabilistic:
+        assert specs[0].draft_logits is not None and specs[1].draft_logits is not None
+        draft_logits = specs[0].draft_logits.values
+        assert draft_logits.any()
+        assert torch.equal(draft_logits, specs[1].draft_logits.values)

tests/v1/spec_decode/test_rejection_sampler_utils.py

@@ -1291,6 +1291,78 @@ def test_block_verification_accepts_at_least_as_many(num_speculative_steps: int)
     )
 
 
+def test_block_verification_is_unbiased_across_steps():
+    """Consecutive block verification steps must keep the target distribution.
+
+    Block verification observes the draft and the acceptance draw of every row,
+    and after a partial accept the next step verifies the same positions again.
+    If those draws are keyed by position alone, the next step reuses draws that
+    verification has already observed, which biases the output. The single-step
+    tests above give every trial fresh positions and cannot see this, so here
+    each request decodes many steps from its own advancing position, drafting
+    with gumbel_sample the way the speculators do.
+    """
+    torch.manual_seed(0)
+    device = "cuda"
+    num_reqs, num_steps, k = 2048, 32, 5
+    vocab_size = NARROW_VOCAB_SIZE
+
+    target_logits_1d = torch.randn(vocab_size, device=device)
+    # A draft close to the target accepts most proposals, so steps often revisit
+    # positions the previous step verified past its accepted prefix.
+    draft_logits_1d = target_logits_1d + torch.randn(vocab_size, device=device)
+
+    reqs = torch.arange(num_reqs, dtype=torch.int32, device=device)
+    rows = torch.arange(k + 1, device=device)
+    draft_cols = rows[:k].to(torch.int32).repeat(num_reqs)
+    temperature = torch.ones(num_reqs, dtype=torch.float32, device=device)
+    seed = torch.arange(num_reqs, dtype=torch.int64, device=device)
+    draft_logits = torch.empty(num_reqs, k, vocab_size, device=device)
+    target_logits = target_logits_1d.expand(num_reqs * (k + 1), -1).contiguous()
+    cu_num_logits = torch.arange(num_reqs + 1, dtype=torch.int32, device=device) * (
+        k + 1
+    )
+
+    start = torch.zeros(num_reqs, dtype=torch.int64, device=device)
+    emitted = []
+    for _ in range(num_steps):
+        pos = start.unsqueeze(1) + rows
+        # A draft is keyed by the position before it, i.e. its verifying row.
+        drafts = gumbel_sample(
+            draft_logits_1d.expand(num_reqs * k, -1),
+            reqs.repeat_interleave(k),
+            temperature,
+            seed,
+            pos[:, :k].reshape(-1),
+            apply_temperature=True,
+            is_drafting=True,
+            logits_cache=draft_logits,
+            logits_cache_col=draft_cols,
+        )
+        draft_sampled = torch.zeros(num_reqs, k + 1, dtype=torch.int64, device=device)
+        draft_sampled[:, 1:] = drafts.view(num_reqs, k)
+        sampled, num_sampled = _rejection_sample(
+            target_logits=target_logits,
+            draft_logits=draft_logits,
+            draft_sampled=draft_sampled.view(-1),
+            cu_num_logits=cu_num_logits,
+            pos=pos.reshape(-1),
+            idx_mapping=reqs,
+            expanded_idx_mapping=reqs.repeat_interleave(k + 1),
+            expanded_local_pos=rows.to(torch.int32).repeat(num_reqs),
+            temperature=temperature,
+            seed=seed,
+            num_speculative_steps=k,
+            use_block_verification=True,
+        )
+        emitted.append(sampled[rows < num_sampled.unsqueeze(1)])
+        start += num_sampled
+
+    _assert_distribution_match(
+        torch.cat(emitted), torch.softmax(target_logits_1d, dim=0), device
+    )
+
+
 @pytest.mark.parametrize("has_draft_logits", [True, False])
 @pytest.mark.parametrize("fly_window_size", [0, 2])
 def test_chunked_requests_match_full_batch(

tests/v1/worker/test_gpu_gumbel_sample.py

@@ -528,3 +528,33 @@ def test_logits_cache_narrower_than_logits_is_rejected():
             logits_cache=cache,
             logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE),
         )
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_gumbel_combine_preserves_first_nan_and_maximum(dtype):
+    """Combining block partials preserves argmax/gather's tie and NaN ordering."""
+    from vllm.v1.worker.gpu.sample.gumbel import _gumbel_combine_kernel
+
+    # Padded rows exercise strides; 197 blocks exercise the masked reduction tail.
+    maxima = torch.randn(6, 211, dtype=dtype, device=DEVICE)[:, :197]
+    maxima[0].fill_(-torch.inf)
+    maxima[1].fill_(1)
+    maxima[2, 7] = maxima[2, 9] = torch.nan
+    maxima[3, 5] = maxima[3, 6] = torch.inf
+    maxima[4].fill_(-torch.inf)
+    maxima[4, 196] = 0
+    token_ids = torch.arange(197, dtype=torch.int64, device=DEVICE).expand(6, -1)
+    sampled = torch.empty(6, dtype=torch.int64, device=DEVICE)
+
+    _gumbel_combine_kernel[(6,)](
+        token_ids,
+        token_ids.stride(0),
+        maxima,
+        maxima.stride(0),
+        sampled,
+        197,
+        PADDED_NUM_BLOCKS=256,
+    )
+
+    expected = token_ids.gather(1, maxima.argmax(-1, keepdim=True)).view(-1)
+    assert torch.equal(sampled, expected)