[Model Runner V2][Spec Decode] Fix draft logits cache column stride in gumbel_sample (#53017)

tests/v1/worker/test_gpu_gumbel_sample.py

@@ -291,3 +291,72 @@ def test_logits_cache_stores_input_logits_bitwise(
     assert torch.equal(_float_bits(stored), _float_bits(logits)), (
         "cached logits differ from the input logits"
     )
+
+
+@pytest.mark.parametrize("extra_cache_cols", [0, 1])
+def test_logits_cache_columns_stay_separate_across_steps(extra_cache_cols: int):
+    """Each drafting step must land in its own cache column.
+
+    `extra_cache_cols=1` is the shape a draft produces when it adds an
+    input-only mask/noise row to its embedding table but keeps a full-width
+    output head: N-wide logits cached into N+1-wide rows. Striding cache
+    columns by the logits width instead of the cache's own row width leaves
+    step 0 correct and silently misaligns every step after it.
+    """
+    torch.manual_seed(0)
+    num_reqs, vocab_size, num_steps = 4, 1031, 3
+    idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
+    temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE)
+    seed = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
+    pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
+
+    cache = torch.zeros(
+        num_reqs, num_steps, vocab_size + extra_cache_cols, device=DEVICE
+    )
+    cols = torch.arange(num_steps, dtype=torch.int32, device=DEVICE)
+    per_step = [
+        torch.randn(num_reqs, vocab_size, device=DEVICE) for _ in range(num_steps)
+    ]
+
+    for step, logits in enumerate(per_step):
+        gumbel_sample(
+            logits,
+            idx_mapping,
+            temp,
+            seed,
+            pos,
+            apply_temperature=True,
+            logits_cache=cache,
+            logits_cache_col=cols[step],
+        )
+
+    for step, logits in enumerate(per_step):
+        stored = cache[:, step, :vocab_size]
+        assert torch.equal(_float_bits(stored), _float_bits(logits)), (
+            f"step {step} was overwritten by a later step"
+        )
+    # Columns past the sampled width belong to no step and stay untouched.
+    assert not cache[:, :, vocab_size:].any()
+
+
+def test_logits_cache_narrower_than_logits_is_rejected():
+    """A cache too narrow to hold a step would silently drop its tail."""
+    num_reqs, vocab_size, num_steps = 2, 64, 3
+    logits = torch.randn(num_reqs, vocab_size, device=DEVICE)
+    idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=DEVICE)
+    temp = torch.ones(num_reqs, dtype=torch.float32, device=DEVICE)
+    seed = torch.zeros(num_reqs, dtype=torch.int64, device=DEVICE)
+    pos = torch.arange(num_reqs, dtype=torch.int64, device=DEVICE)
+    cache = torch.zeros(num_reqs, num_steps, vocab_size - 1, device=DEVICE)
+
+    with pytest.raises(AssertionError, match="narrower"):
+        gumbel_sample(
+            logits,
+            idx_mapping,
+            temp,
+            seed,
+            pos,
+            apply_temperature=True,
+            logits_cache=cache,
+            logits_cache_col=torch.tensor(0, dtype=torch.int32, device=DEVICE),
+        )

vllm/v1/worker/gpu/sample/gumbel.py

@@ -91,8 +91,10 @@ def gumbel_block_argmax(
     temp_ptr,
     seeds_ptr,
     pos_ptr,
+    # [max_num_reqs, num_cols, vocab_size]
     logits_cache_ptr,
-    logits_cache_stride,
+    logits_cache_stride_0,
+    logits_cache_stride_1,
     logits_cache_col_ptr,
     vocab_size,
     APPLY_TEMPERATURE: tl.constexpr,
@@ -115,8 +117,8 @@ def gumbel_block_argmax(
             col = tl.load(logits_cache_col_ptr)
         tl.store(
             logits_cache_ptr
-            + req_state_idx * logits_cache_stride
-            + col * vocab_size
+            + req_state_idx * logits_cache_stride_0
+            + col * logits_cache_stride_1
             + block,
             logits,
             mask=mask & is_valid_req,
@@ -164,8 +166,10 @@ def _gumbel_sample_kernel(
     local_argmax_stride,
     local_max_ptr,
     local_max_stride,
+    # [max_num_reqs, num_cols, vocab_size]
     logits_cache_ptr,
-    logits_cache_stride,
+    logits_cache_stride_0,
+    logits_cache_stride_1,
     logits_cache_col_ptr,
     logits_ptr,
     logits_stride,
@@ -200,7 +204,8 @@ def _gumbel_sample_kernel(
         seeds_ptr,
         pos_ptr,
         logits_cache_ptr,
-        logits_cache_stride,
+        logits_cache_stride_0,
+        logits_cache_stride_1,
         logits_cache_col_ptr,
         vocab_size,
         APPLY_TEMPERATURE=APPLY_TEMPERATURE,
@@ -229,6 +234,12 @@ def gumbel_sample(
     if logits_cache_col is not None:
         logits_cache_col = logits_cache_col.contiguous()
     num_tokens, vocab_size = logits.shape
+    if logits_cache is not None:
+        assert logits_cache.size(-1) >= vocab_size, (
+            f"draft logits cache vocab dim ({logits_cache.size(-1)}) is narrower "
+            f"than the sampled logits ({vocab_size}). Cached logits would be "
+            "truncated."
+        )
     BLOCK_SIZE = 1024
     num_blocks = triton.cdiv(vocab_size, BLOCK_SIZE)
     local_argmax = logits.new_empty(num_tokens, num_blocks, dtype=torch.int64)
@@ -242,6 +253,7 @@ def gumbel_sample(
         local_max.stride(0),
         logits_cache,
         logits_cache.stride(0) if logits_cache is not None else 0,
+        logits_cache.stride(1) if logits_cache is not None else 0,
         logits_cache_col,
         logits,
         logits.stride(0),

vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py

@@ -833,7 +833,8 @@ def _resample_kernel(
         seed_ptr,
         pos_ptr,
         None,  # logits_cache_ptr
-        0,  # logits_cache_stride
+        0,  # logits_cache_stride_0
+        0,  # logits_cache_stride_1
         None,  # logits_cache_col_ptr
         vocab_size,
         APPLY_TEMPERATURE=False,