diff --git a/tests/v1/worker/test_gpu_gumbel_sample.py b/tests/v1/worker/test_gpu_gumbel_sample.py index f64a3f1f24..a87c12416e 100644 --- a/tests/v1/worker/test_gpu_gumbel_sample.py +++ b/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), + ) diff --git a/vllm/v1/worker/gpu/sample/gumbel.py b/vllm/v1/worker/gpu/sample/gumbel.py index 0a930b8183..195b2aadf1 100644 --- a/vllm/v1/worker/gpu/sample/gumbel.py +++ b/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), diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index 3ea9f35ee8..abbf0ee96e 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/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,