diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index cc89397f68a1..8af7f6846b41 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -36,18 +36,34 @@ __global__ void merge_attn_states_kernel( const uint pack_size = 16 / sizeof(scalar_t); const uint threads_per_head = head_size / pack_size; - const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x; + const uint global_idx = blockIdx.x * blockDim.x + threadIdx.x; const uint token_head_threads = num_tokens * num_heads * threads_per_head; - if (global_idx >= token_head_threads) return; - - // global_idx -> token_idx + head_idx + pack_idx + // Derive indices before the block barrier so every thread reaches it. const uint token_head_idx = global_idx / threads_per_head; const uint pack_idx = global_idx % threads_per_head; - const uint token_idx = token_head_idx / num_heads; const uint head_idx = token_head_idx % num_heads; + // A running chunked-attention LSE may be both prefix_lse and output_lse. + // The launcher aligns block boundaries to complete head groups, allowing + // every group to load its LSE values before any thread overwrites them. + __shared__ float shared_prefix_lse[NUM_THREADS]; + __shared__ float shared_suffix_lse[NUM_THREADS]; + const bool is_valid = global_idx < token_head_threads; + const uint group_idx = threadIdx.x / threads_per_head; + + if (is_valid && pack_idx == 0 && token_idx < prefix_num_tokens) { + shared_prefix_lse[group_idx] = + prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + shared_suffix_lse[group_idx] = + suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; + } + __syncthreads(); + if (!is_valid) return; + const uint pack_offset = pack_idx * pack_size; // (0~15)*8, etc. const uint src_head_offset = token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride; @@ -95,11 +111,9 @@ __global__ void merge_attn_states_kernel( return; } - // For tokens within prefix range, merge prefix and suffix - float p_lse = prefix_lse[head_idx * prefix_lse_head_stride + - token_idx * prefix_lse_token_stride]; - float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + - token_idx * suffix_lse_token_stride]; + // For tokens within prefix range, merge prefix and suffix. + float p_lse = shared_prefix_lse[group_idx]; + float s_lse = shared_suffix_lse[group_idx]; p_lse = std::isinf(p_lse) ? -std::numeric_limits::infinity() : p_lse; s_lse = std::isinf(s_lse) ? -std::numeric_limits::infinity() : s_lse; @@ -307,10 +321,19 @@ void merge_attn_states_launcher( // Process one pack elements per thread. for float, the // pack_size is 4 for half/bf16, the pack_size is 8. const uint threads_per_head = head_size / pack_size; + STD_TORCH_CHECK( + threads_per_head <= NUM_THREADS, + "headsize requires more threads than the merge kernel block supports: ", + head_size); const uint total_threads = num_tokens * num_heads * threads_per_head; + // Keep each token-head group inside one block. This is required when + // output_lse aliases prefix_lse because the whole group must read the input + // LSE before its first thread writes the merged value. + const uint block_threads = + (NUM_THREADS / threads_per_head) * threads_per_head; - dim3 block(NUM_THREADS); - dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS); + dim3 block(block_threads); + dim3 grid((total_threads + block_threads - 1) / block_threads); const torch::stable::accelerator::DeviceGuard device_guard( prefix_output.get_device_index()); diff --git a/tests/distributed/test_flashinfer_pcie_all_reduce.py b/tests/distributed/test_flashinfer_pcie_all_reduce.py index 84ee37c2d155..85c3be6a66ef 100644 --- a/tests/distributed/test_flashinfer_pcie_all_reduce.py +++ b/tests/distributed/test_flashinfer_pcie_all_reduce.py @@ -21,6 +21,7 @@ def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs self.destroyed = False self.last_input: torch.Tensor | None = None + self.prepared: list[tuple[list[tuple[int, ...]], torch.dtype]] = [] FakeWorkspace.instances.append(self) def supports(self, inp: torch.Tensor) -> bool: @@ -35,6 +36,9 @@ def all_reduce( out.copy_(inp) return out + def prepare(self, shapes, *, dtype) -> None: + self.prepared.append((list(shapes), dtype)) + def destroy(self) -> None: self.destroyed = True @@ -101,6 +105,7 @@ def test_capture_routes_graph_calls_without_reusing_eager_state() -> None: assert torch.equal(actual, inp) assert len(FakeWorkspace.instances) == 1 assert FakeWorkspace.instances[0].last_input is inp + assert FakeWorkspace.instances[0].prepared == [([(1, 4)], torch.float32)] pool.close() diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 62eba4843a0b..dd2be5cfd6a3 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -257,7 +257,9 @@ def all_gatherv_worker_fn(): device = f"cuda:{pynccl_comm.rank}" assert world_size <= 8 - sizes = [81, 20, 57, 52, 81, 5, 49, 49][:world_size] + # A zero-length rank is required when fewer multimodal inputs than TP + # ranks are distributed across the model-parallel group. + sizes = [81, 0, 57, 52, 81, 5, 49, 49][:world_size] num_elems = sizes[rank] tensor = torch.arange(num_elems, dtype=torch.float32, device=device) + rank * 100 result = torch.zeros(sum(sizes), dtype=torch.float32, device=device) diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 1394ea9df405..4f8eba22d231 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -99,6 +99,64 @@ def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: assert not output.isnan().any() +@pytest.mark.parametrize("num_tokens", [256, 4096]) +@pytest.mark.parametrize("head_size", [128, 192, 512]) +@torch.inference_mode() +def test_merge_attn_states_cuda_inplace_accumulator( + num_tokens: int, head_size: int +) -> None: + """The CUDA kernel supports a running partial as input and destination. + + Chunked attention folds each suffix partial into one prefix allocation. + Both the attention output and its log-sum-exp tensor therefore alias their + corresponding destinations. The result must match a merge into disjoint + output allocations exactly. The 192-element case is Kimi-K3's chunked + context-merge geometry and does not divide the kernel's 128-thread limit; + the 512-element case spans multiple warps per head. + """ + if not current_platform.is_cuda(): + pytest.skip("The custom merge-attention kernel requires CUDA") + + torch.manual_seed(0) + num_heads = 6 + shape = (num_tokens, num_heads, head_size) + + prefix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + suffix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + prefix_lse = torch.randn( + (num_heads, num_tokens), dtype=torch.float32, device="cuda" + ) + suffix_lse = torch.randn( + (num_heads, num_tokens), dtype=torch.float32, device="cuda" + ) + + reference_output = torch.empty_like(prefix_output) + reference_lse = torch.empty_like(prefix_lse) + merge_attn_states_cuda( + reference_output, + prefix_output, + prefix_lse, + suffix_output, + suffix_lse, + reference_lse, + ) + + inplace_output = prefix_output.clone() + inplace_lse = prefix_lse.clone() + merge_attn_states_cuda( + inplace_output, + inplace_output, + inplace_lse, + suffix_output, + suffix_lse, + inplace_lse, + ) + torch.accelerator.synchronize() + + torch.testing.assert_close(inplace_output, reference_output, rtol=0, atol=0) + torch.testing.assert_close(inplace_lse, reference_lse, rtol=0, atol=0) + + def generate_markdown_table(): global all_case_info table_header = ( diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py new file mode 100644 index 000000000000..09b203fc4050 --- /dev/null +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which value the DFlash drafter is fed under AttnRes. + +`_capture_aux_hidden_stream` picks the weights it mixes against from one of +three places depending on where the tapped layer sits, and returns the plain +running prefix when the feature is off. The mixture itself is the kernel's +job and is covered by ``test_attn_res.py``; what is asserted here is the +selection, which is the part that can silently feed the drafter the wrong +tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia import model as k3_model + +END_LAYER = 4 + + +def _weights(tag: float) -> SimpleNamespace: + """A norm/projection pair that is identifiable by value.""" + return SimpleNamespace( + weight=torch.full((2,), tag), + variance_epsilon=tag, + ) + + +def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: + """A stand-in carrying only what the tap reads. + + Constructing the real model needs a distributed init and weights, and none + of it participates in the selection under test. + """ + model = SimpleNamespace( + _aux_attn_res_stream=enabled, + use_attn_res=use_attn_res, + end_layer=END_LAYER, + ) + if not use_attn_res: + return model + + consumers = [] + for i in range(END_LAYER): + consumers.append( + SimpleNamespace( + self_attention_res_norm=_weights(float(i)), + self_attention_res_proj=SimpleNamespace( + weight=torch.full((1, 2), float(i)) + ), + prev_valid_blocks=i, + ) + ) + model.layers = consumers + model.output_attn_res_norm = _weights(99.0) + model.output_attn_res_proj = SimpleNamespace(weight=torch.full((1, 2), 99.0)) + model.num_attn_res_blocks = 99 + return model + + +@pytest.fixture +def recorder(monkeypatch): + """Replace the kernel so the call it would have made is inspectable.""" + calls = [] + + def _fake_attn_res( + prefix, + delta, + block_residual, + norm_weight, + proj_weight, + output_norm_weight, + **kwargs, + ): + calls.append( + SimpleNamespace( + prefix=prefix, + delta=delta, + block_residual=block_residual, + norm_weight=norm_weight, + proj_weight=proj_weight, + kwargs=kwargs, + ) + ) + return torch.full_like(prefix, -1.0) + + monkeypatch.setattr(k3_model, "attn_res", _fake_attn_res) + return calls + + +def _set_last_rank(monkeypatch, is_last: bool): + monkeypatch.setattr( + k3_model, + "get_pp_group", + lambda: SimpleNamespace(is_last_rank=is_last), + ) + + +def _call(stub, layer_idx, prefix_sum, pending_mlp_out, block_residual): + return k3_model.KimiLinearModel._capture_aux_hidden_stream( + stub, layer_idx, prefix_sum, pending_mlp_out, block_residual + ) + + +@pytest.mark.parametrize( + "enabled,use_attn_res", [(False, True), (True, False), (False, False)] +) +def test_disabled_reproduces_the_plain_residual_sum( + recorder, monkeypatch, enabled, use_attn_res +): + """Off, the tap must be exactly the sum it replaced. + + Both conditions matter. `use_attn_res` is what constructs the norm and + projection weights, so without it the lookups below would raise rather + than fall back. + """ + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + got = _call( + _stub_model(enabled=enabled, use_attn_res=use_attn_res), + 0, + prefix_sum, + pending, + torch.zeros(2), + ) + + torch.testing.assert_close(got, prefix_sum + pending) + assert not recorder, "the kernel must not run when the tap is off" + + +def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): + """The value the next layer reads is the mixture against *its* weights, + so the tap has to reach forward rather than use the current layer's.""" + _set_last_rank(monkeypatch, True) + + _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) + + assert len(recorder) == 1 + call = recorder[0] + # Layer 2's weights, not layer 1's. + torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0)) + torch.testing.assert_close(call.proj_weight, torch.full((2,), 2.0)) + assert call.kwargs["num_blocks"] == 2 + + +def test_last_layer_on_the_final_rank_uses_the_output_aggregation( + recorder, monkeypatch +): + """Nothing downstream but the model's own output-side mixture.""" + _set_last_rank(monkeypatch, True) + + _call( + _stub_model(enabled=True), END_LAYER - 1, torch.zeros(2), None, torch.zeros(2) + ) + + assert len(recorder) == 1 + torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0)) + torch.testing.assert_close(recorder[0].proj_weight, torch.full((2,), 99.0)) + assert recorder[0].kwargs["num_blocks"] == 99 + + +def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch): + """The consumer lives on the next rank and the output aggregation only + exists on the last one, so there is nothing here to mix against. + + This is the case that would otherwise reach for weights this rank never + constructs. The forward guard is `layer_idx + 1 < end_layer`, where + `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a + `PPMissingLayer` is unreachable by construction -- the fallback below is + what makes that true rather than merely likely. + """ + _set_last_rank(monkeypatch, False) + prefix_sum = torch.tensor([3.0, 4.0]) + + got = _call( + _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2) + ) + + torch.testing.assert_close(got, prefix_sum) + assert not recorder, "no weights exist on this rank to mix against" + + +def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( + recorder, monkeypatch +): + """The kernel writes an applied delta back into the prefix in place, which + would double-add it into the live residual stream, so the pending output + has to arrive already summed into the prefix with `delta` left None.""" + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2)) + + assert len(recorder) == 1 + assert recorder[0].delta is None + torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending) + # And the caller's tensor is not mutated on the way. + torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0])) diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 7af025c91ae4..0356e304162c 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -22,6 +22,7 @@ def _make_kimi_linear_model() -> KimiLinearModel: object.__setattr__(model, "aux_hidden_state_layers", (2,)) object.__setattr__(model, "use_sequence_parallel", False) object.__setattr__(model, "reuse_attn_res_output", True) + object.__setattr__(model, "use_attn_res", False) return model @@ -190,6 +191,63 @@ def finish_auxiliary_stream(self): ] assert len(aux_hidden_states) == 1 assert aux_hidden_states[0] is projected + + +def test_attn_res_stream_capture_receives_layer_outputs_in_order(monkeypatch): + """Verify the positional contract between ``forward`` and the capture tap.""" + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + captured = torch.tensor([[11.0, 12.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (1,)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr(kimi_model, "attn_res", Mock(return_value=torch.zeros(1, 2))) + monkeypatch.setenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "1") + + capture = Mock(return_value=captured) + monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture) + + _, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args + assert layer_idx == 0 + assert got_prefix is prefix_sum + assert got_pending is layer_hidden_states + assert got_residual is block_residual + torch.testing.assert_close(aux_hidden_states[0], captured) + + def test_kimi_attn_res_workspace_is_reused_and_sliced(): model = _make_kimi_linear_model() object.__setattr__(model, "num_attn_res_blocks", 3) @@ -273,6 +331,7 @@ def _make_attn_res_decoder_layer(*, block_write: bool): object.__setattr__(layer, "use_attn_res", True) object.__setattr__(layer, "reuse_attn_res_output", True) object.__setattr__(layer, "is_block_write_layer", block_write) + object.__setattr__(layer, "is_final_block_write_layer", False) object.__setattr__(layer, "block_write_idx", 0) object.__setattr__(layer, "prev_valid_blocks", 0) object.__setattr__( @@ -361,3 +420,28 @@ def record_attn_res(*args, **kwargs): assert regular_hidden is attention_output assert regular_prefix is prefix assert outputs == [prefix, attention_output] + + +def test_kimi_post_attn_norm_preserves_final_committed_block(monkeypatch): + prefix = torch.randn(2, 4) + attention_output = torch.randn(2, 4) + blocks = torch.randn(2, 3, 4) + allocated_output = torch.randn(2, 4) + outputs = [] + + def record_attn_res(*args, **kwargs): + outputs.append(kwargs["output"]) + return allocated_output + + monkeypatch.setattr(kimi_model, "attn_res", record_attn_res) + layer = _make_attn_res_decoder_layer(block_write=True) + object.__setattr__(layer, "is_final_block_write_layer", True) + + hidden, next_prefix, next_blocks = layer._post_attn_norm( + attention_output, blocks, prefix + ) + + assert hidden is allocated_output + assert next_prefix is attention_output + assert next_blocks is blocks + assert outputs == [None] diff --git a/tests/models/kimi_k3/test_mla_padding.py b/tests/models/kimi_k3/test_mla_padding.py index 8a2d859b931b..007556801152 100644 --- a/tests/models/kimi_k3/test_mla_padding.py +++ b/tests/models/kimi_k3/test_mla_padding.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch @@ -122,6 +123,33 @@ def write_active_prefill(*args): torch.testing.assert_close(output[2:], torch.zeros_like(output[2:])) +@pytest.mark.parametrize("query_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_kimi_mla_context_output_reuses_consumed_query_bytes(query_dtype): + from vllm.models.kimi_k3.nvidia import mla + + query = torch.empty((4, 2, 256), dtype=query_dtype) + output = torch.randn((4, 2, 128), dtype=torch.bfloat16) + + compact = mla._reuse_consumed_query_for_context_output(query, output) + compact.copy_(output) + + assert compact.data_ptr() == query.data_ptr() + assert compact.shape == output.shape + assert compact.dtype == output.dtype + assert compact.is_contiguous() + torch.testing.assert_close(compact, output) + + +def test_kimi_mla_context_output_rejects_insufficient_query_storage(): + from vllm.models.kimi_k3.nvidia import mla + + query = torch.empty((4, 2, 64), dtype=torch.bfloat16) + output = torch.empty((4, 2, 128), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="too small"): + mla._reuse_consumed_query_for_context_output(query, output) + + def test_kimi_mla_caller_output_selection_preserves_decode_and_sp_paths(): from vllm.models.kimi_k3.nvidia import mla diff --git a/tests/models/kimi_k3/test_vision_projector.py b/tests/models/kimi_k3/test_vision_projector.py index 9f2750809a40..c89c79a3b4ce 100644 --- a/tests/models/kimi_k3/test_vision_projector.py +++ b/tests/models/kimi_k3/test_vision_projector.py @@ -80,11 +80,11 @@ def __init__(self): requires_grad=False, ) self.pre_norm = nn.LayerNorm(4, dtype=torch.bfloat16) - self.input_dtype: torch.dtype | None = None + self.inputs: list[torch.Tensor] = [] def forward(self, inputs: torch.Tensor) -> torch.Tensor: - self.input_dtype = inputs.dtype - return inputs + self.inputs.append(inputs) + return inputs + len(self.inputs) def test_kimi_projector_uses_norm_activation_dtype_for_fp8_weights(): @@ -95,9 +95,43 @@ def test_kimi_projector_uses_norm_activation_dtype_for_fp8_weights(): [torch.randn(2, 4), torch.randn(1, 4)], ) - assert projector.input_dtype == torch.bfloat16 + assert len(projector.inputs) == 2 + assert [inputs.shape for inputs in projector.inputs] == [(2, 4), (1, 4)] + assert all(inputs.dtype == torch.bfloat16 for inputs in projector.inputs) assert [output.shape for output in outputs] == [(2, 4), (1, 4)] assert all(output.dtype == torch.bfloat16 for output in outputs) + assert torch.equal(outputs[0], projector.inputs[0] + 1) + assert torch.equal(outputs[1], projector.inputs[1] + 2) + + +def test_kimi_projector_rejects_empty_vision_output(): + with pytest.raises( + ValueError, match="Kimi vision projection requires at least one image feature" + ): + mm_projector_forward(_SerializedFp8Projector(), []) + + +class _DeterministicProjector(nn.Module): + def __init__(self): + super().__init__() + self.pre_norm = nn.LayerNorm(4) + self.linear = nn.Linear(4, 3, bias=False) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return self.linear(self.pre_norm(inputs)) + + +def test_kimi_projector_preserves_batched_projection_results(): + torch.manual_seed(1) + projector = _DeterministicProjector() + inputs = [torch.randn(2, 4), torch.randn(3, 4)] + expected = torch.split(projector(torch.cat(inputs)), [2, 3]) + + outputs = mm_projector_forward(projector, inputs) + + assert len(outputs) == len(expected) + for output, reference in zip(outputs, expected): + torch.testing.assert_close(output, reference) def test_kimi_vision_rope_reuses_packed_qk_buffers(): diff --git a/tests/models/kimi_k3/test_vision_warmup.py b/tests/models/kimi_k3/test_vision_warmup.py index f7e4af39839e..bdd5b88a8f91 100644 --- a/tests/models/kimi_k3/test_vision_warmup.py +++ b/tests/models/kimi_k3/test_vision_warmup.py @@ -13,6 +13,43 @@ ) +def _full_grid_rope_reference( + rope: kimi_k25_vit.Rope2DPosEmbRepeated, + shapes: list[list[int]], +) -> torch.Tensor: + flat_pos = torch.arange(rope.max_height * rope.max_width).float() + x_pos = flat_pos % rope.max_width + y_pos = flat_pos // rope.max_width + dim_range = torch.arange(0, rope.dim, 4).float() + freqs = 1.0 / (rope.theta_base ** (dim_range / rope.dim)) + x_freqs = torch.outer(x_pos, freqs).float() + y_freqs = torch.outer(y_pos, freqs).float() + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) + table = torch.cat( + [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 + ).reshape(rope.max_height, rope.max_width, -1) + return torch.cat( + [table[:h, :w].reshape(-1, rope.dim // 2).repeat(t, 1) for t, h, w in shapes] + ) + + +def test_vision_rope_materializes_only_requested_grids() -> None: + rope = kimi_k25_vit.Rope2DPosEmbRepeated( + dim=32, + max_height=11, + max_width=13, + ) + shapes = [[1, 3, 5], [2, 4, 2], [1, 3, 5]] + + actual = rope.get_freqs_cis(shapes, device=torch.device("cpu")) + expected = _full_grid_rope_reference(rope, shapes) + + assert torch.equal(actual, expected) + assert actual.shape == (46, 16) + assert not hasattr(rope, "freqs_cis") + + def test_warm_vision_position_interpolation(monkeypatch) -> None: model = torch.nn.Sequential( kimi_k25_vit.Learnable2DInterpPosEmbDivided_fixed( diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py index 4f831e851e9b..83ee0b5d0951 100644 --- a/tests/reasoning/test_kimi_k3_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -52,6 +52,7 @@ def test_parser_selection_thinking_disabled(): ) assert parser._thinking_enabled is False + assert parser.thinking_enabled is False def test_extract_reasoning_with_xtml_tags(): @@ -119,6 +120,103 @@ def test_is_reasoning_end_ignores_stale_close_from_prior_turn(): assert not parser.is_reasoning_end([*new_open]) +def test_fresh_assistant_prompt_does_not_inherit_closed_reasoning(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": True, + "continue_final_message": False, + }, + ) + + assert not parser.is_reasoning_end_for_prompt(CLOSE_IDS) + + +def test_continued_assistant_prompt_uses_rendered_reasoning_state(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": False, + "continue_final_message": True, + }, + ) + + assert parser.is_reasoning_end_for_prompt(CLOSE_IDS) + assert not parser.is_reasoning_end_for_prompt(OPEN_IDS) + + +def test_fresh_assistant_stream_classifies_first_tokens_as_reasoning(): + parser = ReasoningOnlyParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": True, + }, + ) + request = ChatCompletionRequest(model="test-model", messages=[]) + + first = parser.parse_delta( + delta_text=".", + delta_token_ids=[9], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=False, + ) + partial_close = parser.parse_delta( + delta_text=f"{CLOSE}think", + delta_token_ids=CLOSE_IDS[:2], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=False, + ) + closed = parser.parse_delta( + delta_text=f"{SEP}{RESPONSE_OPEN}", + delta_token_ids=[CLOSE_IDS[2], 10], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=False, + ) + + assert first is not None + assert first.reasoning == "." + assert first.content is None + assert partial_close is None + assert closed is None + + +def test_content_filter_holds_and_removes_split_protocol_markers(): + parser = ReasoningOnlyParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": False, + "continue_final_message": True, + }, + ) + request = ChatCompletionRequest(model="test-model", messages=[]) + chunks = [".", f"{CLOSE}think", SEP, RESPONSE_OPEN] + messages: list[DeltaMessage] = [] + + for index, chunk in enumerate(chunks): + delta = parser.parse_delta( + delta_text=chunk, + delta_token_ids=[9 + index], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=index == len(chunks) - 1, + ) + if delta is not None: + messages.append(delta) + + content = "".join(message.content or "" for message in messages) + assert content == "." + assert OPEN not in content + assert CLOSE not in content + assert SEP not in content + + def test_streaming_split_open_marker_is_held_back(): parser = KimiK3ReasoningParser(DummyTokenizer()) diff --git a/tests/tool_use/test_kimi_k3_tool_parser.py b/tests/tool_use/test_kimi_k3_tool_parser.py index 118c51db574e..43f29edfc32a 100644 --- a/tests/tool_use/test_kimi_k3_tool_parser.py +++ b/tests/tool_use/test_kimi_k3_tool_parser.py @@ -171,6 +171,48 @@ def test_delegating_parser_preserves_tool_calls_after_reasoning(): assert json.loads(tool_calls[0].arguments) == {"x": 1} +def test_fresh_tool_stream_does_not_inherit_closed_reasoning(): + parser = KimiK3DelegatingParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": True, + }, + ) + request = _request() + messages: list[DeltaMessage] = [] + chunks = [ + (".", [9]), + (f"{CLOSE}think", [4, 2]), + ( + f"{SEP}{_response('')}{_tools(_call('calc', 1))}", + [3, 10], + ), + ] + + for index, (text, token_ids) in enumerate(chunks): + delta = parser.parse_delta( + delta_text=text, + delta_token_ids=token_ids, + request=request, + prompt_token_ids=[4, 2, 3], + finished=index == len(chunks) - 1, + ) + if delta is not None: + messages.append(delta) + + reasoning = "".join(message.reasoning or "" for message in messages) + content = "".join(message.content or "" for message in messages) + tool_calls = [call for message in messages for call in (message.tool_calls or [])] + assert reasoning == "." + assert content == "" + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "calc" + assert OPEN not in reasoning + content + assert CLOSE not in reasoning + content + assert SEP not in reasoning + content + + def test_delegating_parser_required_tool_choice_uses_xtml_parser(): parser = KimiK3DelegatingParser(DummyTokenizer()) request = _request().model_copy(update={"tool_choice": "required"}) @@ -361,9 +403,141 @@ def test_streaming_split_markers_do_not_leak(): assert content == "Hi" assert OPEN not in content assert SEP not in content - assert len(tool_deltas) == 1 - assert tool_deltas[0].function.name == "calc" - assert json.loads(tool_deltas[0].function.arguments) == {"x": 1} + assert [tool_call.function.name for tool_call in tool_deltas if tool_call.id] == [ + "calc" + ] + arguments = "".join(tool_call.function.arguments or "" for tool_call in tool_deltas) + assert json.loads(arguments) == {"x": 1} + + +def test_streaming_emits_argument_text_as_it_arrives(): + """A long string argument must stream, not land in one delta at the close.""" + parser = KimiK3ToolParser(DummyTokenizer()) + request = _request() + value = "word " * 200 + body_chunks = [value[i : i + 5] for i in range(0, len(value), 5)] + chunks = [ + f"{OPEN}tools{SEP}", + f'{OPEN}call tool="write_file" index="1"{SEP}', + f'{OPEN}argument key="content" type="string"{SEP}', + *body_chunks, + f"{CLOSE}argument{SEP}", + f"{CLOSE}call{SEP}", + ] + previous_text = "" + previous_ids: list[int] = [] + messages: list[DeltaMessage] = [] + arguments_before_close = "" + + for i, chunk in enumerate(chunks, start=1): + if chunk == f"{CLOSE}argument{SEP}": + arguments_before_close = "".join( + tool_call.function.arguments or "" + for message in messages + for tool_call in (message.tool_calls or []) + ) + current_text = previous_text + chunk + current_ids = previous_ids + [i] + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=chunk, + previous_token_ids=previous_ids, + current_token_ids=current_ids, + delta_token_ids=[i], + request=request, + ) + if delta is not None: + messages.append(delta) + previous_text = current_text + previous_ids = current_ids + + tool_deltas = [ + tool_call for message in messages for tool_call in (message.tool_calls or []) + ] + + # the name is announced once, up front, and every body chunk moves the stream + assert [tool_call.function.name for tool_call in tool_deltas if tool_call.id] == [ + "write_file" + ] + assert len(tool_deltas) >= len(body_chunks) + assert all(tool_call.index == 0 for tool_call in tool_deltas) + assert ( + arguments_before_close + == json.dumps({"content": value}, ensure_ascii=False)[:-2] + ) + + arguments = "".join(tool_call.function.arguments or "" for tool_call in tool_deltas) + assert json.loads(arguments) == {"content": value} + non_streamed = parser.extract_tool_calls(previous_text, request) + assert arguments == non_streamed.tool_calls[0].function.arguments + + +def test_streaming_holds_whitespace_tolerant_argument_close_fragments(): + parser = KimiK3ToolParser(DummyTokenizer()) + request = _request() + chunks = [ + f"{OPEN}tools{SEP}", + f'{OPEN}call tool="write_file" index="1"{SEP}', + f'{OPEN}argument key="content" type="string"{SEP}', + "payload", + f"{CLOSE} arg", + "ument ", + "<|sep", + "|>", + f"{CLOSE}call{SEP}", + ] + previous_text = "" + previous_ids: list[int] = [] + streamed_arguments = "" + partial_close_snapshots: list[str] = [] + + for i, chunk in enumerate(chunks, start=1): + current_text = previous_text + chunk + current_ids = previous_ids + [i] + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=chunk, + previous_token_ids=previous_ids, + current_token_ids=current_ids, + delta_token_ids=[i], + request=request, + ) + if delta is not None: + streamed_arguments += "".join( + tool_call.function.arguments or "" + for tool_call in (delta.tool_calls or []) + ) + if 5 <= i <= 7: + partial_close_snapshots.append(streamed_arguments) + previous_text = current_text + previous_ids = current_ids + + expected_prefix = json.dumps({"content": "payload"}, ensure_ascii=False)[:-2] + assert partial_close_snapshots == [expected_prefix] * 3 + assert json.loads(streamed_arguments) == {"content": "payload"} + non_streamed = parser.extract_tool_calls(previous_text, request) + assert streamed_arguments == non_streamed.tool_calls[0].function.arguments + + +def test_streaming_ignores_call_shaped_text_after_tools_close(): + parser = KimiK3ToolParser(DummyTokenizer()) + request = _request() + output = _tools() + _call("calc", 1, _arg("x", "number", "1")) + + delta = parser.extract_tool_calls_streaming( + previous_text="", + current_text=output, + delta_text=output, + previous_token_ids=[], + current_token_ids=[1], + delta_token_ids=[1], + request=request, + ) + + assert delta is None + assert parser.extract_tool_calls(output, request).tools_called is False def test_tool_call_ids_are_unique_across_messages(): diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 8f43039face2..54e26459ceeb 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -514,6 +514,79 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): assert moved[0].block_hash_num_tokens == 6 +def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): + """An external mid-block hit must become a running request even when its + first continuation does not need another Mamba block.""" + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + request = make_request("0", [0] * 15, hash_block_size, sha256) + loaded_blocks = manager.allocate_slots( + request, + num_new_tokens=0, + num_external_computed_tokens=10, + delay_cache_blocks=True, + ) + assert loaded_blocks is not None + + request.num_computed_tokens = 10 + first_step_blocks = manager.allocate_slots(request, num_new_tokens=4) + assert first_step_blocks is not None + assert first_step_blocks.get_block_ids()[1] == [] + + source_block_id = manager.get_blocks("0").get_block_ids()[1][1] + partial_hash = request.block_hashes[14 // hash_block_size - 1] + partial_block = manager.block_pool.get_cached_block( + partial_hash, kv_cache_group_ids=[1] + ) + assert partial_block is not None + assert partial_block[0].block_id == source_block_id + + request.num_computed_tokens = 14 + continuation_blocks = manager.allocate_slots(request, num_new_tokens=1) + assert continuation_blocks is not None + + assert continuation_blocks.get_block_ids()[1] == [] + assert manager.get_blocks("0").get_block_ids()[1][1] == source_block_id + copies, _ = manager.take_kv_cache_block_copies() + cow_copy = next(c for c in copies if c.src_block_id == source_block_id) + assert cow_copy.dst_block_id != source_block_id + + moved = manager.block_pool.get_cached_block(partial_hash, kv_cache_group_ids=[1]) + assert moved is not None + assert moved[0].block_id == cow_copy.dst_block_id + + def test_take_partial_tail_offloads_returns_cow_target(): """The connector offload hand-off exposes the mamba CoW *target* block Y (the durable boundary state), not the overwritten source X, and only at diff --git a/tests/v1/core/test_dspark_prefix_cache_policy.py b/tests/v1/core/test_dspark_prefix_cache_policy.py new file mode 100644 index 000000000000..482e5ed722ec --- /dev/null +++ b/tests/v1/core/test_dspark_prefix_cache_policy.py @@ -0,0 +1,41 @@ +from types import SimpleNamespace + +from vllm.v1.core.sched import scheduler as scheduler_module + + +def _spec(method: str, use_eagle: bool = True): + return SimpleNamespace(method=method, use_eagle=lambda: use_eagle) + + +def _groups(*flags: bool): + return [SimpleNamespace(is_eagle_group=flag) for flag in flags] + + +def test_dspark_without_target_eagle_group_does_not_drop_target_cache_tail(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("dspark"), _groups(False, False)) is False + + +def test_dflash_without_target_eagle_group_does_not_drop_target_cache_tail(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("dflash"), _groups(False, False)) is False + + +def test_explicit_target_eagle_group_keeps_cache_drop(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("dspark"), _groups(False, True)) is True + + +def test_classic_eagle_keeps_legacy_unannotated_fallback(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("eagle3"), _groups(False, False)) is True + + +def test_non_eagle_speculation_never_drops_target_cache_tail(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("ngram", use_eagle=False), _groups(True)) is False diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index f080d6aa03e3..517bab5fd65d 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -58,6 +58,7 @@ SlidingWindowMLASpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, + get_kv_cache_dcp_shard_count, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, ) @@ -197,6 +198,16 @@ def new_mamba_spec( ) +def test_mamba_cache_has_one_dcp_token_position_shard(): + spec = new_mamba_spec(block_size=768, num_speculative_blocks=7) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(mamba_cache_mode="align") + ) + + assert get_kv_cache_dcp_shard_count(spec, dcp_world_size=16) == 1 + assert spec.max_num_blocks_per_req(vllm_config, max_len=1_000_000) == 1310 + + def test_unify_kv_cache_spec_page_size_uses_lcm_for_non_divisible_pages(): mimo_spec = FullAttentionSpec( block_size=64, @@ -1309,8 +1320,11 @@ def test_uniform_type_spec_block_table_width_matches_layer_spec( # The runner sizes the block table from the group spec while the metadata # builders are constructed from the per-layer spec, so the aggregate must # report the same width as the layers it wraps. - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) - vllm_config.parallel_config.decode_context_parallel_size = dcp_size + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size), + cache_config=SimpleNamespace(mamba_cache_mode="none"), + model_config=SimpleNamespace(max_model_len=1024), + ) if layer_type == "mla": layer_spec = new_mla_spec() elif layer_type == "replicated": diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index e4188e9167d5..d04a1febbb3c 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -839,6 +839,172 @@ def test_stop_via_update_from_output(): assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID, 10, 11] +@pytest.mark.parametrize("async_scheduling", [False, True]) +@pytest.mark.parametrize( + "filtered_tokens,num_grammar_rejected,expected_accepted", + [ + ([10, 11, 12], 1, 3), + ([10, 11], 2, 2), + ([], 4, 0), + ], +) +def test_speculative_grammar_filter_rolls_back_scheduler_state_and_stats( + async_scheduling: bool, + filtered_tokens: list[int], + num_grammar_rejected: int, + expected_accepted: int, +): + """Rejected suffix tokens remain schedulable and are not accepted drafts.""" + scheduler = create_scheduler( + num_speculative_tokens=3, + speculative_method="ngram_gpu", + async_scheduling=async_scheduling, + ) + request = create_requests(num_requests=1)[0] + request.structured_output_request = Mock() + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + 4 + request.num_output_placeholders = 4 if async_scheduling else 0 + scheduler.requests[request.request_id] = request + scheduler.running.append(request) + + manager = Mock() + manager.filter_speculative_grammar_tokens.return_value = ( + filtered_tokens, + num_grammar_rejected, + ) + manager.should_advance.return_value = False + scheduler.structured_output_manager = manager + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 4}, + total_num_scheduled_tokens=4, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={request.request_id: [10, 11, 12]}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[10, 11, 12, 13]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + outputs = scheduler.update_from_output(scheduler_output, model_output) + + assert request.num_computed_tokens == request.num_tokens + assert request.num_output_placeholders == 0 + assert list(request.output_token_ids) == filtered_tokens + if filtered_tokens: + assert outputs[0].outputs[0].new_token_ids == filtered_tokens + manager.filter_speculative_grammar_tokens.assert_called_once_with( + request, [10, 11, 12, 13] + ) + stats = outputs[0].scheduler_stats.spec_decoding_stats + assert stats is not None + assert stats.num_drafts == 1 + assert stats.num_draft_tokens == 3 + assert stats.num_accepted_tokens == expected_accepted + assert stats.num_accepted_tokens_per_pos == [ + int(position < expected_accepted) for position in range(3) + ] + + +def test_speculative_grammar_filter_is_not_called_for_unstructured_requests(): + """Unstructured speculative decoding does not enter grammar validation.""" + scheduler = create_scheduler(num_speculative_tokens=2) + request = create_requests(num_requests=1)[0] + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + 3 + scheduler.requests[request.request_id] = request + scheduler.running.append(request) + + manager = Mock() + manager.should_advance.return_value = False + scheduler.structured_output_manager = manager + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 3}, + total_num_scheduled_tokens=3, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={request.request_id: [10, 11]}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[10, 11, 12]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + scheduler.update_from_output(scheduler_output, model_output) + + assert list(request.output_token_ids) == [10, 11, 12] + manager.filter_speculative_grammar_tokens.assert_not_called() + + +def test_speculative_grammar_filter_commits_and_advances_only_valid_prefix(): + """Scheduler output and grammar state share the filtered token prefix.""" + scheduler = create_scheduler(num_speculative_tokens=2) + request = create_requests(num_requests=1)[0] + grammar = Mock(spec=StructuredOutputGrammar) + grammar.accept_tokens.side_effect = lambda request_id, tokens: tokens == [10, 11] + request.structured_output_request = Mock(grammar=grammar) + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + 3 + scheduler.requests[request.request_id] = request + scheduler.running.append(request) + + manager = Mock() + manager.filter_speculative_grammar_tokens.return_value = ([10, 11], 1) + manager.should_advance.return_value = True + manager.trim_reasoning_for_advance.return_value = [10, 11] + scheduler.structured_output_manager = manager + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 3}, + total_num_scheduled_tokens=3, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={request.request_id: [10, 11]}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[10, 11, 12]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + outputs = scheduler.update_from_output(scheduler_output, model_output) + + assert request.status == RequestStatus.RUNNING + assert request.num_computed_tokens == request.num_tokens + assert list(request.output_token_ids) == [10, 11] + assert outputs[0].outputs[0].new_token_ids == [10, 11] + grammar.accept_tokens.assert_called_once_with(request.request_id, [10, 11]) + stats = outputs[0].scheduler_stats.spec_decoding_stats + assert stats is not None + assert stats.num_draft_tokens == 2 + assert stats.num_accepted_tokens == 2 + + def test_check_stop_min_tokens(): """Test that requests don't stop when min_tokens requirement isn't met.""" from vllm.v1.core.sched.utils import check_stop diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 77d629729776..14cfec017e67 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -15,8 +15,18 @@ from unittest.mock import Mock import pytest +import torch +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import ( + NixlBaseConnector, +) from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + SlidingWindowSpec, +) from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( @@ -24,6 +34,7 @@ create_request, create_scheduler, create_vllm_config, + make_kv_cache_config, ) pytestmark = pytest.mark.cpu_test @@ -56,6 +67,121 @@ def recompute_scheduler(): return create_scheduler(vllm_config) +def _create_invalid_block_test_scheduler( + scheduler_block_size: int, + group_block_sizes: tuple[int, ...], +) -> Scheduler: + """Create the scheduler state required by invalid-block mapping tests.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.block_size = scheduler_block_size + scheduler.kv_cache_manager = Mock() + scheduler.kv_cache_config = KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + [f"full_attention_{group_idx}"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + for group_idx, block_size in enumerate(group_block_sizes) + ], + ) + return scheduler + + +def test_hybrid_cache_invalid_block_truncates_and_evicts_all_groups(): + """A failed hybrid-cache group invalidates the shared logical suffix.""" + scheduler = _create_invalid_block_test_scheduler(16, (16, 16)) + scheduler.kv_cache_manager.get_block_ids.return_value = ( + [1, 2, 3, 4], + [11, 12, 13, 14], + ) + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.return_value = ( + [1, 2, 3, 4], + [11, 12, 13, 14], + ) + + request = Mock(spec=Request) + request.request_id = "hybrid-request" + request.num_computed_tokens = 64 + + affected, affected_tokens, evicted = scheduler._update_requests_with_invalid_blocks( + [request], + invalid_block_ids={12}, + num_scheduled_tokens={}, + ) + + assert affected == {request.request_id} + assert request.num_computed_tokens == 16 + assert affected_tokens == 48 + assert evicted == {2, 3, 4, 12, 13, 14} + + +def test_hybrid_cache_shared_invalid_block_is_recomputed_once(): + """A shared failed block contributes one recomputation interval.""" + scheduler = _create_invalid_block_test_scheduler(16, (16, 16)) + scheduler.kv_cache_manager.get_block_ids.side_effect = ( + ([1, 2, 3], [11, 12, 13]), + ([21, 22, 23], [31, 12, 33]), + ) + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.side_effect = ( + ([1, 2, 3], [11, 12, 13]), + ([21, 22, 23], [31, 12, 33]), + ) + + first_request = Mock(spec=Request) + first_request.request_id = "first-request" + first_request.num_computed_tokens = 48 + second_request = Mock(spec=Request) + second_request.request_id = "second-request" + second_request.num_computed_tokens = 48 + + affected, affected_tokens, evicted = scheduler._update_requests_with_invalid_blocks( + [first_request, second_request], + invalid_block_ids={12}, + num_scheduled_tokens={}, + ) + + assert affected == {first_request.request_id, second_request.request_id} + assert first_request.num_computed_tokens == 16 + assert second_request.num_computed_tokens == 48 + assert affected_tokens == 32 + assert evicted == {2, 3, 12, 13} + + +def test_hybrid_cache_ignores_invalid_blocks_after_external_prefix(): + """Blocks used only by scheduled local tokens are not load failures.""" + scheduler = _create_invalid_block_test_scheduler(32, (16, 32)) + scheduler.kv_cache_manager.get_block_ids.return_value = ( + [1, 2, 3, 4], + [11, 12], + ) + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.return_value = ( + [1, 2], + [11], + ) + + request = Mock(spec=Request) + request.request_id = "local-suffix" + request.num_computed_tokens = 64 + + affected, affected_tokens, evicted = scheduler._update_requests_with_invalid_blocks( + [request], + invalid_block_ids={3, 12}, + num_scheduled_tokens={request.request_id: 32}, + ) + + assert affected == set() + assert request.num_computed_tokens == 64 + assert affected_tokens == 0 + assert evicted == set() + + def test_sync_recompute_blocks_not_freed_for_running_requests( recompute_scheduler: Scheduler, ): @@ -478,3 +604,296 @@ def cache_blocks_spy(req, num_tokens): # request should be in the running queue assert request in recompute_scheduler.running + + +def test_sync_recompute_handles_invalid_block_in_second_kv_cache_group(): + """Invalid blocks in any hybrid KV group must trigger recomputation.""" + block_size = 16 + num_blocks = 128 + num_prompt_blocks = 8 + num_external_computed_blocks = 7 + invalid_block_idx = 3 + + vllm_config = create_vllm_config( + block_size=block_size, + kv_load_failure_policy="recompute", + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=make_kv_cache_config( + block_size=block_size, + swa_enabled=True, + sw_size=num_prompt_blocks * block_size, + num_blocks=num_blocks, + ), + ) + + request = create_request( + num_tokens=num_prompt_blocks * block_size, + block_size=block_size, + ) + scheduler.add_request(request) + + num_external_computed_tokens = num_external_computed_blocks * block_size + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + { + request.request_id: num_external_computed_tokens, + }, + False, + ) + ) + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.request_finished_all_groups.return_value = ( + False, + None, + ) + scheduler.connector.take_events.return_value = () + + scheduler_output = scheduler.schedule() + + assert request.status == RequestStatus.RUNNING + assert len(scheduler_output.scheduled_new_reqs) == 1 + + block_ids_by_group = scheduler_output.scheduled_new_reqs[0].block_ids + assert len(block_ids_by_group) == 2 + assert all( + len(group_block_ids) > invalid_block_idx + for group_block_ids in block_ids_by_group + ) + + # Report a load failure in the second KV cache group, not the first. + invalid_block_id = block_ids_by_group[1][invalid_block_idx] + model_runner_output = create_model_runner_output( + [request], + invalid_block_ids={invalid_block_id}, + use_eos=False, + ) + + scheduler.update_from_output( + scheduler_output, + model_runner_output, + ) + + assert request.status == RequestStatus.RUNNING + assert request.num_computed_tokens == invalid_block_idx * block_size + assert request in scheduler.running + assert request.request_id in scheduler.requests + + +def test_sync_fail_handles_invalid_block_in_seventeenth_kv_cache_group(): + """A hybrid-cache load failure must fail one request, not the scheduler.""" + block_size = 16 + num_blocks = 512 + num_prompt_blocks = 8 + num_external_computed_blocks = 7 + invalid_block_idx = 3 + num_kv_cache_groups = 17 + + vllm_config = create_vllm_config( + block_size=block_size, + kv_load_failure_policy="fail", + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + *[ + KVCacheGroupSpec( + [f"full_attention_layer_{group_idx}"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + for group_idx in range(num_kv_cache_groups - 1) + ], + KVCacheGroupSpec( + ["sliding_window_layer"], + SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=num_prompt_blocks * block_size, + ), + ), + ], + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + ) + + failed_request = create_request( + num_tokens=num_prompt_blocks * block_size, + block_size=block_size, + ) + scheduler.add_request(failed_request) + + num_external_computed_tokens = num_external_computed_blocks * block_size + scheduler.connector = Mock(spec=NixlBaseConnector) + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + { + failed_request.request_id: num_external_computed_tokens, + }, + False, + ) + ) + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.request_finished_all_groups.return_value = ( + False, + None, + ) + scheduler.connector.take_events.return_value = () + + scheduler_output = scheduler.schedule() + block_ids_by_group = scheduler_output.scheduled_new_reqs[0].block_ids + + assert len(block_ids_by_group) == num_kv_cache_groups + invalid_block_id = block_ids_by_group[-1][invalid_block_idx] + model_runner_output = create_model_runner_output( + [failed_request], + invalid_block_ids={invalid_block_id}, + use_eos=False, + ) + + outputs = scheduler.update_from_output(scheduler_output, model_runner_output) + + assert failed_request.status == RequestStatus.FINISHED_ERROR + assert failed_request.request_id not in scheduler.requests + assert failed_request not in scheduler.running + assert len(outputs) == 1 + engine_output = next(iter(outputs.values())).outputs[0] + assert engine_output.request_id == failed_request.request_id + assert engine_output.finish_reason == FinishReason.ERROR + + healthy_request = create_request( + num_tokens=2 * block_size, + block_size=block_size, + ) + scheduler.add_request(healthy_request) + next_scheduler_output = scheduler.schedule() + + assert healthy_request.request_id in next_scheduler_output.num_scheduled_tokens + assert healthy_request.status == RequestStatus.RUNNING + + +def test_sync_recompute_handles_mixed_kv_group_block_sizes(): + """Recompute from a common boundary for mixed KV block sizes.""" + hash_block_size = 16 + scheduler_block_size = 64 + num_blocks = 512 + num_prompt_tokens = 8 * scheduler_block_size + num_external_computed_tokens = 7 * scheduler_block_size + + # The invalid block begins at token 3 * 32 = 96. + # Scheduling granularity is 64, so recomputation must restart at 64. + invalid_group_idx = 1 + invalid_block_idx = 3 + + vllm_config = create_vllm_config( + block_size=scheduler_block_size, + kv_load_failure_policy="recompute", + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_16"], + FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["window_32"], + SlidingWindowSpec( + block_size=32, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, + sliding_window=num_prompt_tokens, + ), + ), + KVCacheGroupSpec( + ["full_64"], + FullAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=32, + dtype=torch.float16, + ), + ), + ], + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + hash_block_size=hash_block_size, + ) + + request = create_request( + num_tokens=num_prompt_tokens, + block_size=hash_block_size, + ) + scheduler.add_request(request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + { + request.request_id: num_external_computed_tokens, + }, + False, + ) + ) + scheduler.connector.request_finished.return_value = ( + False, + None, + ) + scheduler.connector.request_finished_all_groups.return_value = ( + False, + None, + ) + scheduler.connector.take_events.return_value = () + + scheduler_output = scheduler.schedule() + + assert request.status == RequestStatus.RUNNING + assert len(scheduler_output.scheduled_new_reqs) == 1 + + block_ids_by_group = scheduler_output.scheduled_new_reqs[0].block_ids + assert len(block_ids_by_group) == 3 + assert len(block_ids_by_group[invalid_group_idx]) > invalid_block_idx + + invalid_block_id = block_ids_by_group[invalid_group_idx][invalid_block_idx] + model_runner_output = create_model_runner_output( + [request], + invalid_block_ids={invalid_block_id}, + use_eos=False, + ) + + scheduler.update_from_output( + scheduler_output, + model_runner_output, + ) + + expected_recompute_from = ( + invalid_block_idx * 32 // scheduler_block_size * scheduler_block_size + ) + + assert request.num_computed_tokens == expected_recompute_from + assert request.status == RequestStatus.RUNNING + assert request in scheduler.running + assert request.request_id in scheduler.requests diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index b1fb43353374..a2908b7ab2e4 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -153,6 +153,7 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, + hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -183,6 +184,7 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, + hash_block_size=hash_block_size, ) diff --git a/tests/v1/spec_decode/test_acceptance_length_controller.py b/tests/v1/spec_decode/test_acceptance_length_controller.py index 9d0cc654a788..2e12b5a7c151 100644 --- a/tests/v1/spec_decode/test_acceptance_length_controller.py +++ b/tests/v1/spec_decode/test_acceptance_length_controller.py @@ -260,6 +260,22 @@ def test_runner_v2_limits_drafts_to_adaptive_depth(): ) +def test_runner_v2_allows_scheduler_to_disable_speculation(): + draft_tokens = torch.tensor([[1, 2, 3], [4, 5, 6]]) + + limited = limit_draft_tokens( + draft_tokens, + num_speculative_tokens=0, + max_num_speculative_tokens=3, + ) + + assert limited.shape == (2, 0) + assert ( + limited.untyped_storage().data_ptr() + == draft_tokens.untyped_storage().data_ptr() + ) + + def test_synthetic_scheduler_output_uses_default_speculative_depth(): output = SchedulerOutput.make_empty() diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py index 02e2a0bdbec2..015281027894 100644 --- a/tests/v1/spec_decode/test_dflash_causality.py +++ b/tests/v1/spec_decode/test_dflash_causality.py @@ -10,11 +10,13 @@ from types import SimpleNamespace import pytest +import torch.nn as nn from vllm.model_executor.models.qwen3_dflash import ( _dflash_layer_causal, _get_dflash_fc_input_size, dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( get_eagle3_aux_layers_from_config, @@ -89,3 +91,105 @@ def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name): assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple( layer_ids ) + + +class _TargetRotaryModule(nn.Module): + def __init__(self, is_neox_style: bool): + super().__init__() + self.is_neox_style = is_neox_style + + +class _TargetModel(nn.Module): + def __init__(self, is_neox_style: bool): + super().__init__() + self.rotary = _TargetRotaryModule(is_neox_style) + + +@pytest.mark.parametrize("is_neox_style", [False, True]) +def test_dflash_target_rope_layout_is_discovered(is_neox_style): + target = _TargetModel(is_neox_style) + + assert dflash_target_rope_is_neox_style(target) is is_neox_style + + +def test_dflash_loader_propagates_target_rope_layout(monkeypatch): + """DFlash configures the target rotary layout before draft construction.""" + from vllm.v1.worker.gpu.spec_decode.dflash import utils as loader_module + + draft_hf_config = SimpleNamespace( + num_hidden_layers=1, + layer_types=["sliding_attention"], + dflash_config={"causal": True}, + ) + speculative_config = SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=draft_hf_config), + attention_backend=None, + kv_cache_dtype=None, + draft_load_config=None, + ) + vllm_config = SimpleNamespace( + speculative_config=speculative_config, + attention_config=SimpleNamespace(), + cache_config=SimpleNamespace(), + quant_config=None, + ) + + def fake_replace(obj, **changes): + values = vars(obj).copy() + values.update(changes) + return SimpleNamespace(**values) + + class DraftConstructionObserved(Exception): + pass + + def fake_get_model(**_kwargs): + assert draft_hf_config.is_neox_style is False + raise DraftConstructionObserved + + monkeypatch.setattr(loader_module, "replace", fake_replace) + monkeypatch.setattr(loader_module, "get_model", fake_get_model) + with pytest.raises(DraftConstructionObserved): + loader_module.load_dflash_model(_TargetModel(is_neox_style=False), vllm_config) + + +def test_dspark_loader_preserves_checkpoint_rope_layout(monkeypatch): + """DSpark inference uses the rotary layout encoded by its training model.""" + from vllm.model_executor.models import utils as model_utils + from vllm.v1.worker.gpu.spec_decode.dspark import utils as loader_module + + draft_hf_config = SimpleNamespace( + num_hidden_layers=1, + layer_types=["sliding_attention"], + dflash_config={"causal": True}, + is_neox_style=True, + ) + speculative_config = SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=draft_hf_config), + attention_backend=None, + kv_cache_dtype=None, + draft_load_config=None, + ) + vllm_config = SimpleNamespace( + speculative_config=speculative_config, + attention_config=SimpleNamespace(), + cache_config=SimpleNamespace(), + quant_config=None, + ) + + class DraftConstructionObserved(Exception): + pass + + def fake_get_model(**_kwargs): + assert draft_hf_config.is_neox_style is True + raise DraftConstructionObserved + + monkeypatch.setattr(loader_module, "get_model", fake_get_model) + monkeypatch.setattr( + loader_module, + "_create_draft_vllm_config", + lambda _config: vllm_config, + ) + monkeypatch.setattr(model_utils, "get_draft_quant_config", lambda _config: None) + + with pytest.raises(DraftConstructionObserved): + loader_module.load_dspark_model(_TargetModel(is_neox_style=False), vllm_config) diff --git a/tests/v1/spec_decode/test_dflash_swa.py b/tests/v1/spec_decode/test_dflash_swa.py index 2a0c9d6d6f7e..092e86492cb5 100644 --- a/tests/v1/spec_decode/test_dflash_swa.py +++ b/tests/v1/spec_decode/test_dflash_swa.py @@ -10,6 +10,7 @@ from vllm.model_executor.models.qwen3_dflash import DFlashAttention from vllm.transformers_utils.configs.speculators import SpeculatorsConfig from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata +from vllm.v1.attention.backends import flash_attn as flash_attn_backend from vllm.v1.kv_cache_interface import ( FullAttentionSpec, SlidingWindowSpec, @@ -161,6 +162,53 @@ def test_dflash_swa_layers_keep_sliding_window_kv_cache_spec(monkeypatch): assert spec.dcp_replicated is True +def test_flash_attention_metadata_treats_replicated_kv_as_dcp1(monkeypatch): + """Replicated draft cache metadata uses global sequence lengths locally.""" + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_replicated=True, + ) + model_config = SimpleNamespace( + get_num_attention_heads=lambda _parallel_config: 96, + get_num_kv_heads=lambda _parallel_config: 16, + get_head_size=lambda: 128, + rswa_window=None, + is_mm_prefix_lm=False, + ) + vllm_config = SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace(cp_kv_cache_interleave_size=1), + cache_config=SimpleNamespace(cache_dtype="bfloat16"), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False), + max_cudagraph_capture_size=None, + ), + attention_config=SimpleNamespace( + flash_attn_max_num_splits_for_cuda_graph=0, + ), + scheduler_config=SimpleNamespace(max_num_seqs=1), + ) + + monkeypatch.setattr( + flash_attn_backend, + "get_dcp_group", + lambda: SimpleNamespace(world_size=16, rank_in_group=7), + ) + builder = flash_attn_backend.FlashAttentionMetadataBuilder( + spec, + ["draft.layer"], + vllm_config, + torch.device("cpu"), + ) + + assert builder.dcp_world_size == 1 + assert builder.dcp_rank == 0 + + def test_dflash_swa_layers_use_causal_metadata(): proposer = object.__new__(DFlashProposer) proposer.model = SimpleNamespace(sliding_attention_layer_names={"layer.sw"}) diff --git a/tests/v1/spec_decode/test_dspark_cudagraph_contract.py b/tests/v1/spec_decode/test_dspark_cudagraph_contract.py new file mode 100644 index 000000000000..770407440bb9 --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_cudagraph_contract.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + +def test_dspark_generate_draft_accepts_dflash_capture_contract(): + head_hidden = torch.randn(10, 8) + speculator = SimpleNamespace( + num_query_per_req=5, + capacity_activation_batch_size=1, + _markov_outside_cudagraph=False, + _speculative_steps_for_query_len=Mock(return_value=5), + _run_model=Mock(return_value=head_hidden), + _sample_sequential=Mock(), + ) + + DSparkSpeculator._generate_draft( + speculator, + num_reqs=2, + num_tokens_padded=10, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.FULL, + is_profile=True, + num_query_per_req=5, + capture_only=True, + ) + + speculator._sample_sequential.assert_called_once_with( + 2, + head_hidden, + 5, + 5, + is_profile=True, + use_capacity=True, + ) diff --git a/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py b/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py new file mode 100644 index 000000000000..90328357cce7 --- /dev/null +++ b/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import ( + RemoteK3DSparkSpeculator, + _anchor_positions_from_context, + _build_valid_context_plan, + _contiguous_draft_output, + _RetainedRequestPrefix, +) + + +def test_build_valid_context_plan_drops_rejected_tail_rows(): + batch = SimpleNamespace( + num_reqs=2, + num_scheduled_tokens=np.array([4, 3], dtype=np.int32), + num_computed_tokens_np=np.array([10, 20], dtype=np.int32), + ) + + indices, counts = _build_valid_context_plan(batch, [2, 0]) + + assert indices == [0, 1, 4, 5, 6] + assert counts == [2, 3] + + +def test_anchor_positions_follow_actual_valid_context_rows(): + positions = torch.tensor([24, 25, 26, 80, 81], dtype=torch.int64) + + anchors = _anchor_positions_from_context([3, 2], positions) + + assert anchors == [27, 82] + + +def test_remote_tokens_copy_supports_adaptive_depth(): + proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator) + proxy.device = torch.device("cpu") + proxy.draft_tokens = torch.full((3, 8), -1, dtype=torch.int64) + + proxy._copy_tokens_from_response( + {"tokens": [[11, 12], [21, 22]]}, + active_indices=[0, 2], + num_speculative_tokens=2, + ) + + assert proxy.draft_tokens.tolist() == [ + [11, 12, -1, -1, -1, -1, -1, -1], + [-1, -1, -1, -1, -1, -1, -1, -1], + [21, 22, -1, -1, -1, -1, -1, -1], + ] + + +def test_remote_speculator_accepts_scheduler_selected_zero_depth(): + proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator) + proxy.num_speculative_steps = 3 + proxy.draft_tokens = torch.full((4, 3), -1, dtype=torch.int64) + batch = SimpleNamespace(num_reqs=2) + empty = torch.empty(0) + + output = proxy.propose( + batch, + {}, + {}, + empty, + None, + empty, + empty, + empty, + empty, + empty, + empty, + num_speculative_tokens=0, + ) + + assert output.shape == (2, 0) + assert output.is_contiguous() + + +def test_adaptive_depth_output_is_contiguous_for_tp_broadcast(): + draft_tokens = torch.arange(24, dtype=torch.int64).view(3, 8) + + output = _contiguous_draft_output(draft_tokens, 2, 3) + + assert output.is_contiguous() + assert output.tolist() == [[0, 1, 2], [8, 9, 10]] + + +@pytest.mark.parametrize("rejected", [[5, 0], [-1, 0]]) +def test_build_valid_context_plan_rejects_invalid_counts(rejected): + batch = SimpleNamespace( + num_reqs=2, + num_scheduled_tokens=np.array([4, 3], dtype=np.int32), + num_computed_tokens_np=np.array([0, 0], dtype=np.int32), + ) + + with pytest.raises(ValueError, match="Invalid valid-context length"): + _build_valid_context_plan(batch, rejected) + + +def _make_prefix_matcher() -> RemoteK3DSparkSpeculator: + proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator) + proxy._known_requests = {"old"} + proxy._remote_block_size = 16 + proxy._remote_window_size = 32 + proxy._remote_prefix_cache_tokens = 128 + proxy._retained_prefixes = { + "old": _RetainedRequestPrefix( + token_ids=torch.arange(96, dtype=torch.int32), + committed_end=96, + context_start=0, + serial=1, + ) + } + return proxy + + +def test_remote_prefix_match_requires_exact_token_identity(): + proxy = _make_prefix_matcher() + matching = torch.arange(80, dtype=torch.int32) + + assert proxy._find_reconnect_source(matching, 80, {"new"}) == "old" + + mismatched = matching.clone() + mismatched[40] = -1 + assert proxy._find_reconnect_source(mismatched, 80, {"new"}) is None + + +def test_remote_prefix_match_rejects_range_evicted_from_projected_cache(): + proxy = _make_prefix_matcher() + proxy._remote_prefix_cache_tokens = 48 + matching = torch.arange(40, dtype=torch.int32) + + assert proxy._find_reconnect_source(matching, 40, {"new"}) is None + + +def test_remote_prefix_match_rejects_history_before_cold_bootstrap(): + proxy = _make_prefix_matcher() + proxy._retained_prefixes["old"].context_start = 64 + + assert ( + proxy._find_reconnect_source(torch.arange(80, dtype=torch.int32), 80, {"new"}) + is None + ) + assert ( + proxy._find_reconnect_source(torch.arange(96, dtype=torch.int32), 96, {"new"}) + == "old" + ) diff --git a/tests/v1/spec_decode/test_k3_dspark_standalone.py b/tests/v1/spec_decode/test_k3_dspark_standalone.py new file mode 100644 index 000000000000..39e6f6c7d95c --- /dev/null +++ b/tests/v1/spec_decode/test_k3_dspark_standalone.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest +import torch + +from vllm.entrypoints.k3_dspark_rpc import ( + DraftKVSlotAllocator, + ProjectedContextCache, +) +from vllm.entrypoints.k3_dspark_standalone import ( + EMBED_TENSOR, + LM_HEAD_TENSOR, + resolve_shared_weight_files, +) + + +def _write_index(root, weight_map): + (root / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}) + ) + + +def test_resolve_shared_weight_files_requires_both_target_tensors(tmp_path): + shard = tmp_path / "shared.safetensors" + shard.touch() + _write_index(tmp_path, {EMBED_TENSOR: shard.name}) + + with pytest.raises(KeyError, match=LM_HEAD_TENSOR): + resolve_shared_weight_files(tmp_path) + + +def test_resolve_shared_weight_files_resolves_checkpoint_shards(tmp_path): + shard = tmp_path / "shared.safetensors" + shard.touch() + _write_index( + tmp_path, + { + EMBED_TENSOR: shard.name, + LM_HEAD_TENSOR: shard.name, + }, + ) + + resolved = resolve_shared_weight_files(tmp_path) + + assert resolved == { + EMBED_TENSOR: shard.resolve(), + LM_HEAD_TENSOR: shard.resolve(), + } + + +def test_resolve_shared_weight_files_rejects_path_escape(tmp_path): + outside = tmp_path.parent / "outside.safetensors" + outside.touch() + _write_index( + tmp_path, + { + EMBED_TENSOR: f"../{outside.name}", + LM_HEAD_TENSOR: f"../{outside.name}", + }, + ) + + with pytest.raises(ValueError, match="escapes"): + resolve_shared_weight_files(tmp_path) + + +def test_draft_kv_slot_allocator_keeps_rolling_blocks_unique(): + allocator = DraftKVSlotAllocator( + num_cache_blocks=11, + block_size=4, + window_size=16, + max_requests=2, + ) + first, created = allocator.get_or_allocate("first") + second, _ = allocator.get_or_allocate("second") + + assert created + assert allocator.physical_block_range(first) == slice(1, 6) + assert allocator.physical_block_range(second) == slice(6, 11) + blocks, local_len = allocator.block_table(first, 21) + assert blocks == [2, 3, 4, 5, 1] + assert len(blocks) == len(set(blocks)) + assert local_len == 17 + + +def test_draft_kv_slot_allocator_reuses_freed_request_slot(): + allocator = DraftKVSlotAllocator( + num_cache_blocks=6, + block_size=4, + window_size=16, + max_requests=1, + ) + original, _ = allocator.get_or_allocate("original") + with pytest.raises(RuntimeError, match="capacity exhausted"): + allocator.get_or_allocate("other") + + assert allocator.free("original") is original + replacement, created = allocator.get_or_allocate("replacement") + assert created + assert replacement.slot == original.slot + + +def test_draft_kv_slot_allocator_rebinds_without_changing_slot(): + allocator = DraftKVSlotAllocator( + num_cache_blocks=6, + block_size=4, + window_size=16, + max_requests=1, + ) + original, _ = allocator.get_or_allocate("original") + + rebound = allocator.rebind("original", "replacement") + + assert rebound is original + assert rebound.request_id == "replacement" + assert allocator.get("original") is None + assert allocator.get("replacement") is rebound + + +def test_projected_context_cache_rewinds_and_overwrites_exact_prefix(): + cache = ProjectedContextCache(hidden_size=2, max_tokens=8, chunk_size=4) + initial = torch.arange(16, dtype=torch.bfloat16).view(8, 2) + cache.append(0, initial) + + replacement = torch.tensor([[100, 101], [102, 103]], dtype=torch.bfloat16) + cache.append(5, replacement) + + assert cache.start_position == 0 + assert cache.end_position == 7 + assert torch.equal(cache.read(0, 5), initial[:5]) + assert torch.equal(cache.read(5, 7), replacement) + + +def test_projected_context_cache_evicts_old_rows_without_regaining_them(): + cache = ProjectedContextCache(hidden_size=1, max_tokens=6, chunk_size=4) + cache.append(0, torch.arange(8, dtype=torch.bfloat16).view(8, 1)) + + assert cache.start_position == 2 + assert cache.has_range(2, 8) + assert not cache.has_range(1, 8) + + cache.append(5, torch.tensor([[50], [60]], dtype=torch.bfloat16)) + assert cache.start_position == 2 + assert cache.end_position == 7 + with pytest.raises(ValueError, match="unavailable"): + cache.read(1, 7) + + +def test_projected_context_cache_tracks_configured_device(): + cache = ProjectedContextCache( + hidden_size=2, + max_tokens=4, + chunk_size=2, + device=torch.device("cpu"), + ) + states = torch.arange(8, dtype=torch.bfloat16).view(4, 2) + + cache.append(0, states) + + assert cache.device == torch.device("cpu") + assert cache.read(0, 4).device == cache.device + assert cache.allocated_bytes == states.numel() * states.element_size() + + +def test_projected_context_cache_rejects_device_mismatch(): + cache = ProjectedContextCache(hidden_size=2, max_tokens=4) + states = torch.empty((1, 2), dtype=torch.bfloat16, device="meta") + + with pytest.raises(ValueError, match="device mismatch"): + cache.append(0, states) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 619f3ad6fded..4eca8198222a 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -2,12 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """grammar_bitmask under spec-decode draft padding (#44006).""" +from collections.abc import Iterable, Sequence +from typing import overload +from unittest.mock import Mock + import pytest from transformers import AutoTokenizer from vllm.config import StructuredOutputsConfig, VllmConfig from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig +from vllm.reasoning.step3p5_reasoning_parser import Step3p5ReasoningParser from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -262,6 +267,54 @@ def test_validate_tokens_then_bitmask_round_trip(backend): assert not grammar.is_terminated() +def test_xgrammar_accept_tokens_stops_at_termination(capfd): + """Tokens after a terminating EOS do not reach the matcher.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + processed_before = grammar.num_processed_tokens + + assert grammar.accept_tokens(request.request_id, [eos, trailing]) + assert grammar.is_terminated() + assert grammar.num_processed_tokens == processed_before + 1 + assert "trying to accept new token" not in capfd.readouterr().err + + processed_after_eos = grammar.num_processed_tokens + assert grammar.accept_tokens(request.request_id, [trailing]) + assert grammar.num_processed_tokens == processed_after_eos + assert "trying to accept new token" not in capfd.readouterr().err + + grammar.reset() + assert not grammar.is_terminated() + assert grammar.num_processed_tokens == 0 + + +def test_xgrammar_validate_tokens_stops_at_termination(capfd): + """Validation rolls back after reaching a terminating EOS.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + + assert grammar.validate_tokens([eos, trailing]) == [eos] + assert "trying to accept new token" not in capfd.readouterr().err + # Check matcher state directly to verify validation rolled it back. + assert not grammar.matcher.is_terminated() + + assert grammar.accept_tokens(request.request_id, [eos]) + assert grammar.is_terminated() + + assert grammar.validate_tokens([trailing]) == [] + assert "trying to accept new token" not in capfd.readouterr().err + + class _MarkerReasoner: """Stub reasoner whose reasoning-end marker is a single fixed token.""" @@ -341,3 +394,154 @@ def test_trim_reasoning_for_advance(): next_step = [post, post] request.append_output_token_ids(next_step) assert manager.trim_reasoning_for_advance(request, next_step) == next_step + + +@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) +def test_speculative_grammar_filter_rejects_invalid_boundary_suffix(backend): + """Only the grammar-valid answer prefix may cross the commit boundary.""" + tokenizer, manager, request, _, marker = _setup_boundary_request(backend) + reasoning_token = tokenizer.encode(" ")[0] + valid_answer_token = tokenizer.encode("{")[0] + invalid_answer_token = tokenizer.encode("z")[0] + sampled_tokens = [ + reasoning_token, + marker, + valid_answer_token, + invalid_answer_token, + ] + + filtered, rejected = manager.filter_speculative_grammar_tokens( + request, sampled_tokens + ) + + assert filtered == [reasoning_token, marker, valid_answer_token] + assert rejected == 1 + grammar = request.structured_output_request.grammar + assert grammar.validate_tokens([valid_answer_token]) == [valid_answer_token] + + +@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) +def test_speculative_grammar_filter_rejects_tokens_after_completion(backend): + """A sampled block cannot commit tokens past a completed grammar value.""" + tokenizer, manager, request, _, _ = _setup_boundary_request(backend) + request.structured_output_request.reasoning_ended = True + complete_object = tokenizer.encode("{}") + invalid_suffix = tokenizer.encode("z")[0] + + filtered, rejected = manager.filter_speculative_grammar_tokens( + request, [*complete_object, invalid_suffix] + ) + + assert filtered == complete_object + assert rejected == 1 + + +def test_reasoning_boundary_scan_does_not_copy_committed_history(): + """Boundary detection can expose long history without materializing it.""" + parser_calls = 0 + + class TokenHistory(Sequence[int]): + def __len__(self) -> int: + return 300_000 + + @overload + def __getitem__(self, index: int) -> int: ... + + @overload + def __getitem__(self, index: slice) -> list[int]: ... + + def __getitem__(self, index: int | slice) -> int | list[int]: + raise AssertionError("committed token history was materialized") + + class EndTokenReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + nonlocal parser_calls + parser_calls += 1 + assert len(input_ids) >= 300_000 + return 99 in delta_ids + + reasoner = EndTokenReasoner() + boundary = StructuredOutputManager._find_reasoning_end_offset( + reasoner, TokenHistory(), [10, 99, 20] + ) + + assert boundary == 1 + assert parser_calls == 3 + + +def test_reasoning_boundary_scan_checks_nontransition_block_once(): + """A sampled block without a transition requires one parser call.""" + parser_calls = 0 + + class NoBoundaryReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + nonlocal parser_calls + parser_calls += 1 + return False + + reasoner = NoBoundaryReasoner() + boundary = StructuredOutputManager._find_reasoning_end_offset( + reasoner, [1, 2, 3], [10, 20, 30] + ) + + assert boundary is None + assert parser_calls == 1 + + +def test_reasoning_boundary_scan_preserves_stateful_parser(): + """Pre-commit probing must not consume Step3.5's pending transition.""" + tokenizer = Mock() + tokenizer.get_vocab.return_value = {"": 1, "": 2} + reasoner = Step3p5ReasoningParser(tokenizer) + reasoner._end_token_pending = True + + boundary = StructuredOutputManager._find_reasoning_end_offset( + reasoner, [1, 2, 3], [10, 20] + ) + + assert boundary == 0 + assert reasoner._end_token_pending + + +def test_reasoning_boundary_scan_locates_multi_token_marker(): + """A parser that examines cumulative deltas locates a multi-token marker.""" + + class MultiTokenReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + delta = list(delta_ids) + return any( + delta[index : index + 2] == [20, 30] for index in range(len(delta) - 1) + ) + + boundary = StructuredOutputManager._find_reasoning_end_offset( + MultiTokenReasoner(), [1, 2, 3], [20, 30, 40] + ) + + assert boundary == 1 + + +def test_reasoning_boundary_scan_handles_marker_across_blocks(): + """A multi-token marker may start in history and finish in the new block.""" + + class CrossBlockReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + tokens = list(input_ids) + delta_len = len(list(delta_ids)) + return any( + tokens[index : index + 3] == [7, 8, 9] + for index in range(max(0, len(tokens) - delta_len - 2), len(tokens)) + ) + + boundary = StructuredOutputManager._find_reasoning_end_offset( + CrossBlockReasoner(), [1, 7, 8], [9, 10] + ) + + assert boundary == 0 diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index ad5f1d5d7951..f5deffe3e248 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -15,7 +15,7 @@ class MockReasoner: def __init__(self, tokenizer): - self.is_reasoning_end = Mock(return_value=False) + self.is_reasoning_end_for_prompt = Mock(return_value=False) self.is_reasoning_end_streaming = Mock(return_value=False) @@ -137,7 +137,7 @@ class KwargReasoner: def __init__(self, tokenizer, chat_template_kwargs=None): self.chat_template_kwargs = chat_template_kwargs or {} - def is_reasoning_end(self, input_ids): + def is_reasoning_end_for_prompt(self, input_ids): return not self.chat_template_kwargs.get("enable_thinking", False) manager = StructuredOutputManager(mock_vllm_config) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index c026ab0e4e78..42e97e48f1d9 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -1,8 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import numpy as np import pytest +import torch +from vllm.v1.core.sched.output import GrammarOutput +from vllm.v1.structured_output import utils from vllm.v1.structured_output.backend_xgrammar import ( has_xgrammar_unsupported_json_features, ) @@ -104,3 +110,53 @@ def test_supported_json_features(supported_schema): assert not has_xgrammar_unsupported_json_features(supported_schema), ( "Schema should be supported" ) + + +def test_apply_grammar_bitmask_preserves_source_offsets_after_draft_trimming( + monkeypatch, +): + """A trimmed request must not shift another request's grammar rows. + + The scheduler serializes masks at its scheduled speculative width. A worker + may trim grammar-invalid drafts before applying those masks, so source and + destination offsets must be advanced with their respective widths. + """ + scheduler_output = SimpleNamespace( + scheduled_spec_decode_tokens={ + "trimmed-request": [1], + "full-request": [2, 3, 4], + } + ) + grammar_output = GrammarOutput( + structured_output_request_ids=["trimmed-request", "full-request"], + grammar_bitmask=np.array( + [[10], [11], [12], [13], [20], [21], [22], [23]], + dtype=np.int32, + ), + num_spec_tokens=[3, 3], + ) + input_batch = SimpleNamespace(req_ids=["trimmed-request", "full-request"]) + logits = torch.zeros((6, 32)) + applied_bitmask = None + + def capture_bitmask(logits, bitmask, indices): + nonlocal applied_bitmask + applied_bitmask = bitmask.clone() + assert indices is None + + monkeypatch.setattr( + utils, + "xgr", + SimpleNamespace(apply_token_bitmask_inplace=capture_bitmask), + ) + monkeypatch.setattr(utils, "PIN_MEMORY", False) + + utils.apply_grammar_bitmask( + scheduler_output, + grammar_output, + input_batch, + logits, + ) + + assert applied_bitmask is not None + assert applied_bitmask[:, 0].tolist() == [10, 13, 20, 21, 22, 23] diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index 186d49a038c2..ddcf0c171dca 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -94,3 +94,34 @@ def test_check_attention_cp_compatibility_rejects_no_lse_return(monkeypatch): with pytest.raises(AssertionError, match="requires attention implementations"): cp_utils.check_attention_cp_compatibility(_make_config(dcp_size=2)) + + +def test_replicated_kv_group_executes_attention_as_dcp1(monkeypatch): + """A complete per-rank KV copy must not enter DCP attention collectives.""" + impl = SimpleNamespace( + can_return_lse_for_decode=True, + dcp_world_size=16, + dcp_rank=7, + total_cp_world_size=16, + total_cp_rank=7, + need_to_return_lse_for_decode=True, + supports_pcp=False, + ) + layer = SimpleNamespace( + impl=impl, + get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + ) + + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda vllm_config, layer_type: {"draft.layer": layer}, + ) + + cp_utils.check_attention_cp_compatibility(_make_config(dcp_size=16)) + + assert impl.dcp_world_size == 1 + assert impl.dcp_rank == 0 + assert impl.total_cp_world_size == 1 + assert impl.total_cp_rank == 0 + assert impl.need_to_return_lse_for_decode is False diff --git a/tests/v1/worker/test_gpu_structured_outputs.py b/tests/v1/worker/test_gpu_structured_outputs.py new file mode 100644 index 000000000000..08867468ed0e --- /dev/null +++ b/tests/v1/worker/test_gpu_structured_outputs.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np + +from vllm.v1.worker.gpu.structured_outputs import _build_grammar_row_mapping + + +def test_grammar_mapping_preserves_bonus_rows_after_zero_draft_budget(): + """Zero draft capacity retains each request's scheduled bonus mask.""" + source_indices, logits_indices = _build_grammar_row_mapping( + req_ids=["low", "high", "prefill"], + grammar_req_ids=["low", "high", "prefill"], + grammar_num_spec_tokens=[2, 2, 0], + cu_num_logits_np=np.array([0, 1, 2, 3], dtype=np.int32), + num_draft_tokens_per_req=np.array([0, 0, 0], dtype=np.int32), + num_bonus_tokens=1, + ) + + assert source_indices == [2, 5, 6] + assert logits_indices == [0, 1, 2] + + +def test_grammar_mapping_selects_active_drafts_from_each_source_group(): + """Compaction preserves per-request draft rows and the final bonus row.""" + source_indices, logits_indices = _build_grammar_row_mapping( + req_ids=["plain", "trimmed", "full"], + grammar_req_ids=["trimmed", "full"], + grammar_num_spec_tokens=[3, 3], + cu_num_logits_np=np.array([0, 1, 3, 7], dtype=np.int32), + num_draft_tokens_per_req=np.array([0, 1, 3], dtype=np.int32), + num_bonus_tokens=1, + ) + + assert source_indices == [0, 3, 4, 5, 6, 7] + assert logits_indices == [1, 2, 3, 4, 5, 6] + + +def test_grammar_mapping_supports_non_speculative_batches(): + """A batch without draft tokens maps one bonus row per grammar request.""" + source_indices, logits_indices = _build_grammar_row_mapping( + req_ids=["plain", "grammar"], + grammar_req_ids=["grammar"], + grammar_num_spec_tokens=[0], + cu_num_logits_np=np.array([0, 1, 2], dtype=np.int32), + num_draft_tokens_per_req=None, + num_bonus_tokens=1, + ) + + assert source_indices == [0] + assert logits_indices == [1] diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 545d23f90912..ad6b2e712c04 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from vllm.platforms import current_platform +from vllm.v1.core.sched.output import NewRequestData from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState @@ -27,3 +30,35 @@ def test_postprocess_state_scalar_with_int32_mapping( [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda" ) torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) + + +@pytest.mark.parametrize( + ("num_computed_tokens", "expected_state_index"), + [(0, -1), (110_592, 8), (110_593, 9)], +) +def test_prefix_hit_uses_mamba_checkpoint_cadence( + num_computed_tokens: int, expected_state_index: int +) -> None: + """A resumed request indexes recurrent checkpoints, not attention pages.""" + state = object.__new__(MambaHybridModelState) + state.rope_state = None + state._align_mode = True + state.cache_config = SimpleNamespace(block_size=768, mamba_block_size=12_288) + state._mamba_block_size = 12_288 + state._mamba_state_idx_gpu = torch.zeros(1, dtype=torch.int32) + state.num_accepted_tokens_gpu = torch.full((1,), 9, dtype=torch.int32) + request = NewRequestData( + req_id="prefix-hit", + prompt_token_ids=[], + mm_features=[], + sampling_params=None, + pooling_params=None, + block_ids=(), + num_computed_tokens=num_computed_tokens, + lora_request=None, + ) + + state.add_request(0, request) + + assert state._mamba_state_idx_gpu.item() == expected_state_index + assert state.num_accepted_tokens_gpu.item() == 1 diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a338b934b54f..2ba1ce1c5937 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -18,6 +18,7 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, + batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, preprocess_mamba, @@ -536,6 +537,34 @@ def device(self): def test_config(self): return _TestConfig() + def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device): + batch = 128 + row_bytes = 32 * 1024 + shift = 16 + copy_size = row_bytes - shift + + pattern = (torch.arange(row_bytes, dtype=torch.int32, device=device) % 251).to( + torch.uint8 + ) + state = pattern.expand(batch, -1).clone() + snapshot = state.clone() + + row_stride_bytes = state.stride(0) * state.element_size() + row_offsets = ( + torch.arange(batch, dtype=torch.int64, device=device) * row_stride_bytes + ) + dst_ptrs = (row_offsets + state.data_ptr()).to(torch.uint64) + src_ptrs = (row_offsets + state.data_ptr() + shift).to(torch.uint64) + sizes = torch.full((batch,), copy_size, dtype=torch.int32, device=device) + + expected = snapshot.clone() + expected[:, :copy_size].copy_(snapshot[:, shift:]) + for _ in range(10): + state.copy_(snapshot) + batch_memcpy(src_ptrs, dst_ptrs, sizes) + torch.accelerator.synchronize() + torch.testing.assert_close(state, expected, rtol=0, atol=0) + def test_matches_python_postprocess_mamba(self, device, test_config): """ Golden test: GPU kernel produces identical results to Python impl. @@ -1190,12 +1219,27 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( # --- Verify Python behavior (ground truth) --- dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1 - # Conv state should be modified (shifted copy within block) - conv_changed = not torch.allclose( - conv_state_py[dest_block_id], conv_state_orig[dest_block_id] + # This is an overlapping in-place left shift, so comparing only the + # Python and fused paths can hide the same memcpy race in both. Build + # the memmove result from the untouched snapshot and check each path + # independently. + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-1].copy_( + conv_state_orig[dest_block_id, 1:] ) - assert conv_changed, ( - "Python: Conv state should be modified when accept_token_bias > 0" + torch.testing.assert_close( + conv_state_py, + expected_conv_state, + rtol=0, + atol=0, + msg="Python: overlapping conv copy should have memmove semantics", + ) + torch.testing.assert_close( + conv_state_gpu, + expected_conv_state, + rtol=0, + atol=0, + msg="GPU: overlapping conv copy should have memmove semantics", ) # Temporal state should be modified (copy from different block) @@ -2122,13 +2166,10 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): actual_src_block_idx = src_block_idx + accept_token_bias actual_src_block_id = block_table[req, actual_src_block_idx] - All prior regression tests exercise only ``bias == 1``, i.e. they - only ever read one slot ahead of ``src_block_idx`` in the block - table. An off-by-one (or missing scale) in the address computation - on line 143 of ``mamba_utils.py`` would be invisible to every - existing test but would silently read the wrong physical block on - any speculative-decode cycle that accepts multiple tokens across a - block boundary, feeding a stale hidden state forward one step. + A ``bias == 1`` case only reads one slot ahead of ``src_block_idx`` + in the block table. This test isolates the larger-stride case, where + an off-by-one would read the wrong physical block after multiple + tokens are accepted across a block boundary. Setup (block_size=16): - running = 28 + 2 - 0 = 30 @@ -2141,8 +2182,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): With identity block_ids = [0,1,2,3,...], an off-by-one that used bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3, - producing a clear state-value mismatch against the Python - reference. + producing a clear mismatch against the untouched snapshot. """ cfg = test_config torch.manual_seed(7002) @@ -2166,6 +2206,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): fwd_py, fwd_gpu, ) = _make_dual_layer_state(cfg, device) + conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Python reference --- @@ -2212,12 +2253,22 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): device=device, ) - # --- Ground truth: Python must have sourced temporal from block 3 --- + # --- Ground truth from untouched snapshots --- actual_src_block_id = block_ids_per_req[0][3] # == 3 dest_block_id = block_ids_per_req[0][1] # == 1 + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-2].copy_( + conv_state_orig[dest_block_id, 2:] + ) + torch.testing.assert_close(conv_state_py, expected_conv_state, rtol=0, atol=0) + torch.testing.assert_close(conv_state_gpu, expected_conv_state, rtol=0, atol=0) + + # Python must have sourced temporal from block 3. torch.testing.assert_close( temporal_state_py[dest_block_id], temporal_state_orig[actual_src_block_id], + rtol=0, + atol=0, msg=( "Python reference did not copy from block_ids[src+bias]=3; " "test preconditions are wrong" @@ -2251,21 +2302,44 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): msg="num_accepted_tokens mismatch at accept_token_bias=2", ) - def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( - self, device, test_config, monkeypatch + @pytest.mark.parametrize( + "same_physical_block", [True, False], ids=["same", "distinct"] + ) + @pytest.mark.parametrize("accept_token_bias", [1, 2, 3]) + @pytest.mark.parametrize( + "dtype", + [torch.float16, torch.float32, torch.float64], + ids=["fp16", "fp32", "fp64"], + ) + def test_sd_and_ds_conv_layouts_match_snapshot( + self, + device, + test_config, + monkeypatch, + accept_token_bias, + same_physical_block, + dtype, ): - """DS conv postprocess should match SD when accept_token_bias > 0.""" + """SD and DS copies should independently match memmove semantics.""" from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils cfg = test_config + cfg.dtype = dtype torch.manual_seed(38898) req_ids = ["req_0"] - num_computed_tokens = [30] - num_scheduled_tokens = {"req_0": 1} + # Keep new_num_computed on an aligned boundary while varying how far + # below it the running state starts. This makes the copy bias exactly + # ``accept_token_bias`` for each case. The 32 boundary keeps source and + # destination in logical block 1; the 64 boundary copies block 2 -> 3. + aligned_boundary = 32 if same_physical_block else 64 + num_computed_tokens = [aligned_boundary - 2 * accept_token_bias] + num_scheduled_tokens = {"req_0": accept_token_bias} num_draft_tokens: dict[str, int] = {} - num_accepted_tokens = [2] # Results in accept_token_bias = 1 - mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx + num_accepted_tokens = [accept_token_bias + 1] + dest_block_idx = aligned_boundary // cfg.block_size - 1 + src_block_idx = dest_block_idx if same_physical_block else dest_block_idx - 1 + mamba_state_idx = [src_block_idx] block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] @@ -2288,7 +2362,8 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device ) - # SD GPU path. Default layout is SD. + # SD GPU path. + monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False) model_mamba_utils.get_conv_state_layout.cache_clear() sd_conv = sd_source_conv.clone() sd_temporal = sd_source_temporal.clone() @@ -2312,9 +2387,32 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( ) torch.accelerator.synchronize() - # Sanity: SD path actually modified the state (copy was performed). - assert not torch.equal(sd_conv, sd_source_conv), ( - "SD baseline did not modify conv state; test setup is wrong" + src_block_id = block_ids_per_req[0][src_block_idx] + dest_block_id = block_ids_per_req[0][dest_block_idx] + expected_conv = sd_source_conv.clone() + expected_conv[dest_block_id, :-accept_token_bias].copy_( + sd_source_conv[src_block_id, accept_token_bias:] + ) + torch.testing.assert_close( + sd_conv, + expected_conv, + rtol=0, + atol=0, + msg="SD conv copy did not match the untouched source snapshot", + ) + + actual_temporal_src_idx = src_block_idx + accept_token_bias + actual_temporal_src_id = block_ids_per_req[0][actual_temporal_src_idx] + expected_temporal = sd_source_temporal.clone() + expected_temporal[dest_block_id].copy_( + sd_source_temporal[actual_temporal_src_id] + ) + torch.testing.assert_close( + sd_temporal, + expected_temporal, + rtol=0, + atol=0, + msg="SD temporal copy did not match the untouched source snapshot", ) # DS GPU path on the DS twin. @@ -2346,22 +2444,39 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( # Reset the lru cache so other tests see the default layout again. model_mamba_utils.get_conv_state_layout.cache_clear() - # DS bytes, un-permuted, should match the SD result. + # Validate DS independently against the snapshot; otherwise a shared + # SD/DS bug would remain invisible. + ds_conv_sd_layout = ds_conv.permute(0, 2, 1).contiguous() torch.testing.assert_close( - ds_conv.permute(0, 2, 1).contiguous(), - sd_conv, - msg=( - "DS conv post-kernel does not match SD baseline; the DS " - "row-loop in postprocess_mamba_fused_kernel is wrong." - ), + ds_conv_sd_layout, + expected_conv, + rtol=0, + atol=0, + msg="DS conv copy did not match the untouched source snapshot", ) torch.testing.assert_close( ds_temporal, - sd_temporal, - msg="DS temporal state diverged from SD", + expected_temporal, + rtol=0, + atol=0, + msg="DS temporal copy did not match the untouched source snapshot", + ) + + expected_accepted = 1 if same_physical_block else accept_token_bias + 1 + expected_accepted_tensor = torch.tensor( + [expected_accepted], dtype=torch.int32, device=device ) torch.testing.assert_close( - gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], gpu_ctx_sd.num_accepted_tokens_out[:num_reqs], - msg="DS num_accepted_tokens diverged from SD", + expected_accepted_tensor, + rtol=0, + atol=0, + msg="SD num_accepted_tokens result is wrong", + ) + torch.testing.assert_close( + gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], + expected_accepted_tensor, + rtol=0, + atol=0, + msg="DS num_accepted_tokens result is wrong", ) diff --git a/vllm/distributed/communication_op.py b/vllm/distributed/communication_op.py index 1bc270f3084f..33abaa45ed5d 100644 --- a/vllm/distributed/communication_op.py +++ b/vllm/distributed/communication_op.py @@ -26,6 +26,16 @@ def tensor_model_parallel_all_gather( return get_tp_group().all_gather(input_, dim) +def tensor_model_parallel_all_gatherv( + input_: torch.Tensor, sizes: list[int], dim: int = 0 +) -> torch.Tensor: + """All-gather variable-length tensor slices across the model-parallel group.""" + tp_group = get_tp_group() + if tp_group.world_size == 1: + return input_ + return tp_group.all_gatherv(input_, dim=dim, sizes=sizes) + + def tensor_model_parallel_reduce_scatter( input_: torch.Tensor, dim: int = -1 ) -> torch.Tensor: diff --git a/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py index 1d3b43082c02..53b7c0a86220 100644 --- a/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py @@ -169,11 +169,15 @@ def prepare_graph_all_reduce( ) -> None: del stream channel_id = self._active_channel_id or self._EAGER_CHANNEL_ID - if not self._workspace(channel_id).supports(inp): + workspace = self._workspace(channel_id) + if not workspace.supports(inp): raise ValueError( "FlashInfer PCIe IPC graph warmup received an unsupported " f"shape {tuple(inp.shape)}" ) + hidden = inp.shape[-1] + batch = inp.numel() // hidden + workspace.prepare([(batch, hidden)], dtype=inp.dtype) def all_reduce( self, diff --git a/vllm/entrypoints/k3_dspark_rpc.py b/vllm/entrypoints/k3_dspark_rpc.py new file mode 100644 index 000000000000..96f049205573 --- /dev/null +++ b/vllm/entrypoints/k3_dspark_rpc.py @@ -0,0 +1,1363 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-staged RPC for a dedicated Kimi-K3 draft GPU. + +The verifier and RTX 3090 do not have CUDA peer access on the target host, so +the first transport deliberately uses ZMQ multipart frames backed by host +memory. The protocol is small and versioned so a verifier-side proxy can be +added without coupling the standalone process to the generic EAGLE draft +server protocol. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +import zmq + +from vllm.config.vllm import set_current_vllm_config +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonDecodeMetadata, + MLACommonMetadata, +) +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) +from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id + +if TYPE_CHECKING: + from vllm.entrypoints.k3_dspark_standalone import StandaloneRuntime + +logger = init_logger(__name__) + +PROTOCOL_VERSION = 2 + + +class ProjectedContextCache: + """Bounded, chunked cache of projected DSpark context states. + + The standalone draft server keeps this cache on the draft device. Keeping + the projected rows on CPU forced a blocking D2H copy after every proposal, + even though the rows are normally consumed again by the same GPU during a + prefix reconnect. CPU remains the default for lightweight unit tests and + callers which explicitly want host storage. + """ + + def __init__( + self, + *, + hidden_size: int, + max_tokens: int, + chunk_size: int = 256, + initial_position: int = 0, + device: torch.device | str = "cpu", + ) -> None: + if hidden_size <= 0 or max_tokens <= 0 or chunk_size <= 0: + raise ValueError("Projected context cache dimensions must be positive") + if initial_position < 0: + raise ValueError("Projected context cache position cannot be negative") + self.hidden_size = hidden_size + self.max_tokens = max_tokens + self.chunk_size = chunk_size + self.device = torch.device(device) + self.start_position = initial_position + self.end_position = initial_position + self._chunks: dict[int, torch.Tensor] = {} + + def _truncate(self, end_position: int) -> None: + if not self.start_position <= end_position <= self.end_position: + raise ValueError( + "Cannot truncate projected context outside its retained range: " + f"retained=[{self.start_position}, {self.end_position}), " + f"requested_end={end_position}" + ) + first_discarded_chunk = (end_position + self.chunk_size - 1) // self.chunk_size + for chunk_idx in list(self._chunks): + if chunk_idx >= first_discarded_chunk: + del self._chunks[chunk_idx] + self.end_position = end_position + + def append(self, first_position: int, states: torch.Tensor) -> None: + if states.device != self.device: + raise ValueError( + "Projected context cache device mismatch: " + f"cache={self.device}, states={states.device}" + ) + if states.dtype != torch.bfloat16 or states.ndim != 2: + raise ValueError("Projected context states must be a 2D BF16 tensor") + if states.shape[1] != self.hidden_size: + raise ValueError( + f"Projected context width is {states.shape[1]}, expected " + f"{self.hidden_size}" + ) + if first_position < self.start_position or first_position > self.end_position: + raise ValueError( + "Projected context append is not contiguous with retained state: " + f"retained=[{self.start_position}, {self.end_position}), " + f"first={first_position}" + ) + if first_position < self.end_position: + self._truncate(first_position) + + offset = 0 + num_rows = int(states.shape[0]) + while offset < num_rows: + position = first_position + offset + chunk_idx, chunk_offset = divmod(position, self.chunk_size) + count = min(num_rows - offset, self.chunk_size - chunk_offset) + chunk = self._chunks.get(chunk_idx) + if chunk is None: + chunk = torch.empty( + (self.chunk_size, self.hidden_size), + dtype=torch.bfloat16, + device=self.device, + ) + self._chunks[chunk_idx] = chunk + chunk[chunk_offset : chunk_offset + count].copy_( + states[offset : offset + count] + ) + offset += count + + self.end_position = first_position + num_rows + self.start_position = max( + self.start_position, + self.end_position - self.max_tokens, + ) + for chunk_idx in list(self._chunks): + if (chunk_idx + 1) * self.chunk_size <= self.start_position: + del self._chunks[chunk_idx] + + def has_range(self, start_position: int, end_position: int) -> bool: + return ( + self.start_position <= start_position <= end_position <= self.end_position + ) + + def read(self, start_position: int, end_position: int) -> torch.Tensor: + if not self.has_range(start_position, end_position): + raise ValueError( + "Projected context range is unavailable: " + f"retained=[{self.start_position}, {self.end_position}), " + f"requested=[{start_position}, {end_position})" + ) + output = torch.empty( + (end_position - start_position, self.hidden_size), + dtype=torch.bfloat16, + device=self.device, + ) + offset = 0 + while start_position + offset < end_position: + position = start_position + offset + chunk_idx, chunk_offset = divmod(position, self.chunk_size) + count = min( + end_position - position, + self.chunk_size - chunk_offset, + ) + chunk = self._chunks.get(chunk_idx) + if chunk is None: + raise RuntimeError( + f"Projected context chunk {chunk_idx} is unexpectedly missing" + ) + output[offset : offset + count].copy_( + chunk[chunk_offset : chunk_offset + count] + ) + offset += count + return output + + def truncate(self, end_position: int) -> None: + if end_position < self.end_position: + self._truncate(end_position) + + @property + def allocated_bytes(self) -> int: + return sum( + chunk.numel() * chunk.element_size() for chunk in self._chunks.values() + ) + + +@dataclass +class DraftRequestState: + request_id: str + slot: int + committed_end: int = 0 + context_start: int = 0 + context_cache: ProjectedContextCache | None = None + + +class DraftKVSlotAllocator: + """Assign fixed rolling MLA block ranges to a small request batch.""" + + def __init__( + self, + *, + num_cache_blocks: int, + block_size: int, + window_size: int, + max_requests: int, + ) -> None: + if window_size <= 0 or window_size % block_size != 0: + raise ValueError( + "DSpark KV window must be a positive block-size multiple, got " + f"window={window_size}, block_size={block_size}" + ) + self.block_size = block_size + self.window_size = window_size + self.max_requests = max_requests + # The vLLM rolling window may retain window + block_size - 1 tokens + # while it waits for the next whole-block shift. + self.blocks_per_request = window_size // block_size + 1 + required = 1 + max_requests * self.blocks_per_request + if required > num_cache_blocks: + raise ValueError( + "Dedicated draft KV cache is too small for the requested rolling " + f"slots: required_blocks={required}, available={num_cache_blocks}" + ) + self._free_slots = list(range(max_requests)) + self._states: dict[str, DraftRequestState] = {} + + def get_or_allocate(self, request_id: str) -> tuple[DraftRequestState, bool]: + state = self._states.get(request_id) + if state is not None: + return state, False + if not self._free_slots: + raise RuntimeError( + f"DSpark request capacity exhausted (max={self.max_requests})" + ) + slot = self._free_slots.pop(0) + state = DraftRequestState(request_id=request_id, slot=slot) + self._states[request_id] = state + return state, True + + def free(self, request_id: str) -> DraftRequestState | None: + state = self._states.pop(request_id, None) + if state is not None: + self._free_slots.append(state.slot) + self._free_slots.sort() + return state + + def get(self, request_id: str) -> DraftRequestState | None: + return self._states.get(request_id) + + def rebind(self, source_request_id: str, request_id: str) -> DraftRequestState: + state = self._states.get(source_request_id) + if state is None: + raise KeyError(f"Unknown DSpark source request {source_request_id!r}") + if source_request_id == request_id: + return state + if request_id in self._states: + raise ValueError(f"DSpark request {request_id!r} already exists") + del self._states[source_request_id] + state.request_id = request_id + self._states[request_id] = state + return state + + def physical_block(self, state: DraftRequestState, position: int) -> int: + absolute_block = position // self.block_size + return ( + 1 + + state.slot * self.blocks_per_request + + absolute_block % self.blocks_per_request + ) + + def cache_slot(self, state: DraftRequestState, position: int) -> int: + return self.physical_block(state, position) * self.block_size + ( + position % self.block_size + ) + + def block_table( + self, state: DraftRequestState, sequence_end: int + ) -> tuple[list[int], int]: + if sequence_end <= 0: + raise ValueError(f"sequence_end must be positive, got {sequence_end}") + first_block = ( + max(state.context_start, sequence_end - self.window_size) // self.block_size + ) + end_block = (sequence_end + self.block_size - 1) // self.block_size + blocks = [ + 1 + + state.slot * self.blocks_per_request + + absolute_block % self.blocks_per_request + for absolute_block in range(first_block, end_block) + ] + local_sequence_len = sequence_end - first_block * self.block_size + if local_sequence_len > self.window_size + self.block_size - 1: + raise AssertionError("rolling DSpark sequence length exceeded its window") + if len(blocks) > self.blocks_per_request: + raise AssertionError("rolling DSpark block table aliases a live block") + return blocks, local_sequence_len + + def physical_block_range(self, state: DraftRequestState) -> slice: + start = 1 + state.slot * self.blocks_per_request + return slice(start, start + self.blocks_per_request) + + @property + def active_requests(self) -> int: + return len(self._states) + + @property + def request_ids(self) -> list[str]: + return list(self._states) + + +@dataclass +class _DraftCudaGraphState: + """Persistent inputs and output for one standalone draft graph shape.""" + + batch_size: int + num_speculative_tokens: int + query_len: int + input_ids: torch.Tensor + positions: torch.Tensor + slots: torch.Tensor + seq_lens: torch.Tensor + block_table: torch.Tensor + output_tokens: torch.Tensor + input_ids_host: torch.Tensor + positions_host: torch.Tensor + slots_host: torch.Tensor + seq_lens_host: torch.Tensor + block_table_host: torch.Tensor + attn_metadata: dict[str, Any] + slot_mapping: dict[str, torch.Tensor] + graph: torch.cuda.CUDAGraph | None = None + captured_hidden: torch.Tensor | None = None + captured_logits: torch.Tensor | None = None + + def stage( + self, + *, + input_ids: list[int], + positions: list[int], + slots: list[int], + block_rows: list[list[int]], + seq_lens: list[int], + ) -> None: + """Copy one request batch into address-stable graph inputs.""" + expected_tokens = self.batch_size * self.query_len + if not ( + len(input_ids) == len(positions) == len(slots) == expected_tokens + and len(block_rows) == len(seq_lens) == self.batch_size + ): + raise ValueError("Draft CUDA graph input shape mismatch") + + self.input_ids_host.copy_(torch.tensor(input_ids, dtype=torch.int64)) + self.positions_host.copy_(torch.tensor(positions, dtype=torch.int64)) + self.slots_host.copy_(torch.tensor(slots, dtype=torch.int64)) + self.seq_lens_host.copy_(torch.tensor(seq_lens, dtype=torch.int32)) + self.block_table_host.zero_() + for row_idx, row in enumerate(block_rows): + if len(row) > self.block_table_host.shape[1]: + raise ValueError( + "Draft block table exceeds CUDA graph capacity: " + f"row={len(row)}, capacity={self.block_table_host.shape[1]}" + ) + self.block_table_host[row_idx, : len(row)].copy_( + torch.tensor(row, dtype=torch.int32) + ) + + # All copies and the replay are enqueued on the same stream. The + # proposal's final query event synchronizes before these pinned host + # buffers can be reused by the next (serialized) RPC. + self.input_ids.copy_(self.input_ids_host, non_blocking=True) + self.positions.copy_(self.positions_host, non_blocking=True) + self.slots.copy_(self.slots_host, non_blocking=True) + self.seq_lens.copy_(self.seq_lens_host, non_blocking=True) + self.block_table.copy_(self.block_table_host, non_blocking=True) + + +class K3DSparkDraftEngine: + """Minimal greedy K3 draft scheduler backed by the 3090 KV cache.""" + + def __init__( + self, + runtime: StandaloneRuntime, + *, + max_requests: int, + window_size: int, + device: torch.device, + ) -> None: + self.runtime = runtime + self.model = runtime.model + self.method = runtime.method + self.device = device + self.max_model_len = int(runtime.vllm_config.model_config.max_model_len) + first_cache = next(iter(runtime.kv_caches.values())) + self.allocator = DraftKVSlotAllocator( + num_cache_blocks=int(first_cache.shape[0]), + block_size=runtime.kv_cache_block_size, + window_size=window_size, + max_requests=max_requests, + ) + speculative_config = runtime.vllm_config.speculative_config + assert speculative_config is not None + draft_config = speculative_config.draft_model_config.hf_config + self.hidden_size = int(draft_config.hidden_size) + aux_layers = get_eagle3_aux_layers_from_config(speculative_config) + if not aux_layers: + raise ValueError( + f"K3 {self.method} config does not declare target auxiliary layers" + ) + self.num_aux_layers = len(aux_layers) + target_hidden_size = int( + getattr(draft_config, "target_hidden_size", None) + or draft_config.hidden_size + ) + self.raw_context_width = int(target_hidden_size * self.num_aux_layers) + self.mask_token_id = get_parallel_drafting_token_id(draft_config) + self.max_speculative_tokens = int( + runtime.vllm_config.speculative_config.num_speculative_tokens + ) + self.max_context_tokens = int( + runtime.vllm_config.scheduler_config.max_num_batched_tokens + ) + self.prefix_cache_tokens = int( + os.environ.get( + "VLLM_K3_DRAFT_PREFIX_CACHE_TOKENS", + os.environ.get("VLLM_K3_DSPARK_PREFIX_CACHE_TOKENS", "131072"), + ) + ) + if self.prefix_cache_tokens < self.allocator.window_size: + raise ValueError( + "VLLM_K3_DRAFT_PREFIX_CACHE_TOKENS must be at least the " + f"draft KV window ({self.allocator.window_size}), got " + f"{self.prefix_cache_tokens}" + ) + self._positions_staging = torch.empty( + self.max_context_tokens, + dtype=torch.int64, + pin_memory=True, + ) + self._context_staging = torch.empty( + self.max_context_tokens * self.raw_context_width, + dtype=torch.bfloat16, + pin_memory=True, + ) + self._lock = threading.Lock() + self.proposal_count = 0 + self.last_latency_ms = 0.0 + self.last_timing_ms: dict[str, float] = {} + self._timing_totals_ms: dict[str, float] = {} + self.cold_bootstrap_count = 0 + self.reconnect_count = 0 + self.last_reconnect_latency_ms = 0.0 + self.cuda_graph_enabled = False + self.cuda_graph_capture_seconds = 0.0 + self.cuda_graph_memory_gib = 0.0 + self.cuda_graph_replay_count = 0 + self.cuda_graph_eager_fallback_count = 0 + self._cuda_graphs: dict[tuple[int, int], _DraftCudaGraphState] = {} + + def _make_cuda_graph_state( + self, + batch_size: int, + num_speculative_tokens: int, + ) -> _DraftCudaGraphState: + if self.method == "dflash" and self.runtime.attn_metadata_builder is None: + raise RuntimeError("K3 DFlash attention metadata builder is missing") + + query_len = ( + num_speculative_tokens + if self.method == "dspark" + else 1 + num_speculative_tokens + ) + num_tokens = batch_size * query_len + max_blocks = self.allocator.blocks_per_request + input_ids = torch.empty(num_tokens, dtype=torch.int64, device=self.device) + positions = torch.empty(num_tokens, dtype=torch.int64, device=self.device) + slots = torch.empty(num_tokens, dtype=torch.int64, device=self.device) + seq_lens = torch.empty(batch_size, dtype=torch.int32, device=self.device) + block_table = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + output_tokens = torch.empty( + (batch_size, num_speculative_tokens), + dtype=torch.int64, + device=self.device, + ) + + input_ids_host = torch.empty(num_tokens, dtype=torch.int64, pin_memory=True) + positions_host = torch.empty(num_tokens, dtype=torch.int64, pin_memory=True) + slots_host = torch.empty(num_tokens, dtype=torch.int64, pin_memory=True) + seq_lens_host = torch.empty(batch_size, dtype=torch.int32, pin_memory=True) + block_table_host = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, pin_memory=True + ) + + # The graph's launch topology is shape-static, while seq_lens remains + # a live tensor. Triton BF16 attention uses seq_lens to bound the KV + # scan; this conservative upper bound therefore does not force every + # replay to scan the full rolling window. + max_seq_len = self.allocator.window_size + self.allocator.block_size - 1 + query_start_cpu = torch.arange( + 0, + (batch_size + 1) * query_len, + query_len, + dtype=torch.int32, + ) + query_start_gpu = query_start_cpu.to(self.device) + if self.method == "dflash": + common = CommonAttentionMetadata( + query_start_loc=query_start_gpu, + query_start_loc_cpu=query_start_cpu, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=torch.full( + (batch_size,), max_seq_len, dtype=torch.int32 + ), + max_seq_len=max_seq_len, + num_reqs=batch_size, + num_actual_tokens=num_tokens, + max_query_len=query_len, + block_table_tensor=block_table, + slot_mapping=slots, + causal=True, + ) + assert self.runtime.attn_metadata_builder is not None + metadata = self.runtime.attn_metadata_builder.build(0, common) + else: + metadata = MLACommonMetadata( + num_reqs=batch_size, + max_query_len=query_len, + max_seq_len=max_seq_len, + num_actual_tokens=num_tokens, + query_start_loc=query_start_gpu, + slot_mapping=slots, + num_decodes=batch_size, + num_decode_tokens=num_tokens, + num_prefills=0, + causal=False, + head_dim=int(next(iter(self.runtime.kv_caches.values())).shape[-1]), + prefill=None, + decode=MLACommonDecodeMetadata( + block_table=block_table, + seq_lens=seq_lens, + dcp_tot_seq_lens=None, + ), + ) + attn_metadata = {layer_name: metadata for layer_name in self.runtime.kv_caches} + slot_mapping = {layer_name: slots for layer_name in self.runtime.kv_caches} + state = _DraftCudaGraphState( + batch_size=batch_size, + num_speculative_tokens=num_speculative_tokens, + query_len=query_len, + input_ids=input_ids, + positions=positions, + slots=slots, + seq_lens=seq_lens, + block_table=block_table, + output_tokens=output_tokens, + input_ids_host=input_ids_host, + positions_host=positions_host, + slots_host=slots_host, + seq_lens_host=seq_lens_host, + block_table_host=block_table_host, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping, + ) + + # Seed capture with valid, isolated dummy sequences. Runtime replay + # overwrites every graph input before use. + dummy_input_ids: list[int] = [] + dummy_positions: list[int] = [] + dummy_slots: list[int] = [] + dummy_rows: list[list[int]] = [] + for request_idx in range(batch_size): + dummy_input_ids.append(0) + dummy_input_ids.extend([self.mask_token_id] * (query_len - 1)) + dummy_positions.extend(range(query_len)) + dummy_slots.extend( + request_idx * self.allocator.block_size + position + for position in range(query_len) + ) + dummy_rows.append([request_idx]) + state.stage( + input_ids=dummy_input_ids, + positions=dummy_positions, + slots=dummy_slots, + block_rows=dummy_rows, + seq_lens=[query_len] * batch_size, + ) + return state + + def _run_cuda_graph_state( + self, + state: _DraftCudaGraphState, + ) -> None: + num_tokens = state.batch_size * state.query_len + with ( + set_current_vllm_config(self.runtime.draft_vllm_config), + set_forward_context( + state.attn_metadata, + self.runtime.draft_vllm_config, + num_tokens=num_tokens, + skip_compiled=True, + slot_mapping=state.slot_mapping, + ), + ): + hidden = self.model( + input_ids=state.input_ids, + positions=state.positions, + ) + if self.method == "dflash": + sample_hidden = hidden.view(state.batch_size, state.query_len, -1)[:, 1:] + logits = self.model.compute_logits( + sample_hidden.reshape( + state.batch_size * state.num_speculative_tokens, -1 + ) + ) + state.output_tokens.copy_( + logits.argmax(dim=-1).view( + state.batch_size, state.num_speculative_tokens + ) + ) + else: + logits = self.model.compute_draft_logits(hidden).view( + state.batch_size, state.query_len, -1 + ) + previous = state.input_ids.view(state.batch_size, state.query_len)[:, 0] + for step in range(state.query_len): + markov = self.model.markov_bias(self.model.markov_embed(previous)) + previous = (logits[:, step] + markov).argmax(dim=-1) + state.output_tokens[:, step].copy_(previous) + # Retain graph-owned outputs so their backing allocations cannot be + # recycled while captured nodes still reference them. + state.captured_hidden = hidden + state.captured_logits = logits + + @torch.inference_mode() + def capture_cuda_graphs(self, *, warmups: int = 2) -> None: + """Capture all DSpark or DFlash shapes selectable by this server.""" + if warmups < 1: + raise ValueError("Draft CUDA graph capture requires at least one warmup") + if self._cuda_graphs: + return + + started = time.perf_counter() + allocated_before = torch.cuda.memory_allocated(self.device) + capture_stream = torch.cuda.Stream(device=self.device) + capture_stream.wait_stream(torch.cuda.current_stream(self.device)) + # Capture the largest shape first. Triton MLA grows shared workspace + # buffers on demand; capturing a smaller shape first and then resizing + # that workspace for B2/K3 leaves the earlier graph with stale device + # pointers and causes an illegal access on replay. + shapes = [ + (batch_size, depth) + for batch_size in range( + self.runtime.vllm_config.scheduler_config.max_num_seqs, + 0, + -1, + ) + for depth in range(self.max_speculative_tokens, 0, -1) + ] + logger.info( + "Capturing %d standalone K3 %s CUDA graphs on %s: %s", + len(shapes), + self.method, + self.device, + ", ".join(f"B{batch_size}K{depth}" for batch_size, depth in shapes), + ) + with torch.cuda.stream(capture_stream): + for batch_size, depth in shapes: + state = self._make_cuda_graph_state(batch_size, depth) + for _ in range(warmups): + self._run_cuda_graph_state(state) + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph( + graph, + stream=capture_stream, + ): + self._run_cuda_graph_state(state) + state.graph = graph + self._cuda_graphs[(batch_size, depth)] = state + + torch.cuda.current_stream(self.device).wait_stream(capture_stream) + torch.cuda.synchronize(self.device) + # Dummy capture rows only touch these reserved low blocks. Block zero + # is never assigned; live request allocation clears its own full range. + max_dummy_blocks = int(self.runtime.vllm_config.scheduler_config.max_num_seqs) + for cache in self.runtime.kv_caches.values(): + cache[:max_dummy_blocks].zero_() + torch.cuda.synchronize(self.device) + + self.cuda_graph_enabled = True + self.cuda_graph_capture_seconds = time.perf_counter() - started + self.cuda_graph_memory_gib = max( + 0.0, + (torch.cuda.memory_allocated(self.device) - allocated_before) / 1024**3, + ) + logger.info( + "Standalone K3 %s CUDA graphs ready in %.2fs; allocated_delta=%.3f GiB", + self.method, + self.cuda_graph_capture_seconds, + self.cuda_graph_memory_gib, + ) + + @property + def cuda_graph_shapes(self) -> list[str]: + return [ + f"B{batch_size}K{depth}" for batch_size, depth in sorted(self._cuda_graphs) + ] + + def _clear_state_cache( + self, + state: DraftRequestState, + *, + clear_context: bool = True, + ) -> None: + block_range = self.allocator.physical_block_range(state) + for cache in self.runtime.kv_caches.values(): + cache[block_range].zero_() + state.committed_end = 0 + state.context_start = 0 + if clear_context: + state.context_cache = None + + def reset(self, request_ids: list[str]) -> None: + with self._lock: + for request_id in request_ids: + state, _ = self.allocator.get_or_allocate(request_id) + self._clear_state_cache(state) + + def free(self, request_ids: list[str]) -> None: + with self._lock: + for request_id in request_ids: + self.allocator.free(request_id) + + def clear(self) -> None: + self.free(self.allocator.request_ids) + + @property + def prefix_cache_bytes(self) -> int: + return sum( + state.context_cache.allocated_bytes + for request_id in self.allocator.request_ids + if (state := self.allocator.get(request_id)) is not None + and state.context_cache is not None + ) + + @property + def prefix_cache_host_bytes(self) -> int: + return sum( + state.context_cache.allocated_bytes + for request_id in self.allocator.request_ids + if (state := self.allocator.get(request_id)) is not None + and state.context_cache is not None + and state.context_cache.device.type == "cpu" + ) + + @property + def prefix_cache_device_bytes(self) -> int: + return self.prefix_cache_bytes - self.prefix_cache_host_bytes + + @property + def mean_timing_ms(self) -> dict[str, float]: + if self.proposal_count <= 0: + return {} + return { + key: value / self.proposal_count + for key, value in self._timing_totals_ms.items() + } + + def _record_timing(self, timing_ms: dict[str, float]) -> None: + self.last_timing_ms = timing_ms + for key, value in timing_ms.items(): + self._timing_totals_ms[key] = self._timing_totals_ms.get(key, 0.0) + value + + def _restore_projected_context( + self, + state: DraftRequestState, + prefix_end: int, + ) -> int: + context_cache = state.context_cache + if context_cache is None: + raise ValueError( + f"No projected context is retained for {state.request_id!r}" + ) + restore_start = max(0, prefix_end - self.allocator.window_size) + restore_start = ( + restore_start // self.allocator.block_size * self.allocator.block_size + ) + if not context_cache.has_range(restore_start, prefix_end): + raise ValueError( + f"Projected context for {state.request_id!r} cannot restore " + f"prefix_end={prefix_end}; retained=" + f"[{context_cache.start_position}, {context_cache.end_position})" + ) + + self._clear_state_cache(state, clear_context=False) + for start in range(restore_start, prefix_end, self.max_context_tokens): + end = min(prefix_end, start + self.max_context_tokens) + context_states = context_cache.read(start, end) + positions = torch.arange(start, end, dtype=torch.int64) + context_gpu = context_states.to(self.device, non_blocking=True) + positions_gpu = positions.to(self.device, non_blocking=True) + slots = torch.tensor( + [ + self.allocator.cache_slot(state, position) + for position in range(start, end) + ], + dtype=torch.int64, + device=self.device, + ) + self.model.precompute_and_store_context_kv( + context_gpu, + positions_gpu, + slots, + ) + state.committed_end = prefix_end + state.context_start = restore_start + context_cache.truncate(prefix_end) + return restore_start + + def reconnect( + self, + source_request_id: str, + request_id: str, + prefix_end: int, + ) -> dict[str, Any]: + started = time.perf_counter() + with self._lock: + state = self.allocator.get(source_request_id) + if state is None: + raise KeyError(f"Unknown DSpark source request {source_request_id!r}") + context_cache = state.context_cache + restore_start = max(0, prefix_end - self.allocator.window_size) + restore_start = ( + restore_start // self.allocator.block_size * self.allocator.block_size + ) + if ( + prefix_end <= 0 + or context_cache is None + or not context_cache.has_range(restore_start, prefix_end) + ): + retained = ( + None + if context_cache is None + else [context_cache.start_position, context_cache.end_position] + ) + raise ValueError( + f"Cannot reconnect {source_request_id!r} at {prefix_end}; " + f"retained={retained}" + ) + state = self.allocator.rebind(source_request_id, request_id) + restored_start = self._restore_projected_context(state, prefix_end) + torch.accelerator.synchronize() + self.reconnect_count += 1 + self.last_reconnect_latency_ms = (time.perf_counter() - started) * 1000 + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "request_id": request_id, + "restored_start": restored_start, + "prefix_end": prefix_end, + "latency_ms": self.last_reconnect_latency_ms, + "active_requests": self.allocator.active_requests, + } + + def _decode_host_tensor( + self, + frame: bytes, + *, + dtype: torch.dtype, + shape: tuple[int, ...], + ) -> torch.Tensor: + element_size = torch.tensor([], dtype=dtype).element_size() + expected = element_size + for dim in shape: + expected *= dim + if len(frame) != expected: + raise ValueError( + f"Tensor frame has {len(frame)} bytes, expected {expected} for " + f"shape={shape}, dtype={dtype}" + ) + source = torch.frombuffer(frame, dtype=dtype) + if dtype == torch.int64: + staging = self._positions_staging + elif dtype == torch.bfloat16: + staging = self._context_staging + else: + raise TypeError(f"Unsupported DSpark RPC tensor dtype: {dtype}") + if source.numel() > staging.numel(): + raise ValueError( + f"Tensor frame exceeds pinned staging capacity: " + f"elements={source.numel()}, capacity={staging.numel()}" + ) + output = staging[: source.numel()] + output.copy_(source) + return output.view(shape) + + def _append_context( + self, + states: list[DraftRequestState], + context_counts: list[int], + positions_cpu: torch.Tensor, + context_cpu: torch.Tensor, + *, + projected: bool, + ) -> None: + if positions_cpu.numel() == 0: + return + positions = positions_cpu.to(self.device, non_blocking=True) + context_input = context_cpu.to(self.device, non_blocking=True) + context_states = ( + context_input + if projected + else self.model.combine_hidden_states(context_input) + ) + + slots: list[int] = [] + offset = 0 + for state, count in zip(states, context_counts): + req_positions = positions_cpu[offset : offset + count] + if count: + first = int(req_positions[0]) + if first > state.committed_end: + raise ValueError( + f"Context gap for {state.request_id!r}: " + f"expected <= {state.committed_end}, got {first}" + ) + expected = torch.arange(first, first + count, dtype=torch.int64) + if not torch.equal(req_positions, expected): + raise ValueError( + f"Context positions for {state.request_id!r} are not contiguous" + ) + state.committed_end = int(req_positions[-1]) + 1 + slots.extend( + self.allocator.cache_slot(state, int(position)) + for position in req_positions + ) + offset += count + slot_mapping = torch.tensor(slots, dtype=torch.int64, device=self.device) + self.model.precompute_and_store_context_kv( + context_states, + positions, + slot_mapping, + ) + offset = 0 + for state, count in zip(states, context_counts): + if count: + first_position = int(positions_cpu[offset]) + if state.context_cache is None: + state.context_cache = ProjectedContextCache( + hidden_size=self.hidden_size, + max_tokens=self.prefix_cache_tokens, + initial_position=first_position, + device=self.device, + ) + state.context_cache.append( + first_position, + context_states[offset : offset + count], + ) + offset += count + + def _run_query_block( + self, + states: list[DraftRequestState], + anchor_positions: list[int], + anchor_token_ids: list[int], + num_speculative_tokens: int, + ) -> torch.Tensor: + batch_size = len(states) + sample_from_anchor = self.method == "dspark" + query_len = ( + num_speculative_tokens if sample_from_anchor else 1 + num_speculative_tokens + ) + input_ids: list[int] = [] + positions: list[int] = [] + slots: list[int] = [] + block_rows: list[list[int]] = [] + seq_lens: list[int] = [] + + for state, anchor_position, anchor_token_id in zip( + states, anchor_positions, anchor_token_ids + ): + if anchor_position != state.committed_end: + raise ValueError( + f"Anchor position for {state.request_id!r} must equal the " + f"committed context end ({state.committed_end}), got " + f"{anchor_position}" + ) + sequence_end = anchor_position + query_len + if sequence_end > self.max_model_len: + raise ValueError( + f"Draft query exceeds max_model_len={self.max_model_len}: " + f"end={sequence_end}" + ) + input_ids.append(anchor_token_id) + input_ids.extend([self.mask_token_id] * (query_len - 1)) + req_positions = range(anchor_position, sequence_end) + positions.extend(req_positions) + slots.extend( + self.allocator.cache_slot(state, position) + for position in range(anchor_position, sequence_end) + ) + row, local_seq_len = self.allocator.block_table(state, sequence_end) + block_rows.append(row) + seq_lens.append(local_seq_len) + + if self.cuda_graph_enabled: + graph_state = self._cuda_graphs.get((batch_size, num_speculative_tokens)) + if graph_state is not None: + graph_state.stage( + input_ids=input_ids, + positions=positions, + slots=slots, + block_rows=block_rows, + seq_lens=seq_lens, + ) + assert graph_state.graph is not None + graph_state.graph.replay() + self.cuda_graph_replay_count += 1 + return graph_state.output_tokens + self.cuda_graph_eager_fallback_count += 1 + + max_blocks = max(len(row) for row in block_rows) + block_table = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + for row_idx, row in enumerate(block_rows): + block_table[row_idx, : len(row)] = torch.tensor( + row, dtype=torch.int32, device=self.device + ) + input_ids_gpu = torch.tensor(input_ids, dtype=torch.int64, device=self.device) + positions_gpu = torch.tensor(positions, dtype=torch.int64, device=self.device) + slots_gpu = torch.tensor(slots, dtype=torch.int64, device=self.device) + seq_lens_gpu = torch.tensor(seq_lens, dtype=torch.int32, device=self.device) + query_start_loc = torch.arange( + 0, + (batch_size + 1) * query_len, + query_len, + dtype=torch.int32, + device=self.device, + ) + if self.method == "dflash": + query_start_loc_cpu = torch.arange( + 0, + (batch_size + 1) * query_len, + query_len, + dtype=torch.int32, + ) + common = CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=seq_lens_gpu, + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32), + max_seq_len=max(seq_lens), + num_reqs=batch_size, + num_actual_tokens=batch_size * query_len, + max_query_len=query_len, + block_table_tensor=block_table, + slot_mapping=slots_gpu, + causal=True, + ) + if self.runtime.attn_metadata_builder is None: + raise RuntimeError("K3 DFlash attention metadata builder is missing") + metadata = self.runtime.attn_metadata_builder.build(0, common) + else: + metadata = MLACommonMetadata( + num_reqs=batch_size, + max_query_len=query_len, + max_seq_len=max(seq_lens), + num_actual_tokens=batch_size * query_len, + query_start_loc=query_start_loc, + slot_mapping=slots_gpu, + num_decodes=batch_size, + num_decode_tokens=batch_size * query_len, + num_prefills=0, + causal=False, + head_dim=int(next(iter(self.runtime.kv_caches.values())).shape[-1]), + prefill=None, + decode=MLACommonDecodeMetadata( + block_table=block_table, + seq_lens=seq_lens_gpu, + dcp_tot_seq_lens=None, + ), + ) + attn_metadata = {layer_name: metadata for layer_name in self.runtime.kv_caches} + slot_mapping = {layer_name: slots_gpu for layer_name in self.runtime.kv_caches} + with ( + set_current_vllm_config(self.runtime.draft_vllm_config), + set_forward_context( + attn_metadata, + self.runtime.draft_vllm_config, + num_tokens=batch_size * query_len, + skip_compiled=True, + slot_mapping=slot_mapping, + ), + ): + hidden = self.model(input_ids=input_ids_gpu, positions=positions_gpu) + + if self.method == "dflash": + sample_hidden = hidden.view(batch_size, query_len, -1)[:, 1:] + return ( + self.model.compute_logits( + sample_hidden.reshape(batch_size * num_speculative_tokens, -1) + ) + .argmax(dim=-1) + .view(batch_size, num_speculative_tokens) + ) + + base_logits = self.model.compute_draft_logits(hidden).view( + batch_size, query_len, -1 + ) + previous = torch.tensor(anchor_token_ids, dtype=torch.int64, device=self.device) + draft_tokens = torch.empty( + (batch_size, query_len), dtype=torch.int64, device=self.device + ) + for step in range(query_len): + markov = self.model.markov_bias(self.model.markov_embed(previous)) + previous = (base_logits[:, step] + markov).argmax(dim=-1) + draft_tokens[:, step].copy_(previous) + return draft_tokens + + @torch.inference_mode() + def propose(self, header: dict[str, Any], frames: list[bytes]) -> dict[str, Any]: + started = time.perf_counter() + requests = header.get("requests") + if not isinstance(requests, list) or not requests: + raise ValueError("PROPOSE requires a non-empty requests list") + if len(requests) > self.allocator.max_requests: + raise ValueError( + f"Batch has {len(requests)} requests, max is " + f"{self.allocator.max_requests}" + ) + num_speculative_tokens = int( + header.get("num_speculative_tokens", self.max_speculative_tokens) + ) + if not 1 <= num_speculative_tokens <= self.max_speculative_tokens: + raise ValueError( + f"num_speculative_tokens must be in [1, {self.max_speculative_tokens}]" + ) + projected = bool(header.get("projected", False)) + context_counts = [int(req.get("context_count", 0)) for req in requests] + if any(count < 0 for count in context_counts): + raise ValueError("context_count cannot be negative") + total_context = sum(context_counts) + expected_frames = 2 if total_context else 0 + if len(frames) != expected_frames: + raise ValueError( + f"PROPOSE expected {expected_frames} tensor frames, got {len(frames)}" + ) + context_width = self.hidden_size if projected else self.raw_context_width + if total_context: + positions_cpu = self._decode_host_tensor( + frames[0], dtype=torch.int64, shape=(total_context,) + ) + context_cpu = self._decode_host_tensor( + frames[1], + dtype=torch.bfloat16, + shape=(total_context, context_width), + ) + else: + positions_cpu = torch.empty(0, dtype=torch.int64) + context_cpu = torch.empty((0, context_width), dtype=torch.bfloat16) + decoded_at = time.perf_counter() + + lock_started = time.perf_counter() + with self._lock: + lock_acquired = time.perf_counter() + gpu_start = torch.cuda.Event(enable_timing=True) + context_end = torch.cuda.Event(enable_timing=True) + query_end = torch.cuda.Event(enable_timing=True) + gpu_start.record() + states: list[DraftRequestState] = [] + for req in requests: + request_id = req.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("Every request requires a non-empty request_id") + state, created = self.allocator.get_or_allocate(request_id) + if bool(req.get("reset", False)) or created: + self._clear_state_cache(state) + reset_position = int(req.get("reset_position", 0)) + if not 0 <= reset_position <= self.max_model_len: + raise ValueError( + "Draft reset_position must be in model bounds, got " + f"{reset_position}" + ) + state.committed_end = reset_position + state.context_start = reset_position + if reset_position: + self.cold_bootstrap_count += 1 + states.append(state) + self._append_context( + states, + context_counts, + positions_cpu, + context_cpu, + projected=projected, + ) + context_end.record() + anchor_positions = [int(req["anchor_position"]) for req in requests] + anchor_token_ids = [int(req["anchor_token_id"]) for req in requests] + draft_tokens = self._run_query_block( + states, + anchor_positions, + anchor_token_ids, + num_speculative_tokens, + ) + query_end.record() + submit_done = time.perf_counter() + query_end.synchronize() + sync_done = time.perf_counter() + tokens = draft_tokens.cpu().tolist() + tokens_copied = time.perf_counter() + self.proposal_count += 1 + self.last_latency_ms = (tokens_copied - started) * 1000 + timing_ms = { + "decode_frames": (decoded_at - started) * 1000, + "lock_wait": (lock_acquired - lock_started) * 1000, + "host_submit": (submit_done - lock_acquired) * 1000, + "gpu_context": gpu_start.elapsed_time(context_end), + "gpu_query": context_end.elapsed_time(query_end), + "gpu_wait": (sync_done - submit_done) * 1000, + "tokens_d2h": (tokens_copied - sync_done) * 1000, + "total": self.last_latency_ms, + } + self._record_timing(timing_ms) + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "tokens": tokens, + "latency_ms": self.last_latency_ms, + "timing_ms": timing_ms, + "active_requests": self.allocator.active_requests, + } + + +class K3DSparkZMQServer: + def __init__( + self, + engine: K3DSparkDraftEngine, + *, + address: str, + stop: threading.Event, + ) -> None: + self.engine = engine + self.address = address + self.stop = stop + self.ready = threading.Event() + self.error: str | None = None + self._thread = threading.Thread( + target=self._run, + name="k3-draft-zmq", + daemon=True, + ) + + def start(self) -> None: + self._thread.start() + if not self.ready.wait(timeout=10): + raise TimeoutError(f"K3 draft proposal socket did not bind: {self.address}") + if self.error is not None: + raise RuntimeError(self.error) + + def join(self, timeout: float = 5.0) -> None: + self._thread.join(timeout=timeout) + + def _handle(self, parts: list[bytes]) -> dict[str, Any]: + if not parts: + raise ValueError("Empty DSpark RPC message") + header = json.loads(parts[0]) + if not isinstance(header, dict): + raise ValueError("DSpark RPC header must be a JSON object") + if int(header.get("protocol", -1)) != PROTOCOL_VERSION: + raise ValueError( + f"Unsupported protocol {header.get('protocol')}; " + f"expected {PROTOCOL_VERSION}" + ) + op = str(header.get("op", "")).upper() + if op == "PING": + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "op": "PONG", + "method": self.engine.method, + "active_requests": self.engine.allocator.active_requests, + "max_requests": self.engine.allocator.max_requests, + "block_size": self.engine.allocator.block_size, + "window_size": self.engine.allocator.window_size, + "prefix_cache_tokens": self.engine.prefix_cache_tokens, + "prefix_cache_device": str(self.engine.device), + "cold_bootstrap_count": self.engine.cold_bootstrap_count, + "cuda_graph_enabled": self.engine.cuda_graph_enabled, + "cuda_graph_shapes": self.engine.cuda_graph_shapes, + } + if op == "CLEAR": + self.engine.clear() + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "active_requests": 0, + } + if op in ("RESET", "FREE"): + request_ids = header.get("request_ids") + if not isinstance(request_ids, list) or not all( + isinstance(req_id, str) and req_id for req_id in request_ids + ): + raise ValueError(f"{op} requires request_ids: list[str]") + if op == "RESET": + self.engine.reset(request_ids) + else: + self.engine.free(request_ids) + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "active_requests": self.engine.allocator.active_requests, + } + if op == "RECONNECT": + source_request_id = header.get("source_request_id") + request_id = header.get("request_id") + if not isinstance(source_request_id, str) or not source_request_id: + raise ValueError("RECONNECT requires source_request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("RECONNECT requires request_id") + return self.engine.reconnect( + source_request_id, + request_id, + int(header.get("prefix_end", 0)), + ) + if op == "PROPOSE": + return self.engine.propose(header, parts[1:]) + raise ValueError(f"Unknown K3 draft RPC operation: {op!r}") + + def _run(self) -> None: + context = zmq.Context() + socket = context.socket(zmq.REP) + socket.setsockopt(zmq.LINGER, 0) + try: + socket.bind(self.address) + logger.info( + "K3 %s proposal RPC listening on %s", + self.engine.method, + self.address, + ) + self.ready.set() + poller = zmq.Poller() + poller.register(socket, zmq.POLLIN) + while not self.stop.is_set(): + if not dict(poller.poll(250)).get(socket): + continue + try: + response = self._handle(socket.recv_multipart()) + except Exception as exc: + logger.exception("K3 DSpark proposal request failed") + response = { + "ok": False, + "protocol": PROTOCOL_VERSION, + "error": f"{type(exc).__name__}: {exc}", + } + socket.send_json(response) + except Exception as exc: + self.error = f"{type(exc).__name__}: {exc}" + logger.exception("K3 DSpark proposal server failed") + self.ready.set() + finally: + socket.close() + context.term() diff --git a/vllm/entrypoints/k3_dspark_standalone.py b/vllm/entrypoints/k3_dspark_standalone.py new file mode 100644 index 000000000000..8f3fe695a6ed --- /dev/null +++ b/vllm/entrypoints/k3_dspark_standalone.py @@ -0,0 +1,877 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Load a Kimi-K3 draft model on a dedicated single GPU. + +The process loads either DSpark or DFlash plus the target embedding and LM +head, without loading the Kimi-K3 target transformer. Draft weights, KV cache, +attention workspace, and proposal compute stay on this process's GPU. +""" + +from __future__ import annotations + +import argparse +import json +import os +import signal +import threading +import time +from dataclasses import asdict, dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from safetensors import safe_open + +from vllm.config.vllm import set_current_vllm_config +from vllm.engine.arg_utils import EngineArgs +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.utils.torch_utils import set_default_torch_dtype +from vllm.v1.worker.workspace import init_workspace_manager + +logger = init_logger(__name__) + +EMBED_TENSOR = "language_model.model.embed_tokens.weight" +LM_HEAD_TENSOR = "language_model.lm_head.weight" +SHARED_TENSORS = (EMBED_TENSOR, LM_HEAD_TENSOR) + + +@dataclass +class RuntimeStatus: + phase: str = "starting" + ready: bool = False + proposal_transport_ready: bool = False + device: str = "" + compute_capability: str = "" + torch_version: str = torch.__version__ + torch_arches: tuple[str, ...] = () + draft_model: str = "" + method: str = "" + target_weights: str = "" + num_speculative_tokens: int = 0 + max_model_len: int = 0 + allocated_gib: float = 0.0 + reserved_gib: float = 0.0 + draft_kv_cache_gib: float = 0.0 + draft_kv_cache_blocks: int = 0 + draft_kv_cache_token_capacity: int = 0 + draft_kv_cache_smoke: bool = False + draft_kv_window: int = 0 + proposal_address: str | None = None + proposal_count: int = 0 + proposal_active_requests: int = 0 + proposal_last_latency_ms: float = 0.0 + smoke_token: int | None = None + load_seconds: float = 0.0 + error: str | None = None + + +class _TargetLanguageModel(nn.Module): + def __init__(self, vocab_size: int, hidden_size: int) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding( + vocab_size, + hidden_size, + params_dtype=torch.bfloat16, + prefix="model.embed_tokens", + disable_tp=True, + ) + + +class StandaloneTargetFacade(nn.Module): + """Only the two frozen target modules that DSpark shares.""" + + def __init__(self, vocab_size: int, hidden_size: int) -> None: + super().__init__() + self.model = _TargetLanguageModel(vocab_size, hidden_size) + self.lm_head = ParallelLMHead( + vocab_size, + hidden_size, + params_dtype=torch.bfloat16, + prefix="lm_head", + disable_tp=True, + ) + + def get_language_model(self) -> StandaloneTargetFacade: + return self + + +@dataclass +class StandaloneRuntime: + vllm_config: Any + draft_vllm_config: Any + target_facade: StandaloneTargetFacade + model: nn.Module + kv_caches: dict[str, torch.Tensor] + kv_cache_block_size: int + method: str + attn_metadata_builder: Any | None = None + + +def resolve_shared_weight_files(target_weights: Path) -> dict[str, Path]: + """Resolve the two shared tensors through a safetensors index.""" + root = target_weights.resolve() + index_path = root / "model.safetensors.index.json" + if not index_path.is_file(): + raise FileNotFoundError(f"Target weight index is missing: {index_path}") + + index = json.loads(index_path.read_text()) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"Invalid safetensors weight_map in {index_path}") + + resolved: dict[str, Path] = {} + for tensor_name in SHARED_TENSORS: + relative = weight_map.get(tensor_name) + if not isinstance(relative, str): + raise KeyError(f"Target tensor is absent from the index: {tensor_name}") + tensor_path = (root / relative).resolve() + if not tensor_path.is_relative_to(root): + raise ValueError( + f"Target tensor path escapes the checkpoint root: {tensor_path}" + ) + if not tensor_path.is_file(): + raise FileNotFoundError(f"Target tensor file is missing: {tensor_path}") + resolved[tensor_name] = tensor_path + return resolved + + +def _load_module_weight( + module: VocabParallelEmbedding, + tensor_name: str, + tensor_path: Path, + device: torch.device, +) -> None: + logger.info("Loading shared target tensor %s from %s", tensor_name, tensor_path) + with safe_open(str(tensor_path), framework="pt", device=str(device)) as handle: + if tensor_name not in set(handle.keys()): + raise KeyError(f"{tensor_name} is absent from {tensor_path}") + loaded_weight = handle.get_tensor(tensor_name) + if loaded_weight.dtype != torch.bfloat16: + raise TypeError( + f"{tensor_name} must be bfloat16, got {loaded_weight.dtype}" + ) + module.weight_loader(module.weight, loaded_weight) + del loaded_weight + torch.accelerator.empty_cache() + + +def load_shared_target_weights( + target: StandaloneTargetFacade, + target_weights: Path, + device: torch.device, +) -> None: + files = resolve_shared_weight_files(target_weights) + _load_module_weight( + target.model.embed_tokens, + EMBED_TENSOR, + files[EMBED_TENSOR], + device, + ) + _load_module_weight( + target.lm_head, + LM_HEAD_TENSOR, + files[LM_HEAD_TENSOR], + device, + ) + + +def _validate_cuda_runtime(device: torch.device) -> tuple[str, tuple[str, ...]]: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available in the DSpark container") + major, minor = torch.cuda.get_device_capability(device) + expected_arch = f"sm_{major}{minor}" + arches = tuple(torch.cuda.get_arch_list()) + if expected_arch not in arches: + raise RuntimeError( + f"PyTorch does not contain {expected_arch}; compiled arches are {arches}" + ) + if major < 8: + raise RuntimeError( + f"Kimi-K3 DSpark BF16 requires compute capability >= 8.0, got " + f"{major}.{minor}" + ) + return f"{major}.{minor}", arches + + +def _init_single_gpu_distributed() -> None: + from vllm.distributed.parallel_state import ( + ensure_model_parallel_initialized, + init_distributed_environment, + model_parallel_is_initialized, + ) + from vllm.utils.network_utils import get_open_port + + if model_parallel_is_initialized(): + return + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", + backend="gloo", + ) + ensure_model_parallel_initialized( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + +def _build_vllm_config(args: argparse.Namespace): + attention_backend = "TRITON_ATTN" if args.method == "dflash" else "TRITON_MLA" + speculative_config = { + "model": str(args.draft_model), + "method": args.method, + "num_speculative_tokens": args.num_speculative_tokens, + "attention_backend": attention_backend, + "kv_cache_dtype": "bfloat16", + "draft_sample_method": "greedy", + "rejection_sample_method": "block", + "max_model_len": args.max_model_len, + } + if args.draft_quantization != "none": + speculative_config["quantization"] = args.draft_quantization + engine_args = EngineArgs( + model=str(args.target_config), + tokenizer_mode="skip", + skip_tokenizer_init=True, + dtype="bfloat16", + kv_cache_dtype="bfloat16", + max_model_len=args.max_model_len, + tensor_parallel_size=1, + decode_context_parallel_size=1, + max_num_batched_tokens=args.max_num_batched_tokens, + max_num_seqs=args.max_num_seqs, + block_size=16, + enable_prefix_caching=False, + enforce_eager=True, + compilation_config={"custom_ops": ["none"]}, + kernel_config={ + "ir_op_priority": { + "rms_norm": ["native"], + "fused_add_rms_norm": ["native"], + }, + "linear_backend": "torch", + }, + load_format="safetensors", + use_tqdm_on_load=False, + speculative_config=speculative_config, + ) + return engine_args.create_engine_config(headless=True) + + +def _allocate_draft_kv_cache( + model: nn.Module, + *, + method: str, + block_size: int, + cache_gib: float, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], int]: + """Allocate and bind the draft model's private BF16 KV cache.""" + if cache_gib <= 0: + raise ValueError(f"--draft-kv-cache-gib must be positive, got {cache_gib}") + + if method == "dflash": + attentions = [layer.self_attn.attn for layer in model.model.layers] + else: + attentions = [layer.self_attn for layer in model.model.layers] + if not attentions: + raise RuntimeError("K3 draft model does not expose any attention layers") + + if method == "dflash": + cache_shapes = { + ( + int(attn.impl.num_kv_heads), + int(attn.impl.head_size), + ) + for attn in attentions + } + if len(cache_shapes) != 1: + raise ValueError(f"K3 DFlash layers have mixed KV shapes: {cache_shapes}") + num_kv_heads, head_size = cache_shapes.pop() + bytes_per_block_all_layers = ( + len(attentions) + * block_size + * num_kv_heads + * 2 + * head_size + * torch.tensor([], dtype=torch.bfloat16).element_size() + ) + else: + latent_widths = { + int(attn.kv_lora_rank + attn.qk_rope_head_dim) for attn in attentions + } + if len(latent_widths) != 1: + raise ValueError( + f"K3 DSpark draft layers have mixed MLA widths: {latent_widths}" + ) + latent_width = latent_widths.pop() + bytes_per_block_all_layers = ( + len(attentions) + * block_size + * latent_width + * torch.tensor([], dtype=torch.bfloat16).element_size() + ) + num_blocks = int(cache_gib * 1024**3) // bytes_per_block_all_layers + # Block zero is deliberately never assigned to a live request. + if num_blocks < 2: + minimum_mib = 2 * bytes_per_block_all_layers / 1024**2 + raise ValueError( + "The draft KV cache must fit a null block plus one data block; " + f"minimum={minimum_mib:.2f} MiB" + ) + + caches: dict[str, torch.Tensor] = {} + for attn in attentions: + if method == "dflash": + # TritonAttention exposes logical B,H,N,2D while its preferred + # physical layout is B,N,H,2D on CUDA. + physical = torch.zeros( + (num_blocks, block_size, num_kv_heads, 2 * head_size), + dtype=torch.bfloat16, + device=device, + ) + cache = physical.permute(0, 2, 1, 3) + else: + cache = torch.zeros( + (num_blocks, block_size, latent_width), + dtype=torch.bfloat16, + device=device, + ) + attn.bind_kv_cache(cache) + caches[attn.layer_name] = cache + return caches, num_blocks + + +def _build_dflash_metadata_builder( + model: nn.Module, + draft_vllm_config: Any, + device: torch.device, +) -> Any: + from vllm.v1.worker.utils import AttentionGroup + + attentions = [layer.self_attn.attn for layer in model.model.layers] + layer_names = [attn.layer_name for attn in attentions] + first = attentions[0] + backend = first.get_attn_backend() + if backend.get_name() != "TRITON_ATTN": + raise ValueError( + f"Standalone K3 DFlash requires TRITON_ATTN, got {backend.get_name()}" + ) + kv_spec = first.get_kv_cache_spec(draft_vllm_config) + if kv_spec is None: + raise RuntimeError("K3 DFlash attention did not return a KV cache spec") + group = AttentionGroup(backend, layer_names, kv_spec, 0) + group.create_metadata_builders( + draft_vllm_config, + device, + kernel_block_size=int(draft_vllm_config.cache_config.block_size), + ) + return group.get_metadata_builder() + + +def _load_runtime(args: argparse.Namespace, status: RuntimeStatus) -> StandaloneRuntime: + start = time.perf_counter() + device = torch.device("cuda", args.device) + torch.cuda.set_device(device) + status.device = torch.cuda.get_device_name(device) + status.compute_capability, status.torch_arches = _validate_cuda_runtime(device) + status.phase = "building_config" + + vllm_config = _build_vllm_config(args) + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + draft_config = speculative_config.draft_model_config.hf_config + + status.phase = "loading_shared_target_weights" + with set_current_vllm_config(vllm_config): + _init_single_gpu_distributed() + init_workspace_manager(device, num_lanes=2) + with torch.device(device), set_default_torch_dtype(torch.bfloat16): + target = StandaloneTargetFacade( + vocab_size=draft_config.vocab_size, + hidden_size=draft_config.hidden_size, + ) + load_shared_target_weights(target, args.target_weights, device) + + status.phase = f"loading_{args.method}" + from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( + _create_draft_vllm_config, + ) + + if args.method == "dflash": + from vllm.v1.worker.gpu.spec_decode.dflash.utils import ( + load_dflash_model, + maybe_load_mask_embedding, + ) + + model = load_dflash_model(target, vllm_config) + maybe_load_mask_embedding( + model, + str(args.draft_model), + int(draft_config.dflash_config["mask_token_id"]), + ) + else: + from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model + + model = load_dspark_model(target, vllm_config) + model.eval() + draft_vllm_config = _create_draft_vllm_config(vllm_config) + + status.phase = "allocating_draft_kv_cache" + block_size = int(vllm_config.cache_config.block_size) + kv_caches, num_blocks = _allocate_draft_kv_cache( + model, + method=args.method, + block_size=block_size, + cache_gib=args.draft_kv_cache_gib, + device=device, + ) + status.draft_kv_cache_blocks = num_blocks + status.draft_kv_cache_token_capacity = (num_blocks - 1) * block_size + status.draft_kv_cache_gib = ( + sum(cache.numel() * cache.element_size() for cache in kv_caches.values()) + / 1024**3 + ) + + torch.accelerator.synchronize() + status.load_seconds = time.perf_counter() - start + status.allocated_gib = torch.cuda.memory_allocated(device) / 1024**3 + status.reserved_gib = torch.cuda.memory_reserved(device) / 1024**3 + metadata_builder = ( + _build_dflash_metadata_builder(model, draft_vllm_config, device) + if args.method == "dflash" + else None + ) + return StandaloneRuntime( + vllm_config, + draft_vllm_config, + target, + model, + kv_caches, + block_size, + args.method, + metadata_builder, + ) + + +@torch.inference_mode() +def _run_eager_smoke(runtime: StandaloneRuntime, device: torch.device) -> int: + if runtime.method == "dflash": + return _run_dflash_eager_smoke(runtime, device) + + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonDecodeMetadata, + MLACommonMetadata, + ) + from vllm.v1.worker.gpu.spec_decode.utils import ( + get_parallel_drafting_token_id, + ) + + model = runtime.model + draft_config = runtime.vllm_config.speculative_config.draft_model_config.hf_config + num_aux_layers = int(draft_config.num_target_layers) + hidden_size = int(draft_config.hidden_size) + + aux = torch.zeros( + (1, num_aux_layers * hidden_size), + dtype=torch.bfloat16, + device=device, + ) + context_len = 1 + positions = torch.zeros(context_len, dtype=torch.int64, device=device) + context = model.combine_hidden_states(aux) + data_block = 1 + context_slots = torch.tensor( + [data_block * runtime.kv_cache_block_size], + dtype=torch.int64, + device=device, + ) + model.precompute_and_store_context_kv(context, positions, context_slots) + + sample_from_anchor = bool(getattr(draft_config, "sample_from_anchor", True)) + query_len = runtime.vllm_config.speculative_config.num_speculative_tokens + if not sample_from_anchor: + query_len += 1 + mask_token_id = get_parallel_drafting_token_id(draft_config) + input_ids = torch.tensor( + [draft_config.bos_token_id] + [mask_token_id] * (query_len - 1), + dtype=torch.int64, + device=device, + ) + query_positions = torch.arange( + context_len, + context_len + query_len, + dtype=torch.int64, + device=device, + ) + query_slots = torch.arange( + data_block * runtime.kv_cache_block_size + context_len, + data_block * runtime.kv_cache_block_size + context_len + query_len, + dtype=torch.int64, + device=device, + ) + query_start_loc = torch.tensor([0, query_len], dtype=torch.int32, device=device) + block_table = torch.tensor([[data_block]], dtype=torch.int32, device=device) + seq_lens = torch.tensor([context_len + query_len], dtype=torch.int32, device=device) + metadata = MLACommonMetadata( + num_reqs=1, + max_query_len=query_len, + max_seq_len=context_len + query_len, + num_actual_tokens=query_len, + query_start_loc=query_start_loc, + slot_mapping=query_slots, + num_decodes=1, + num_decode_tokens=query_len, + num_prefills=0, + causal=False, + # MLA metadata validates the cached latent width, not the full Q/K + # projection width. + head_dim=int(next(iter(runtime.kv_caches.values())).shape[-1]), + prefill=None, + decode=MLACommonDecodeMetadata( + block_table=block_table, + seq_lens=seq_lens, + dcp_tot_seq_lens=None, + ), + ) + attn_metadata = {layer_name: metadata for layer_name in runtime.kv_caches} + slot_mapping = {layer_name: query_slots for layer_name in runtime.kv_caches} + with ( + set_current_vllm_config(runtime.draft_vllm_config), + set_forward_context( + attn_metadata, + runtime.draft_vllm_config, + num_tokens=query_len, + skip_compiled=True, + slot_mapping=slot_mapping, + ), + ): + hidden = model(input_ids=input_ids, positions=query_positions) + base_logits = model.compute_draft_logits(hidden[-1:]) + markov = model.markov_bias(model.markov_embed(input_ids[-1:])) + token = int((base_logits + markov).argmax(dim=-1).item()) + torch.accelerator.synchronize() + for cache in runtime.kv_caches.values(): + cache[data_block].zero_() + return token + + +@torch.inference_mode() +def _run_dflash_eager_smoke( + runtime: StandaloneRuntime, + device: torch.device, +) -> int: + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, + ) + from vllm.v1.worker.gpu.spec_decode.utils import ( + get_parallel_drafting_token_id, + ) + + speculative_config = runtime.vllm_config.speculative_config + assert speculative_config is not None + draft_config = speculative_config.draft_model_config.hf_config + aux_layers = get_eagle3_aux_layers_from_config(speculative_config) + if not aux_layers: + raise ValueError("K3 DFlash config does not declare target auxiliary layers") + + raw_context = torch.zeros( + (1, len(aux_layers) * int(draft_config.hidden_size)), + dtype=torch.bfloat16, + device=device, + ) + context = runtime.model.combine_hidden_states(raw_context) + block_size = runtime.kv_cache_block_size + context_position = torch.zeros(1, dtype=torch.int64, device=device) + context_slot = torch.tensor([block_size], dtype=torch.int64, device=device) + runtime.model.precompute_and_store_context_kv( + context, + context_position, + context_slot, + ) + + query_len = int(speculative_config.num_speculative_tokens) + 1 + sequence_end = 1 + query_len + block_ids = list(range(1, 1 + (sequence_end + block_size - 1) // block_size)) + input_ids = torch.tensor( + [draft_config.bos_token_id] + + [get_parallel_drafting_token_id(draft_config)] * (query_len - 1), + dtype=torch.int64, + device=device, + ) + positions = torch.arange(1, sequence_end, dtype=torch.int64, device=device) + slots = torch.tensor( + [ + block_ids[position // block_size] * block_size + position % block_size + for position in range(1, sequence_end) + ], + dtype=torch.int64, + device=device, + ) + query_start_cpu = torch.tensor([0, query_len], dtype=torch.int32) + query_start_gpu = query_start_cpu.to(device) + seq_lens = torch.tensor([sequence_end], dtype=torch.int32, device=device) + block_table = torch.tensor([block_ids], dtype=torch.int32, device=device) + common = CommonAttentionMetadata( + query_start_loc=query_start_gpu, + query_start_loc_cpu=query_start_cpu, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=torch.tensor([sequence_end], dtype=torch.int32), + max_seq_len=sequence_end, + num_reqs=1, + num_actual_tokens=query_len, + max_query_len=query_len, + block_table_tensor=block_table, + slot_mapping=slots, + causal=True, + ) + assert runtime.attn_metadata_builder is not None + metadata = runtime.attn_metadata_builder.build(0, common) + attn_metadata = {layer_name: metadata for layer_name in runtime.kv_caches} + slot_mapping = {layer_name: slots for layer_name in runtime.kv_caches} + with ( + set_current_vllm_config(runtime.draft_vllm_config), + set_forward_context( + attn_metadata, + runtime.draft_vllm_config, + num_tokens=query_len, + skip_compiled=True, + slot_mapping=slot_mapping, + ), + ): + hidden = runtime.model(input_ids=input_ids, positions=positions) + token = int(runtime.model.compute_logits(hidden[1:2]).argmax(dim=-1).item()) + torch.accelerator.synchronize() + for cache in runtime.kv_caches.values(): + cache[1 : 1 + len(block_ids)].zero_() + return token + + +def _make_handler(status: RuntimeStatus, proposal_engine: Any | None = None): + class StatusHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 + if self.path not in ("/", "/healthz", "/readyz", "/v1/status"): + self.send_error(404) + return + payload = asdict(status) + if proposal_engine is not None: + payload["proposal_count"] = proposal_engine.proposal_count + payload["proposal_active_requests"] = ( + proposal_engine.allocator.active_requests + ) + payload["proposal_max_requests"] = ( + proposal_engine.allocator.max_requests + ) + payload["proposal_last_latency_ms"] = proposal_engine.last_latency_ms + payload["proposal_last_timing_ms"] = proposal_engine.last_timing_ms + payload["proposal_mean_timing_ms"] = proposal_engine.mean_timing_ms + payload["proposal_cold_bootstrap_count"] = ( + proposal_engine.cold_bootstrap_count + ) + payload["proposal_reconnect_count"] = proposal_engine.reconnect_count + payload["proposal_last_reconnect_latency_ms"] = ( + proposal_engine.last_reconnect_latency_ms + ) + payload["proposal_prefix_cache_tokens"] = ( + proposal_engine.prefix_cache_tokens + ) + payload["proposal_prefix_cache_host_gib"] = ( + proposal_engine.prefix_cache_host_bytes / 1024**3 + ) + payload["proposal_prefix_cache_gpu_gib"] = ( + proposal_engine.prefix_cache_device_bytes / 1024**3 + ) + payload["proposal_cuda_graph_enabled"] = ( + proposal_engine.cuda_graph_enabled + ) + payload["proposal_cuda_graph_shapes"] = ( + proposal_engine.cuda_graph_shapes + ) + payload["proposal_cuda_graph_capture_seconds"] = ( + proposal_engine.cuda_graph_capture_seconds + ) + payload["proposal_cuda_graph_memory_gib"] = ( + proposal_engine.cuda_graph_memory_gib + ) + payload["proposal_cuda_graph_replay_count"] = ( + proposal_engine.cuda_graph_replay_count + ) + payload["proposal_cuda_graph_eager_fallback_count"] = ( + proposal_engine.cuda_graph_eager_fallback_count + ) + body = json.dumps(payload, sort_keys=True).encode() + code = 200 if status.ready else 503 + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + logger.info("DSpark status HTTP: %s", format % args) + + return StatusHandler + + +def _serve_status( + args: argparse.Namespace, + status: RuntimeStatus, + stop: threading.Event, + proposal_engine: Any | None = None, +) -> None: + server = ThreadingHTTPServer( + (args.host, args.port), _make_handler(status, proposal_engine) + ) + server.timeout = 0.5 + + def request_stop(signum: int, _frame: object) -> None: + logger.info("Received signal %d; stopping DSpark status server", signum) + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + logger.info( + "DSpark status endpoint listening on http://%s:%d", args.host, args.port + ) + while not stop.is_set(): + server.handle_request() + server.server_close() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--method", choices=("dspark", "dflash"), default="dspark") + parser.add_argument("--draft-model", type=Path, required=True) + parser.add_argument("--target-weights", type=Path, required=True) + parser.add_argument("--target-config", type=Path, required=True) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8091) + parser.add_argument("--num-speculative-tokens", type=int, default=3) + parser.add_argument("--max-model-len", type=int, default=32768) + parser.add_argument("--max-num-batched-tokens", type=int, default=64) + parser.add_argument("--max-num-seqs", type=int, default=2) + parser.add_argument("--max-retained-requests", type=int) + parser.add_argument("--draft-kv-cache-gib", type=float, default=1.0) + parser.add_argument("--draft-kv-window", type=int, default=32768) + parser.add_argument("--proposal-address", default="tcp://127.0.0.1:8092") + parser.add_argument("--disable-proposal-transport", action="store_true") + parser.add_argument( + "--draft-quantization", + choices=("none", "fp8_per_channel", "fp8_per_tensor", "mxfp8"), + default="none", + help=( + "Online weight quantization for the draft's linear layers (vLLM " + "online-quantization shorthand): fp8_per_channel quantizes each " + "BF16 weight to float8_e4m3 with one scale per output channel and " + "dynamic per-token activation scaling. Draft-time only: the " + "target's verification pass is unchanged, so accepted outputs " + "keep the target distribution; only the proposal quality (and " + "therefore acceptance length) can move." + ), + ) + parser.add_argument( + "--draft-fp8-head", + action="store_true", + help=( + "Score draft proposals with a rowwise-fp8 copy of the LM head " + "(VLLM_DSPARK_FP8_DRAFT_HEAD=1); same draft-time-only contract " + "as --draft-quantization." + ), + ) + parser.add_argument("--enable-cuda-graph", action="store_true") + parser.add_argument("--cuda-graph-warmups", type=int, default=2) + parser.add_argument("--skip-smoke-test", action="store_true") + parser.add_argument("--exit-after-load", action="store_true") + args = parser.parse_args() + if ( + args.max_retained_requests is not None + and args.max_retained_requests < args.max_num_seqs + ): + parser.error("--max-retained-requests must be >= --max-num-seqs") + return args + + +def main() -> None: + args = _parse_args() + if args.draft_fp8_head: + # Read by the draft loader (envs.VLLM_DSPARK_FP8_DRAFT_HEAD) while + # _load_runtime builds the draft head, so it is set before that. + os.environ["VLLM_DSPARK_FP8_DRAFT_HEAD"] = "1" + status = RuntimeStatus( + draft_model=str(args.draft_model), + method=args.method, + target_weights=str(args.target_weights), + num_speculative_tokens=args.num_speculative_tokens, + max_model_len=args.max_model_len, + draft_kv_window=args.draft_kv_window, + ) + try: + runtime = _load_runtime(args, status) + if not args.skip_smoke_test: + status.phase = "eager_smoke_test" + status.smoke_token = _run_eager_smoke( + runtime, torch.device("cuda", args.device) + ) + status.draft_kv_cache_smoke = True + stop = threading.Event() + proposal_engine = None + proposal_server = None + if not args.disable_proposal_transport: + status.phase = "starting_proposal_transport" + from vllm.entrypoints.k3_dspark_rpc import ( + K3DSparkDraftEngine, + K3DSparkZMQServer, + ) + + proposal_engine = K3DSparkDraftEngine( + runtime, + max_requests=( + args.max_retained_requests + if args.max_retained_requests is not None + else args.max_num_seqs + ), + window_size=args.draft_kv_window, + device=torch.device("cuda", args.device), + ) + if args.enable_cuda_graph: + status.phase = f"capturing_{args.method}_cuda_graphs" + proposal_engine.capture_cuda_graphs(warmups=args.cuda_graph_warmups) + proposal_server = K3DSparkZMQServer( + proposal_engine, + address=args.proposal_address, + stop=stop, + ) + proposal_server.start() + status.proposal_transport_ready = True + status.proposal_address = args.proposal_address + status.phase = "ready" + else: + status.phase = "ready_without_transport" + status.ready = True + status.allocated_gib = torch.cuda.memory_allocated(args.device) / 1024**3 + status.reserved_gib = torch.cuda.memory_reserved(args.device) / 1024**3 + logger.info("Standalone K3 draft is loaded: %s", json.dumps(asdict(status))) + if args.exit_after_load: + print(json.dumps(asdict(status), sort_keys=True), flush=True) + return + _serve_status(args, status, stop, proposal_engine) + if proposal_server is not None: + proposal_server.join() + except Exception as exc: + status.phase = "failed" + status.error = f"{type(exc).__name__}: {exc}" + logger.exception("Standalone K3 draft startup failed") + raise + + +if __name__ == "__main__": + main() diff --git a/vllm/envs.py b/vllm/envs.py index 8bcd9d0bc7f2..92d8c3e0098a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -247,6 +247,7 @@ VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True VLLM_MOE_SKIP_PADDING: bool = True VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False + VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1844,6 +1845,13 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool( int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0")) ), + # Kimi K3 only, and unrelated to the MoE flags above. Tap the pre-norm + # AttnRes mixture, rather than the post-mixture sum, as the auxiliary + # hidden state handed to a DFlash drafter. This changes the numerics the + # speculator sees, so it is off by default while the effect is measured. + "VLLM_KIMI_K3_AUX_ATTN_RES_STREAM": lambda: bool( + int(os.getenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "0")) + ), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index a3da0fe82a79..6c0bbd1deb33 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -282,27 +282,28 @@ def extra_repr(self): f"max_width={self.max_width}, theta_base={self.theta_base}" ) - def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: - """Calculate the cis(freqs) for each position in the 2D grid.""" - N = self.max_height * self.max_width - flat_pos = torch.arange(0, N).float().to(device) - x_pos = flat_pos % self.max_width - y_pos = flat_pos // self.max_width - dim_range = ( - torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) - ) # C/4 + def _compute_grid_freqs_cis( + self, height: int, width: int, device: torch.device + ) -> torch.Tensor: + """Calculate rotary frequencies for one requested image grid.""" + x_pos = torch.arange(width, dtype=torch.float32, device=device) + y_pos = torch.arange(height, dtype=torch.float32, device=device) + dim_range = torch.arange(0, self.dim, 4, dtype=torch.float32, device=device)[ + : (self.dim // 4) + ] freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) - x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 - y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 - x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 - y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 - # N, C/4, 2 - freqs_cis = torch.cat( - [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 - ) - # max_height, max_width, C/2 - freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) - return freqs_cis + x_freqs = torch.outer(x_pos, freqs).float() + y_freqs = torch.outer(y_pos, freqs).float() + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) + freqs_cis = torch.stack( + ( + x_cis.unsqueeze(0).expand(height, -1, -1), + y_cis.unsqueeze(1).expand(-1, width, -1), + ), + dim=-1, + ) + return freqs_cis.reshape(height * width, self.dim // 2) def get_freqs_cis( self, grid_thws: torch.Tensor | list[list[int]], device: torch.device @@ -314,11 +315,6 @@ def get_freqs_cis( Returns: freqs_cis: tensor of shape (sum(t * height * width), dim//2) """ - if not hasattr(self, "freqs_cis"): - self.register_buffer( - "freqs_cis", self._precompute_freqs_cis(device), persistent=False - ) - shapes = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() assert all( 1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes @@ -327,14 +323,15 @@ def get_freqs_cis( self.max_height, self.max_width, ) - freqs_cis = torch.cat( - [ - self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) - for t, h, w in shapes - ], - dim=0, - ) - return freqs_cis + grids: dict[tuple[int, int], torch.Tensor] = {} + result = [] + for t, h, w in shapes: + grid = grids.get((h, w)) + if grid is None: + grid = self._compute_grid_freqs_cis(h, w, device) + grids[(h, w)] = grid + result.append(grid.repeat(t, 1)) + return torch.cat(result, dim=0) class MLP2(nn.Module): @@ -843,21 +840,24 @@ def prepare_encoder_cudagraph_metadata( @torch.inference_mode() def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]): - """Apply MM projector to vision tower outputs.""" - num_embedding_list = [x.shape[0] for x in vt_output] - batched = torch.cat(vt_output, dim=0) + """Apply the projector without concatenating independent image features.""" + if not vt_output: + raise ValueError("Kimi vision projection requires at least one image feature") + projector_norm = getattr(mm_projector, "pre_norm", None) if projector_norm is None: projector_norm = getattr(mm_projector, "post_norm", None) projector_dtype = ( - projector_norm.weight.dtype if projector_norm is not None else batched.dtype + projector_norm.weight.dtype if projector_norm is not None else None ) - if batched.dtype != projector_dtype: - batched = batched.to(projector_dtype) - proj_out = mm_projector(batched) - proj_out = proj_out.reshape(-1, proj_out.shape[-1]) - proj_out = torch.split(proj_out, num_embedding_list) - return proj_out + + projected = [] + for image_features in vt_output: + if projector_dtype is not None and image_features.dtype != projector_dtype: + image_features = image_features.to(projector_dtype) + output = mm_projector(image_features) + projected.append(output.reshape(-1, output.shape[-1])) + return tuple(projected) @torch.inference_mode() diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 58c59fcbb5cb..2c86b4af801b 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -11,6 +11,7 @@ from transformers import Qwen3Config from vllm import _custom_ops as ops +from vllm import envs from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.distributed import ( @@ -86,6 +87,33 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool: ) +def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: + """Return the target model's rotary-embedding layout when it is exposed. + + A DFlash-family draft must rotate query and key tensors with the same + dimension layout as the target model used during hidden-state extraction. + Draft checkpoints do not encode this property. A mismatch changes every + drafted attention result without raising an error and collapses token + acceptance. + + Args: + target_model: Target model that can expose rotary layout modules. + + Returns: + The exposed NeoX rotary-layout setting, or ``None`` when unavailable. + """ + language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + for module in language_model.modules(): + style = getattr(module, "is_neox_style", None) + if isinstance(style, bool): + return style + return None + + def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: spec_config = vllm_config.speculative_config config = spec_config.draft_model_config.hf_config @@ -160,6 +188,32 @@ def _resolve_layer_attention( return sliding_window, _dflash_layer_causal(config, layer_idx) +def _qkv_weight_out_major(qkv_proj: nn.Module) -> torch.Tensor: + """Return the QKV projection weight as ``[out_features, in_features]``. + + Online fp8 quantization (``Fp8PtpcOnlineLinearMethod``) replaces the BF16 + ``[out, in]`` parameter with a transposed float8 ``[in, out]`` tensor and + a ``[out, 1]`` per-channel scale. The fused context K/V projection reads + the weight directly, so dequantize it back to the model dtype here; the + values are exactly the ones the quantized query path multiplies with. + """ + weight = qkv_proj.weight + if weight.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2): + return weight + scale = getattr(qkv_proj, "weight_scale", None) + if scale is None: + raise RuntimeError("fp8 qkv_proj has no weight_scale to dequantize with") + dtype = ( + qkv_proj.params_dtype + if hasattr(qkv_proj, "params_dtype") + else torch.bfloat16 + ) + if scale.dim() == 2 and scale.shape[0] == weight.shape[1]: + # Transposed [in, out] fp8 with [out, 1] scales. + return (weight.to(dtype) * scale.to(dtype).t()).t().contiguous() + return (weight.to(dtype) * scale.to(dtype)).contiguous() + + class DFlashAttention(Attention): """Attention with DFlash-specific KV allocation semantics. @@ -216,6 +270,7 @@ def __init__( add_swa_attention_sink_bias: bool = False, sliding_window: int | None = None, causal: bool = False, + is_neox_style: bool = True, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -259,6 +314,7 @@ def __init__( self.rotary_emb = get_rope( self.head_dim, max_position=max_position, + is_neox_style=is_neox_style, rope_parameters=rope_parameters, ) @@ -342,6 +398,12 @@ def __init__( # non-causal) from the draft config. sliding_window, causal = _resolve_layer_attention(config, layer_idx) + # The loader copies this value from the built target model. Kimi-K3 + # uses interleaved rotary dimensions while Qwen3 defaults to NeoX + # rotary dimensions, so relying on the Qwen3 default is not valid for + # a Kimi-trained DFlash-family checkpoint. + is_neox_style = getattr(config, "is_neox_style", True) + self.self_attn = DFlashQwen3Attention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, @@ -352,6 +414,7 @@ def __init__( add_swa_attention_sink_bias=add_swa_attention_sink_bias, sliding_window=sliding_window, causal=causal, + is_neox_style=is_neox_style, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, @@ -508,7 +571,9 @@ def _build_context_kv_buffers( self._hidden_norm_weight = self.hidden_norm.weight.data # KV projection weights: [num_layers * 2 * kv_size, hidden_size] - kv_weights = [a.qkv_proj.weight[a.q_size :] for a in layers_attn] + kv_weights = [ + _qkv_weight_out_major(a.qkv_proj)[a.q_size :] for a in layers_attn + ] self._fused_kv_weight = torch.cat(kv_weights, dim=0) if has_bias: kv_biases = [a.qkv_proj.bias[a.q_size :] for a in layers_attn] @@ -770,6 +835,40 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ) else: self.draft_id_to_target_id = None + # Rowwise-fp8 copy of the (possibly shared) LM head, materialized by + # maybe_init_fp8_draft_head(); None keeps the BF16 head. + self._fp8_draft_head = None + self._logit_scale = float(logit_scale) + + def maybe_init_fp8_draft_head(self) -> None: + """Materialize the rowwise-fp8 draft lm_head copy (opt-in). + + Called by ``load_dflash_model`` after the target's lm_head may have + been aliased onto this model, and before the proposal CUDA graphs + are captured: the quantized copy must exist when the graph records + ``compute_logits``. Draft-time only: proposals are scored with the + fp8 head, the target's verification never sees it, so accepted + tokens keep the target distribution; a rare argmax flip costs one + rejected draft token. + """ + from vllm.model_executor.layers.fp8_draft_head import ( + fp8_draft_head_supported, + quantize_draft_head, + ) + + if not envs.VLLM_DSPARK_FP8_DRAFT_HEAD: + return + if not fp8_draft_head_supported(self.lm_head.weight.device): + logger.warning( + "VLLM_DSPARK_FP8_DRAFT_HEAD is set but this device has no " + "fp8 support (SM89+ required); using the BF16 draft lm_head." + ) + return + self._fp8_draft_head = quantize_draft_head(self.lm_head.weight) + logger.info_once( + "DFlash draft logits use a rowwise-fp8 copy of the lm_head " + "(draft-time only; the target's verify pass is untouched)." + ) def embed_input_ids( self, @@ -799,7 +898,21 @@ def compute_logits( self, hidden_states: torch.Tensor, ) -> torch.Tensor | None: - logits = self.logits_processor(self.lm_head, hidden_states) + if self._fp8_draft_head is not None: + from vllm.model_executor.layers.fp8_draft_head import ( + fp8_draft_head_logits, + ) + + # Mirrors LogitsProcessor._get_logits: local (shard) logits, the + # same TP gather and vocab-padding slice, then the logit scale. + local_logits = fp8_draft_head_logits(hidden_states, self._fp8_draft_head) + logits = self.logits_processor._gather_logits(local_logits) + if logits is not None: + logits = logits[..., : self.logits_processor.org_vocab_size] + if self._logit_scale != 1.0: + logits = logits * self._logit_scale + else: + logits = self.logits_processor(self.lm_head, hidden_states) if self.draft_id_to_target_id is None: return logits diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index ff62b34ec787..8598ac4c6386 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -17,6 +17,7 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, + tensor_model_parallel_all_gatherv, ) from vllm.logger import init_logger from vllm.platforms import current_platform @@ -442,9 +443,8 @@ def run_dp_sharded_mrope_vision_model( # Get load balancing assignment with all metadata # image_to_tp_rank = [0, 2, 1, 3] # gpu_sample_counts = [1, 3] - # grouped_pixel_values_len = [1000, 350] - (image_to_tp_rank, gpu_sample_counts, grouped_pixel_values_len) = ( - get_load_balance_assignment(patches_per_image, tp_size) + (image_to_tp_rank, gpu_sample_counts, _) = get_load_balance_assignment( + patches_per_image, tp_size ) # cu_gpu_sample_counts = [0, 1, 4] @@ -481,11 +481,23 @@ def run_dp_sharded_mrope_vision_model( vision_model.spatial_merge_size * vision_model.spatial_merge_size ) - # Find the max length across all ranks - # The output embedding of every DP rank has to be - # padded to this length for tensor_model_parallel_all_gather - # to work - max_len_per_rank = max(grouped_pixel_values_len) // embed_dim_reduction_factor + patches_per_output_image = [ + patch_size // embed_dim_reduction_factor for patch_size in patches_per_image + ] + output_lengths_per_rank = [] + rank_image_offset = 0 + for count in gpu_sample_counts: + rank_image_indices = image_to_tp_rank[ + rank_image_offset : rank_image_offset + count + ] + output_lengths_per_rank.append( + sum(patches_per_output_image[i] for i in rank_image_indices) + ) + rank_image_offset += count + + if not grid_thw_list: + return () + local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local] # Run the vision model on the local pixel_values_local @@ -514,46 +526,21 @@ def run_dp_sharded_mrope_vision_model( dtype=pixel_values.dtype, ) - # Pad the output based on max_len_per_rank - # for tensor_model_parallel_all_gather to work - current_len = image_embeds_local.shape[0] - if current_len < max_len_per_rank: - padding_size = max_len_per_rank - current_len - if rope_type == "rope_2d": - padding = torch.empty( - ( - padding_size, - image_embeds_local.shape[1], - image_embeds_local.shape[2], - ), - dtype=image_embeds_local.dtype, - device=image_embeds_local.device, - ) - else: - padding = torch.empty( - (padding_size, image_embeds_local.shape[1]), - dtype=image_embeds_local.dtype, - device=image_embeds_local.device, - ) - image_embeds_local_padded = torch.cat([image_embeds_local, padding], dim=0) - else: - image_embeds_local_padded = image_embeds_local - - # Do all_gather to collect embeddings from all ranks - gathered_embeds = tensor_model_parallel_all_gather(image_embeds_local_padded, dim=0) - - # Remove padding and reconstruct per-rank embeddings - rank_embeddings = list[torch.Tensor]() - for rank in range(tp_size): - start_idx = rank * max_len_per_rank - end_idx = start_idx + ( - grouped_pixel_values_len[rank] // embed_dim_reduction_factor + expected_local_len = output_lengths_per_rank[tp_rank_local] + if image_embeds_local.shape[0] != expected_local_len: + raise ValueError( + "Vision encoder output length does not match the image-grid metadata: " + f"rank {tp_rank_local} produced {image_embeds_local.shape[0]} rows, " + f"expected {expected_local_len}" ) - rank_embeddings.append(gathered_embeds[start_idx:end_idx]) - patches_per_output_image = [ - (patch_size // embed_dim_reduction_factor) for patch_size in patches_per_image - ] + # Gather only the rows produced by each rank. Padding every rank to the + # largest shard can multiply the transient allocation by the TP size when + # a request contains fewer images than tensor-parallel ranks. + gathered_embeds = tensor_model_parallel_all_gatherv( + image_embeds_local, sizes=output_lengths_per_rank, dim=0 + ) + rank_embeddings = list(gathered_embeds.split(output_lengths_per_rank, dim=0)) # Reconstruct embeddings in the original order original_order_embeddings = [None] * len(grid_thw_list) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 8e107646ad66..766bd1570bfc 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -151,6 +151,24 @@ def _restore_merged_output_order( ) +def _reuse_consumed_query_for_context_output( + query: torch.Tensor, + output: torch.Tensor, +) -> torch.Tensor: + """Return contiguous semantic-output storage backed by a consumed query.""" + if not query.is_contiguous(): + raise ValueError("Kimi-K3 MLA prefill query storage must be contiguous") + required_bytes = output.numel() * output.element_size() + query_bytes = query.view(torch.uint8).flatten() + if query_bytes.numel() < required_bytes: + raise ValueError( + "Kimi-K3 MLA prefill query storage is too small for compact context " + f"output: available={query_bytes.numel()} bytes, " + f"required={required_bytes} bytes" + ) + return query_bytes[:required_bytes].view(output.dtype).view_as(output) + + class KimiShardedMergedColumnParallelLinear(MergedColumnParallelLinear): """Merged column projection with one gather and logical-shard reorder.""" @@ -1066,6 +1084,16 @@ def _forward_prefill_fused( ) if has_context: + suffix_output, suffix_lse = output_prefill + out = out.view(-1, self.num_local_heads, self.v_head_dim) + # FlashAttention 2 pads Kimi-K3's 128-wide V to the 256-wide + # query/key head dimension. Preserve only the semantic V slice in + # caller-owned output storage before context attention allocates + # its equally large padded result. The merge kernel supports + # output aliasing its suffix input, so both padded results never + # need to be live at the same time. + out.copy_(suffix_output[..., : self.v_head_dim]) + del output_prefill, suffix_output if self.dcp_world_size > 1: context_output, context_lse = ( self.impl._context_parallel_compute_prefill_context( # type: ignore[attr-defined] @@ -1080,13 +1108,14 @@ def _forward_prefill_fused( context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] q, self._attn_read_kv_cache(), attn_metadata, self._k_scale ) - suffix_output, suffix_lse = output_prefill - out = out.view(-1, self.num_local_heads, self.v_head_dim) + compact_context_output = _reuse_consumed_query_for_context_output(q, out) + compact_context_output.copy_(context_output[..., : self.v_head_dim]) + del context_output merge_attn_states( output=out, - prefix_output=context_output[..., : self.v_head_dim], + prefix_output=compact_context_output, prefix_lse=context_lse, - suffix_output=suffix_output[..., : self.v_head_dim], + suffix_output=out, suffix_lse=suffix_lse, ) elif not writes_out: diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index bbceb40cec19..30ed84ffc256 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1579,6 +1579,11 @@ def __init__( self.attn_res_block_size = attn_res_block_size self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 self.block_write_idx = layer_idx // self.attn_res_block_size + self.is_final_block_write_layer = ( + self.is_block_write_layer + and self.block_write_idx + == cdiv(config.num_hidden_layers, self.attn_res_block_size) - 1 + ) self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) self.self_attention_res_norm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -1694,7 +1699,14 @@ def _post_attn_norm( assert prefix_sum is not None if self.is_block_write_layer: - output = prefix_sum if self.reuse_attn_res_output else None + # The old prefix becomes the last committed residual block at the + # final block boundary. It must remain immutable for every later + # AttnRes mixture and therefore cannot also hold the new delta. + output = ( + prefix_sum + if self.reuse_attn_res_output and not self.is_final_block_write_layer + else None + ) prefix_sum = hidden_states prefix_delta = None else: @@ -1891,6 +1903,94 @@ def make_empty_intermediate_tensors( } ) + def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + super()._set_aux_hidden_state_layers(layers) + if self.use_attn_res: + # Emitted once, at configuration time. Which layers are tapped and + # which convention is in force are the two things you need to + # confirm from a running process, and neither is recoverable from + # the served output. + logger.info_once( + "Kimi-K3 aux hidden capture: layers=%s mode=%s " + "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)", + layers, + "attn_res_stream" if self._aux_attn_res_stream else "prefix_only", + int(self._aux_attn_res_stream), + ) + + @property + def _aux_attn_res_stream(self) -> bool: + return envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM + + def _capture_aux_hidden_stream( + self, + layer_idx: int, + prefix_sum: torch.Tensor, + pending_mlp_out: torch.Tensor | None, + block_residual: torch.Tensor, + ) -> torch.Tensor: + """Auxiliary feature tapped after ``layer_idx`` under AttnRes. + + The wire between layers only carries the current block's running prefix; + the committed blocks live in the bank. The value the next consumer + actually reads is the pre-norm AttnRes mixture over + ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were + trained against. ``attn_res`` with no delta, no block write and no + output norm computes exactly that and leaves both the prefix and the + bank untouched. + + Folding the pending MLP output into the prefix rather than passing it as + ``delta`` is deliberate: the kernel writes an applied delta back into + the prefix in place, which would double-add it into the live residual + stream. + + Args: + layer_idx: Index of the layer that produced the pending MLP output. + prefix_sum: Running prefix for the active AttnRes block. + pending_mlp_out: MLP output to fold into the running prefix, if any. + block_residual: Committed AttnRes block bank for the active rows. + + Returns: + Auxiliary hidden states for the configured DFlash capture mode. + """ + prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out + # `use_attn_res` is what constructs the norm and projection weights this + # reads; without it there is no mixture to compute and the attribute + # lookups below would raise. + if not (self._aux_attn_res_stream and self.use_attn_res): + return prefix + + if layer_idx + 1 < self.end_layer: + consumer = self.layers[layer_idx + 1] + score_norm = consumer.self_attention_res_norm + score_proj = consumer.self_attention_res_proj + num_blocks = consumer.prev_valid_blocks + elif get_pp_group().is_last_rank: + # Nothing downstream but the model's own output-side aggregation. + score_norm = self.output_attn_res_norm + score_proj = self.output_attn_res_proj + num_blocks = self.num_attn_res_blocks + else: + # Last layer of a non-final pipeline stage: the consumer lives on + # the next rank and the output-side aggregation only exists on the + # last one, so there is nothing here to mix against. Falling back + # to the running prefix keeps the tap defined rather than reaching + # for weights this rank does not construct. + return prefix + + return attn_res( + prefix, + None, + block_residual, + score_norm.weight, + score_proj.weight.squeeze(0), + None, + num_blocks=num_blocks, + block_write_idx=-1, + eps=score_norm.variance_epsilon, + output_norm_eps=0.0, + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -1993,6 +2093,9 @@ def forward( pp_group = get_pp_group() stream_aux_hidden_states = bool( projector is not None + # The streaming projector consumes plain residual sums. DFlash + # AttnRes capture requires the pre-norm mixture computed below. + and not self._aux_attn_res_stream and not self.use_sequence_parallel and pp_group.is_first_rank and pp_group.is_last_rank @@ -2054,7 +2157,10 @@ def forward( projector.accumulate_auxiliary_state(hidden_states, residual) elif self.use_attn_res: assert prefix_sum is not None - aux_hidden_state = prefix_sum + hidden_states + assert residual is not None + aux_hidden_state = self._capture_aux_hidden_stream( + layer_idx, prefix_sum, hidden_states, residual + ) aux_hidden_states.append(aux_hidden_state) else: assert residual is not None diff --git a/vllm/parser/kimi_k3.py b/vllm/parser/kimi_k3.py index 380e2cce5f91..e3f58d25f4d7 100644 --- a/vllm/parser/kimi_k3.py +++ b/vllm/parser/kimi_k3.py @@ -12,6 +12,7 @@ ) from vllm.parser.abstract_parser import DelegatingParser from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser +from vllm.tool_parsers.utils import partial_tag_overlap if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -23,6 +24,45 @@ class KimiK3Parser(DelegatingParser): """Compose the Kimi K3 reasoning and tool parsers for XTML output.""" + _CONTENT_PROTOCOL_MARKERS = ( + "<|open|>think<|sep|>", + "<|close|>think<|sep|>", + "<|open|>response<|sep|>", + "<|close|>response<|sep|>", + "<|close|>message<|sep|>", + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._pending_content_protocol = "" + + def _strip_content_protocol( + self, content: str | None, *, finished: bool + ) -> str | None: + """Remove Kimi XTML control markers from streamed API content. + + Reasoning and tool parsers normally consume these markers before this + method runs. The final filter preserves the API invariant when a + malformed model transition or an inconsistent initial parser phase + routes a marker through the content field. Marker prefixes are held + across stream chunks so partial control tokens are never exposed. + """ + pending = self._pending_content_protocol + (content or "") + for marker in self._CONTENT_PROTOCOL_MARKERS: + pending = pending.replace(marker, "") + + overlap = max( + partial_tag_overlap(pending, marker) + for marker in self._CONTENT_PROTOCOL_MARKERS + ) + if overlap: + emitted = pending[:-overlap] + self._pending_content_protocol = "" if finished else pending[-overlap:] + else: + emitted = pending + self._pending_content_protocol = "" + return emitted or None + # TODO: Switch Kimi K3 to the parser engine once its XTML reasoning/tool # path is covered there. def _extract_tool_calls( @@ -113,18 +153,25 @@ def parse_delta( ) if ( - self._tool_parser is not None - or not isinstance(self._reasoning_parser, KimiK3ReasoningParser) - or not state.reasoning_ended - or delta_message is None + delta_message is not None + and self._tool_parser is None + and isinstance(self._reasoning_parser, KimiK3ReasoningParser) + and state.reasoning_ended ): - return delta_message + stripped = self._reasoning_parser.strip_content_streaming( + previous_text=previous_content, + current_text=state.previous_text, + ) + delta_message.content = stripped.content if stripped is not None else None + + if delta_message is None: + if finished: + self._pending_content_protocol = "" + return None - stripped = self._reasoning_parser.strip_content_streaming( - previous_text=previous_content, - current_text=state.previous_text, + delta_message.content = self._strip_content_protocol( + delta_message.content, finished=finished ) - delta_message.content = stripped.content if stripped is not None else None if ( delta_message.role is None and delta_message.content is None diff --git a/vllm/reasoning/kimi_k3_reasoning_parser.py b/vllm/reasoning/kimi_k3_reasoning_parser.py index f05d5c044805..d170652a20e8 100644 --- a/vllm/reasoning/kimi_k3_reasoning_parser.py +++ b/vllm/reasoning/kimi_k3_reasoning_parser.py @@ -95,6 +95,9 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): if thinking is None: thinking = chat_kwargs.get("enable_thinking", True) self._thinking_enabled = bool(thinking) + self._starts_new_assistant_message = bool( + chat_kwargs.get("add_generation_prompt", True) + ) and not bool(chat_kwargs.get("continue_final_message", False)) # XTML markers as literal strings (skip_special_tokens=False at serve time) self._think_open = "<|open|>think<|sep|>" @@ -145,6 +148,11 @@ def reasoning_start_str(self) -> str | None: def reasoning_end_str(self) -> str | None: return self._think_close + @property + def thinking_enabled(self) -> bool: + """Whether the request opens a Kimi K3 reasoning channel.""" + return self._thinking_enabled + def adjust_request( self, request: "ChatCompletionRequest | ResponsesRequest", @@ -170,6 +178,24 @@ def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: _newest_marker(input_ids, self._think_close_ids, self._think_open_ids) == 0 ) + def is_reasoning_end_for_prompt(self, input_ids: Sequence[int]) -> bool: + """Return the reasoning phase at the generation boundary. + + A fresh Kimi K3 assistant message starts in the think channel whenever + thinking is enabled. That request contract is authoritative even when + the token list exposed to the output parser does not contain the + generation-prefix marker; closed think channels from historical turns + cannot determine the phase of a fresh assistant message. + + ``continue_final_message`` does not open a fresh channel, so its prompt + markers remain authoritative. + """ + if not self._thinking_enabled: + return True + if self._starts_new_assistant_message: + return False + return self.is_reasoning_end(input_ids) + def is_reasoning_end_streaming( self, input_ids: Sequence[int], delta_ids: Iterable[int] ) -> bool: diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py index 9f688f2270e4..e1fc5a156996 100644 --- a/vllm/tool_parsers/kimi_k3_tool_parser.py +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -37,6 +37,7 @@ import regex as re from openai.types.responses import ToolChoiceFunction +from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, @@ -69,6 +70,25 @@ def _partial_tag_overlap(text: str, tag: str) -> int: return 0 +def _partial_pattern_overlap(text: str, pattern: re.Pattern) -> int: + """Return a trailing prefix accepted by a partial regular expression. + + Args: + text: Generated text that may end inside a structural marker. + pattern: Complete structural-marker expression. + + Returns: + The number of trailing characters that form an incomplete match. + """ + candidate = text.rfind("<") + while candidate >= 0: + match = pattern.fullmatch(text[candidate:], partial=True) + if match is not None and match.partial: + return len(text) - candidate + candidate = text.rfind("<", 0, candidate) + return 0 + + class KimiK3ToolParser(ToolParser): supports_required_and_named = False # Enables the vLLM-side XTML structural tag builder @@ -86,6 +106,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): self.tools_close = "<|close|>tools<|sep|>" self.response_open = "<|open|>response<|sep|>" self.response_close = "<|close|>response<|sep|>" + self.argument_close = "<|close|>argument<|sep|>" # Regexes operate on detokenized text. The XTML markers reach us as the # literal strings <|open|>/<|close|>/<|sep|>. adjust_request keeps them @@ -127,6 +148,16 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): + _S, re.DOTALL, ) + # Half-open variants: streaming needs to see a call/argument as soon as + # its opening marker lands, before the matching close marker exists. + self._call_open_re = re.compile( + _O + r"\s*call\s+(?P" + _TEXT_UNTIL_SEP + r")" + _S + ) + self._call_close_re = re.compile(_C + r"\s*call\s*" + _S) + self._arg_open_re = re.compile( + _O + r"\s*argument\s+(?P" + _TEXT_UNTIL_SEP + r")" + _S + ) + self._arg_close_re = re.compile(_C + r"\s*argument\s*" + _S) # attr segment: key="value" (value already escaped on the encode side) self._attr_re = re.compile(r'(?P\w+)="(?P[^"]*)"') self._response_re = re.compile( @@ -136,7 +167,8 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): # streaming state self._sent_content_idx = 0 - self._sent_tool_call_count = 0 + # arguments already streamed, per tool-call index + self._streamed_args: list[str] = [] if not self.model_tokenizer: raise ValueError( @@ -218,6 +250,17 @@ def _decode_call(self, attrs: str, body: str) -> ToolCall | None: """ call_attrs = self._attrs(attrs) tool_name = call_attrs.get("tool", "") + if not tool_name: + return None + return ToolCall( + type="function", + function=FunctionCall( + name=tool_name, + arguments=json.dumps(self._decode_arguments(body), ensure_ascii=False), + ), + ) + + def _decode_arguments(self, body: str) -> dict: arguments: dict = {} for arg_match in self._arg_re.finditer(body): arg_attrs = self._attrs(arg_match["attrs"]) @@ -231,15 +274,37 @@ def _decode_call(self, attrs: str, body: str) -> ToolCall | None: arguments[key] = json.loads(raw_value) except json.JSONDecodeError: arguments[key] = raw_value - if not tool_name: - return None - return ToolCall( - type="function", - function=FunctionCall( - name=tool_name, - arguments=json.dumps(arguments, ensure_ascii=False), - ), - ) + return arguments + + def _partial_arguments(self, body: str) -> str: + """Serialize the arguments of a ``call`` block that has not closed yet. + + Args: + body: The incomplete XTML call body. + + Returns: + A prefix of the completed call's JSON arguments. String values can + stream as generated. Other argument types remain buffered until + their complete JSON literal is available. + """ + closed_end = 0 + for arg_match in self._arg_re.finditer(body): + closed_end = arg_match.end() + arguments = self._decode_arguments(body[:closed_end]) + + m_open = self._arg_open_re.search(body, closed_end) + if m_open is None: + # drop the closing "}" of the object + return json.dumps(arguments, ensure_ascii=False)[:-1] + + arg_attrs = self._attrs(m_open["attrs"]) + if arg_attrs.get("type", "string") != "string": + return json.dumps(arguments, ensure_ascii=False)[:-1] + raw_value = body[m_open.end() :] + held = _partial_pattern_overlap(raw_value, self._arg_close_re) + arguments[arg_attrs.get("key", "")] = raw_value[: len(raw_value) - held] + # drop the closing quote of the open value and the "}" of the object + return json.dumps(arguments, ensure_ascii=False)[:-2] def _strip_response_content(self, text: str) -> str | None: """Strip XTML response/message markers from generated response text. @@ -361,35 +426,66 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - # Conservative streaming: stream unwrapped response-channel text, then - # buffer tool calls and emit each call once its block closes. + # Stream unwrapped response-channel text, then stream each tool call as + # it is generated: the open marker carries the name, and argument text + # is forwarded as it arrives rather than buffered until the call closes. content = self._extract_response_content(current_text) - # tools channel is open: parse fully-closed calls we have not emitted yet m_tools = self._tools_open_re.search(current_text) if m_tools is None: return DeltaMessage(content=content) if content else None - section = current_text[m_tools.end() :] - calls = [ - tc - for m in self._call_re.finditer(section) - if (tc := self._decode_call(m["attrs"], m["body"])) is not None - ] - if len(calls) <= self._sent_tool_call_count: - return DeltaMessage(content=content) if content else None - new = calls[self._sent_tool_call_count :] - - deltas = [ - DeltaToolCall( - index=self._sent_tool_call_count + i, - id=tc.id, - type="function", - function=DeltaFunctionCall( - name=tc.function.name, arguments=tc.function.arguments - ).model_dump(exclude_none=True), + section_start = m_tools.end() + m_tools_close = self._tools_close_re.search(current_text, section_start) + section = current_text[ + section_start : ( + len(current_text) if m_tools_close is None else m_tools_close.start() ) - for i, tc in enumerate(new) ] - self._sent_tool_call_count = len(calls) + opens = list(self._call_open_re.finditer(section)) + deltas: list[DeltaToolCall] = [] + index = 0 + for i, m_open in enumerate(opens): + end = opens[i + 1].start() if i + 1 < len(opens) else len(section) + body = section[m_open.end() : end] + name = self._attrs(m_open["attrs"]).get("tool", "") + if not name: + # an empty/garbage block is dropped, as in _decode_call + continue + m_close = self._call_close_re.search(body) + if m_close is None: + arguments = self._partial_arguments(body) + else: + arguments = json.dumps( + self._decode_arguments(body[: m_close.start()]), ensure_ascii=False + ) + + if index == len(self._streamed_args): + self._streamed_args.append(arguments) + deltas.append( + DeltaToolCall( + index=index, + id=make_tool_call_id(), + type="function", + function=DeltaFunctionCall( + name=name, arguments=arguments + ).model_dump(exclude_none=True), + ) + ) + else: + sent = self._streamed_args[index] + if arguments != sent and arguments.startswith(sent): + self._streamed_args[index] = arguments + deltas.append( + DeltaToolCall( + index=index, + function=DeltaFunctionCall( + arguments=arguments[len(sent) :] + ).model_dump(exclude_none=True), + ) + ) + index += 1 + + if not deltas: + return DeltaMessage(content=content) if content else None return DeltaMessage(content=content, tool_calls=deltas) diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 97f30736fda3..a8444f7bbb7b 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -395,15 +395,20 @@ def __init__( self.max_num_splits = 0 # No upper bound on the number of splits. self.aot_schedule = get_flash_attn_version() == 3 - try: - from vllm.distributed.parallel_state import get_dcp_group - - self.dcp_world_size = get_dcp_group().world_size - self.dcp_rank = get_dcp_group().rank_in_group - except AssertionError: - # DCP might not be initialized in testing + # A DCP-replicated KV group stores the complete sequence on every + # rank. Build ordinary local-attention metadata for that group instead + # of partitioning its sequence lengths for a second time. + if getattr(kv_cache_spec, "dcp_replicated", False): self.dcp_world_size = 1 self.dcp_rank = 0 + else: + try: + self.dcp_world_size = get_dcp_group().world_size + self.dcp_rank = get_dcp_group().rank_in_group + except AssertionError: + # DCP might not be initialized in testing + self.dcp_world_size = 1 + self.dcp_rank = 0 self.cp_kv_cache_interleave_size = ( self.parallel_config.cp_kv_cache_interleave_size diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index e0d3c2a1a6cc..da11eb1dddfb 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -295,3 +295,7 @@ class GrammarOutput: structured_output_request_ids: list[str] # Bitmask ordered as structured_output_request_ids. grammar_bitmask: "npt.NDArray[np.int32]" + # Number of speculative rows represented in grammar_bitmask for each + # structured output request. Worker-side draft trimming may reduce the + # number of logits without changing this compact source layout. + num_spec_tokens: list[int] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ff576e2bc5e5..e7596ebdd955 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,6 +70,31 @@ logger = init_logger(__name__) +def use_eagle_for_target_cache( + speculative_config: Any | None, + kv_cache_groups: Iterable[Any], +) -> bool: + """Whether target KV groups require EAGLE's last-hash drop. + + DSpark and DFlash use EAGLE-style scheduling/lookahead but can keep their + draft KV outside the target cache groups. In that layout, falling back from + an empty ``is_eagle_group`` set to all target groups drops one valid target + prefix hash. Classic EAGLE keeps the historical unannotated fallback. + + Args: + speculative_config: Active speculative-decoding configuration, if any. + kv_cache_groups: Target-cache groups inspected for explicit EAGLE state. + + Returns: + Whether target-cache prefix matching must drop EAGLE's trailing hash. + """ + if speculative_config is None or not speculative_config.use_eagle(): + return False + if any(group.is_eagle_group for group in kv_cache_groups): + return True + return speculative_config.method not in ("dspark", "dflash") + + class Scheduler(SchedulerInterface): def __init__( self, @@ -270,6 +295,11 @@ def __init__( ) self.use_eagle = speculative_config.use_eagle() + self.use_eagle_for_target_cache = use_eagle_for_target_cache( + speculative_config, + kv_cache_config.kv_cache_groups, + ) + # Create the KV cache manager. if hash_block_size is None: hash_block_size = block_size @@ -279,7 +309,7 @@ def __init__( max_model_len=self.max_model_len, max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, - use_eagle=self.use_eagle, + use_eagle=self.use_eagle_for_target_cache, log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, @@ -406,7 +436,7 @@ def _mamba_block_aligned_split( # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle: + if getattr(self, "use_eagle_for_target_cache", self.use_eagle): # EAGLE drops the last complete draft-attention block. Convert the # resulting maximum reusable prefix to the recurrent-state grid. # Older configurations without an annotated EAGLE group retain @@ -1262,6 +1292,30 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: scheduled_encoder_inputs ) + # Skip speculative decoding when every scheduled request needs at most + # one output token. Drafting cannot pay off for a 1-token output, so the + # draft pass plus verification is pure added latency. Only fires when no + # running request is scheduled, so an in-flight multi-token generation + # can never lose its draft tokens. + if ( + num_spec_tokens_to_schedule > 0 + and scheduled_new_reqs + and not scheduled_running_reqs + and all(req.max_tokens <= 1 for req in scheduled_new_reqs) + ): + num_spec_tokens_to_schedule = 0 + # NOTE: allocate_slots (called above for each request) already + # reserved KV blocks for self.num_lookahead_tokens speculative + # slots. We cannot retroactively shrink that reservation here + # because it was made per-request during the scheduling loop, + # before this aggregate skip condition is evaluated. The + # reservation is harmless: these are single-token requests + # (max_tokens <= 1) that finish immediately after one decode + # step, so the extra blocks are released at the next scheduling + # iteration. Zeroing num_spec_tokens_to_schedule prevents the + # draft model from running, which is the source of the wasted + # latency we are avoiding. + scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, @@ -1723,7 +1777,15 @@ def get_grammar_bitmask( structured_output_request_ids, scheduler_output.scheduled_spec_decode_tokens, ) - return GrammarOutput(structured_output_request_ids, bitmask) + num_spec_tokens = [ + len(scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())) + for req_id in structured_output_request_ids + ] + return GrammarOutput( + structured_output_request_ids, + bitmask, + num_spec_tokens, + ) def update_from_output( self, @@ -1835,9 +1897,13 @@ def update_from_output( scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) + observed_spec_decode = False + num_draft_tokens = 0 + num_accepted = 0 if scheduled_spec_token_ids and ( generated_token_ids or self.num_sampled_tokens_per_step == 0 ): + observed_spec_decode = True num_draft_tokens = len(scheduled_spec_token_ids) num_sampled = self.num_sampled_tokens_per_step num_accepted = max(len(generated_token_ids) - num_sampled, 0) @@ -1855,13 +1921,6 @@ def update_from_output( request.num_computed_tokens -= num_rejected if request.num_output_placeholders > 0: request.num_output_placeholders -= num_rejected - spec_decoding_stats = self.make_spec_decoding_stats( - spec_decoding_stats, - num_draft_tokens=num_draft_tokens, - num_accepted_tokens=num_accepted, - num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, - request_id=req_id, - ) # Free encoder inputs only after the step has actually executed. if request.has_encoder_inputs: @@ -1877,6 +1936,40 @@ def update_from_output( status_before_stop = request.status num_output_tokens_before = len(request._output_token_ids) + if ( + len(new_token_ids) > 1 + and scheduled_spec_token_ids + and request.use_structured_output + and not output_is_stale + ): + new_token_ids, num_grammar_rejected = ( + self.structured_output_manager.filter_speculative_grammar_tokens( + request, new_token_ids + ) + ) + if num_grammar_rejected > 0: + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_grammar_rejected + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_grammar_rejected + # Target-sampled tokens occupy the end of the output block. + # Removing that tail changes the accepted-draft count only + # when the rejected suffix extends into draft positions. + num_accepted -= max( + num_grammar_rejected - self.num_sampled_tokens_per_step, + 0, + ) + assert num_accepted >= 0 + + if observed_spec_decode: + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, + ) + # Check for stop and update request status. if new_token_ids: new_token_ids, stopped = self._update_request_with_output( @@ -2946,50 +3039,73 @@ def _update_requests_with_invalid_blocks( is_affected = False marked_invalid_block = False req_id = request.request_id - # TODO (davidb): add support for hybrid memory allocator - (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) - # We iterate only over blocks that may contain externally computed - # tokens + req_block_ids_by_group = self.kv_cache_manager.get_block_ids(req_id) + # We iterate only over blocks that may contain externally + # computed tokens. req_num_computed_tokens = ( request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) ) + computed_block_ids_by_group = ( + self.kv_cache_manager.get_block_ids_for_computed_tokens( + req_id, + req_num_computed_tokens, + ) + ) - req_num_computed_blocks = ( - req_num_computed_tokens + self.block_size - 1 - ) // self.block_size - for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids): - if block_id not in invalid_block_ids: - continue - - is_affected = True - - if block_id in marked_invalid_block_ids: - # This invalid block is shared with a previous request - # and was already marked for recomputation. - # This means this request can still consider this block - # as computed when rescheduled. - # Currently this only applies to sync loading; Async - # loading does not yet support block sharing - continue - - marked_invalid_block_ids.add(block_id) + # Map each invalid block to the earliest scheduler-aligned token + # boundary from which this request must be recomputed. + invalid_block_boundaries: dict[int, int] = {} + for group, group_block_ids in zip( + self.kv_cache_config.kv_cache_groups, + computed_block_ids_by_group, + strict=True, + ): + group_block_size = group.kv_cache_spec.block_size + for block_idx, block_id in enumerate(group_block_ids): + if block_id not in invalid_block_ids: + continue - if marked_invalid_block: - # This request has already marked an invalid block for - # recomputation and updated its num_computed_tokens. - continue + block_start = block_idx * group_block_size + recompute_from = block_start // self.block_size * self.block_size + previous_boundary = invalid_block_boundaries.get(block_id) + invalid_block_boundaries[block_id] = ( + recompute_from + if previous_boundary is None + else min(previous_boundary, recompute_from) + ) - marked_invalid_block = True - # Truncate the computed tokens at the first failed block - request.num_computed_tokens = idx * self.block_size - num_affected_tokens = ( - req_num_computed_tokens - request.num_computed_tokens + if invalid_block_boundaries: + is_affected = True + new_invalid_block_ids = ( + invalid_block_boundaries.keys() - marked_invalid_block_ids ) - total_affected_tokens += num_affected_tokens + marked_invalid_block_ids.update(new_invalid_block_ids) - # collect invalid block and all downstream dependent blocks - if evict_blocks: - blocks_to_evict.update(req_block_ids[idx:]) + if new_invalid_block_ids: + marked_invalid_block = True + request.num_computed_tokens = min( + invalid_block_boundaries[block_id] + for block_id in new_invalid_block_ids + ) + num_affected_tokens = ( + req_num_computed_tokens - request.num_computed_tokens + ) + total_affected_tokens += num_affected_tokens + + # Every KV group after the common recomputation boundary + # depends on the failed prefix, so collect downstream + # blocks from all groups. + if evict_blocks: + for group, group_block_ids in zip( + self.kv_cache_config.kv_cache_groups, + req_block_ids_by_group, + strict=True, + ): + group_block_size = group.kv_cache_spec.block_size + first_block_idx = ( + request.num_computed_tokens // group_block_size + ) + blocks_to_evict.update(group_block_ids[first_block_idx:]) if is_affected: if not marked_invalid_block: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index a334a8dde9da..686cae8257e6 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1597,6 +1597,7 @@ def allocate_new_blocks( # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. if num_required_blocks <= len(req_blocks) and not has_partial_hit: + self._allocated_block_reqs.add(request_id) return [] else: prev_block_len = len(req_blocks) diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 14f238b4c735..d3ebdb6fe3ec 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -148,6 +148,21 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: """ return cdiv(max_len, self.block_size) + def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int: + """Return the number of unique token-position shards under DCP. + + Cache types that store recurrent or otherwise rank-local state do not + shard that state by token position. Attention cache specifications + override this method because their default layout is DCP-sharded. + """ + configured_dcp = int(dcp_world_size) + if configured_dcp < 1: + raise ValueError( + "Configured decode-context-parallel size must be positive: " + f"{configured_dcp}" + ) + return 1 + def copy_with_new_block_size(self, block_size: int) -> Self: """ Create a new KVCacheSpec from self but replacing the block size. @@ -240,11 +255,38 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - kv_shard_count = get_kv_cache_dcp_shard_count( - self, parallel_config.decode_context_parallel_size + kv_shard_count = self.get_num_dcp_kv_shards( + parallel_config.decode_context_parallel_size ) return cdiv(max_len, self.block_size * kv_shard_count) + def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int: + """Return the configured or explicitly overridden attention shard count.""" + configured_dcp = int(dcp_world_size) + if configured_dcp < 1: + raise ValueError( + "Configured decode-context-parallel size must be positive: " + f"{configured_dcp}" + ) + replicated = bool(getattr(self, "dcp_replicated", False)) + override = getattr(self, "dcp_kv_shard_count", None) + if replicated: + if override not in (None, 1): + raise ValueError( + "dcp_replicated cannot be combined with " + f"dcp_kv_shard_count={override}" + ) + return 1 + if override is None: + return configured_dcp + override = int(override) + if override < 1 or override > configured_dcp or configured_dcp % override != 0: + raise ValueError( + "dcp_kv_shard_count must be a positive divisor of the configured " + f"DCP size, got shards={override}, DCP={configured_dcp}" + ) + return override + @dataclass(frozen=True, kw_only=True) class FullAttentionSpec(AttentionSpec): @@ -398,36 +440,24 @@ def get_kv_cache_dcp_shard_count( dcp_world_size: int, ) -> int: """Return the number of unique DCP token-position shards for a cache group.""" - configured_dcp = int(dcp_world_size) - if configured_dcp < 1: - raise ValueError( - f"Configured decode-context-parallel size must be positive: " - f"{configured_dcp}" - ) - replicated = bool(getattr(spec, "dcp_replicated", False)) - override = getattr(spec, "dcp_kv_shard_count", None) - if replicated: - if override not in (None, 1): - raise ValueError( - f"dcp_replicated cannot be combined with dcp_kv_shard_count={override}" - ) - return 1 - if override is None: - return configured_dcp - override = int(override) - if override < 1 or override > configured_dcp or configured_dcp % override != 0: - raise ValueError( - "dcp_kv_shard_count must be a positive divisor of the configured " - f"DCP size, got shards={override}, DCP={configured_dcp}" - ) - return override + return spec.get_num_dcp_kv_shards(dcp_world_size) def has_nondefault_kv_dcp_layout( spec: KVCacheSpec, dcp_world_size: int, ) -> bool: - return get_kv_cache_dcp_shard_count(spec, dcp_world_size) != int(dcp_world_size) + layer_specs = ( + spec.kv_cache_specs.values() + if isinstance(spec, UniformTypeKVCacheSpecs) + else (spec,) + ) + is_attention_group = all( + isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs + ) + return is_attention_group and ( + get_kv_cache_dcp_shard_count(spec, dcp_world_size) != int(dcp_world_size) + ) def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): @@ -1046,6 +1076,18 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: ) return next(iter(widths)) + def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int: + shard_counts = { + spec.get_num_dcp_kv_shards(dcp_world_size) + for spec in self.kv_cache_specs.values() + } + if len(shard_counts) != 1: + raise ValueError( + "All layers in a uniform KV cache group must use the same " + f"number of DCP KV shards, got {sorted(shard_counts)}." + ) + return next(iter(shard_counts)) + @classmethod def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: """ diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 5ba4ad5b77e7..79867280935b 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -4,7 +4,8 @@ import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING +from copy import copy +from typing import TYPE_CHECKING, overload from vllm.config import VllmConfig from vllm.logger import init_logger @@ -32,6 +33,34 @@ logger = init_logger(__name__) +class _TokenSequenceView(Sequence[int]): + """Read-only concatenation that does not copy committed token history.""" + + def __init__(self, prefix: Sequence[int], suffix: Sequence[int]) -> None: + self.prefix = prefix + self.suffix = suffix + + def __len__(self) -> int: + return len(self.prefix) + len(self.suffix) + + @overload + def __getitem__(self, index: int) -> int: ... + + @overload + def __getitem__(self, index: slice) -> list[int]: ... + + def __getitem__(self, index: int | slice) -> int | list[int]: + if isinstance(index, slice): + return [self[i] for i in range(*index.indices(len(self)))] + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + if index < len(self.prefix): + return self.prefix[index] + return self.suffix[index - len(self.prefix)] + + class StructuredOutputManager: """Engine-level manager for structured output requests.""" @@ -380,9 +409,7 @@ def should_fill_bitmask(self, request: "Request") -> bool: # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. request.structured_output_request.reasoning_ended = ( - reasoner.is_reasoning_end_for_prompt( - request.prompt_token_ids or [] - ) + reasoner.is_reasoning_end_for_prompt(request.prompt_token_ids or []) ) return request.structured_output_request.reasoning_ended return True @@ -494,6 +521,106 @@ def trim_reasoning_for_advance( return new_token_ids return new_token_ids[num_reasoning:] + @staticmethod + def _find_reasoning_end_offset( + reasoner: "ReasoningParser", + prior_token_ids: Sequence[int], + new_token_ids: list[int], + ) -> int | None: + """Return where a reasoning-end marker completes in a sampled block. + + The first parser call keeps reasoning-only blocks at one check. When + the block contains a transition, cumulative delta prefixes locate + single-token and multi-token markers, including markers that start in + ``prior_token_ids`` and finish in ``new_token_ids``. + + Args: + reasoner: Request-scoped parser that identifies the transition. + prior_token_ids: Tokens committed before the sampled block. + new_token_ids: Accepted tokens awaiting commit. + + Returns: + The zero-based offset where the marker completes, or ``None`` if + the sampled block remains inside reasoning. + """ + complete_tokens = _TokenSequenceView(prior_token_ids, new_token_ids) + if not copy(reasoner).is_reasoning_end_streaming( + complete_tokens, new_token_ids + ): + return None + + for offset in range(len(new_token_ids)): + delta_ids = new_token_ids[: offset + 1] + token_ids = _TokenSequenceView(prior_token_ids, delta_ids) + # Streaming parsers may retain request-local transition state. + # A probe must not consume that state before the normal commit path. + if copy(reasoner).is_reasoning_end_streaming(token_ids, delta_ids): + return offset + + # Some parsers report only that the complete block crossed a boundary. + # Treating the full block as reasoning avoids exposing an unknown suffix + # to the answer grammar. + return len(new_token_ids) - 1 + + def filter_speculative_grammar_tokens( + self, + request: "Request", + new_token_ids: list[int], + ) -> tuple[list[int], int]: + """Validate an accepted speculative block before it is committed. + + Grammar masks cover scheduled draft positions, but an accepted block + can contain an unconstrained token immediately after a reasoning + transition or after the grammar completes. The filter retains the + grammar-valid prefix and reports the trailing tokens that must be + resampled. ``validate_tokens`` restores the grammar state before + returning, so the scheduler's normal commit path remains the only + operation that advances the matcher. + + Args: + request: Request that owns the grammar and reasoning state. + new_token_ids: Accepted tokens awaiting scheduler commit. + + Returns: + The tokens that are safe to commit and the rejected suffix length. + """ + if self.vllm_config.speculative_config is None: + return new_token_ids, 0 + if len(new_token_ids) < 2 or not request.use_structured_output: + return new_token_ids, 0 + + structured_request = request.structured_output_request + if structured_request is None: + return new_token_ids, 0 + grammar = structured_request.grammar + if not isinstance(grammar, StructuredOutputGrammar): + return new_token_ids, 0 + + reasoner = self._get_reasoner(request) + grammar_start = 0 + if ( + reasoner is not None + and not self.enable_in_reasoning + and not structured_request.reasoning_ended + ): + boundary = self._find_reasoning_end_offset( + reasoner, + request.all_token_ids, + new_token_ids, + ) + if boundary is None: + return new_token_ids, 0 + grammar_start = boundary + 1 + + grammar_tokens = new_token_ids[grammar_start:] + if not grammar_tokens: + return new_token_ids, 0 + valid_grammar_tokens = grammar.validate_tokens(grammar_tokens) + rejected = len(grammar_tokens) - len(valid_grammar_tokens) + if rejected == 0: + return new_token_ids, 0 + return new_token_ids[:grammar_start] + valid_grammar_tokens, rejected + def clear_backend(self) -> None: if self.backend is not None: self.backend.destroy() diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 58b726bf72a4..29659c4e3ba3 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -156,11 +156,12 @@ class XgrammarGrammar(StructuredOutputGrammar): def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """Accepts a list of tokens and advances the FSM. - Returns True if the FSM was advanced successfully. - Returns False if the FSM failed to advance. + Returns True if all grammar-constrained tokens were accepted. + Tokens after termination are ignored. Returns False if the FSM + failed to advance. """ if self._is_terminated: - return False + return True for token in tokens: if not self.matcher.accept_token(token): logger.error( @@ -171,7 +172,9 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: ) return False self.num_processed_tokens += 1 - self._is_terminated = self.matcher.is_terminated() + self._is_terminated = self.matcher.is_terminated() + if self._is_terminated: + break return True def validate_tokens(self, tokens: list[int]) -> list[int]: @@ -180,10 +183,15 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: Returns the prefix list of tokens that are accepted by the FSM. """ + if self._is_terminated: + return [] + accepted_tokens = [] for token in tokens: if self.matcher.accept_token(token): accepted_tokens.append(token) + if self.matcher.is_terminated(): + break else: break if len(accepted_tokens) > 0: @@ -203,8 +211,9 @@ def is_terminated(self) -> bool: return self._is_terminated def reset(self): - self.num_processed_tokens = 0 self.matcher.reset() + self.num_processed_tokens = 0 + self._is_terminated = False # cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 0629a6d2e0f6..c2b445a52dbb 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -131,14 +131,24 @@ def apply_grammar_bitmask( ) sorted_bitmask = sorted_bitmask_tensor.numpy() cumulative_index = 0 - for req_id in grammar_output.structured_output_request_ids: - num_spec_tokens = len(spec_tokens.get(req_id, ())) + for req_id, num_grammar_spec_tokens in zip( + grammar_output.structured_output_request_ids, + grammar_output.num_spec_tokens, + strict=True, + ): + num_worker_spec_tokens = len(spec_tokens.get(req_id, ())) + assert num_worker_spec_tokens <= num_grammar_spec_tokens if (logit_idx := struct_out_req_batch_indices.get(req_id)) is not None: - for i in range(1 + num_spec_tokens): + for i in range(num_worker_spec_tokens): bitmask_index = logit_idx + i sorted_bitmask[bitmask_index] = grammar_bitmask[cumulative_index + i] out_indices.append(bitmask_index) - cumulative_index += 1 + num_spec_tokens + bonus_index = logit_idx + num_worker_spec_tokens + sorted_bitmask[bonus_index] = grammar_bitmask[ + cumulative_index + num_grammar_spec_tokens + ] + out_indices.append(bonus_index) + cumulative_index += 1 + num_grammar_spec_tokens # Copy async to device. grammar_bitmask = sorted_bitmask_tensor.to(logits.device, non_blocking=True) diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 243d8f19052c..f6725cba470c 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -47,6 +47,15 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: except Exception: spec = None if getattr(spec, "dcp_replicated", False): + # A replicated KV group contains the complete sequence on + # every rank. Its attention kernel must therefore execute + # as a local DCP1 operation; applying DCP collectives would + # partition and reduce the same cache a second time. + layer_impl.dcp_world_size = 1 + layer_impl.dcp_rank = 0 + layer_impl.total_cp_world_size = 1 + layer_impl.total_cp_rank = 0 + layer_impl.need_to_return_lse_for_decode = False continue if vllm_config.speculative_config is not None and interleave_size > 1: assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, ( diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index f8336fa0749d..ff4cdde87bc7 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -152,6 +152,12 @@ def __init__( self.write_starts = new_buffer(self.num_rows, dtype=torch.int32) self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32) + @property + def cpu(self) -> torch.Tensor | None: + """Return the host backing tensor when this tensor uses UVA.""" + uva_buf = getattr(self, "_uva_buf", None) + return None if uva_buf is None else uva_buf.cpu + def stage_write( self, index: int, start: int, x: Iterable[int] | Iterable[float] ) -> None: diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 9bd83781583b..d1bedf75a79e 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -105,6 +105,11 @@ class InputBatch: max_req_tokens: int | None = None valid_num_draft_tokens_per_req: np.ndarray | None = None + # Optional host view of the request token table. Remote speculators use + # this to verify that a target prefix-cache hit belongs to retained draft + # state before reconnecting it. The tensor is shared, not copied. + all_token_ids_cpu: torch.Tensor | None = None + # When > 0, dummy batches carry seeded-random token ids instead of zeros. # All-zero ids embed identically, so every dummy token routes to the SAME # MoE experts — grouped expert reads collapse and the profiled cost of a diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 16685bd7a2a4..28f782620e2c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -65,7 +65,6 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask -from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -511,6 +510,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: max_num_logits=self.max_num_reqs * self.decode_query_len, vocab_size=self.vocab_size, device=self.device, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, ) if self.is_pooling_model and self.is_last_pp_rank: @@ -619,19 +619,15 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) - # One local block covers `block_size * dcp_shard_count` tokens in - # the global sequence. Replicated groups keep the full cache on - # every rank instead. group_cp_size = get_kv_cache_dcp_shard_count(spec, self.dcp_size) group_cp_sizes.append(group_cp_size) - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * group_cp_size + # Cache specifications own their block-table geometry. Attention + # caches account for their token-position DCP shards, while + # recurrent state and replicated attention caches remain unscaled. + max_num_blocks = spec.max_num_blocks_per_req( + self.vllm_config, block_table_max_model_len ) - # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens if isinstance(spec, MambaSpec): - max_num_blocks = ( - max_num_blocks if self.cache_config.enable_prefix_caching else 1 - ) + spec.num_speculative_blocks max_num_blocks = get_block_table_width( max_num_blocks, spec.block_size, token_alignment=None ) @@ -1562,6 +1558,7 @@ def prepare_inputs( prompt_lens=prompt_lens, max_req_tokens=max_req_tokens, valid_num_draft_tokens_per_req=valid_num_draft_tokens_per_req, + all_token_ids_cpu=self.req_states.all_token_ids.cpu, ) # InputBuffers are reused across real, dummy, and captured batches. # Clear stale padding before a capacity manager optionally marks a @@ -1646,6 +1643,7 @@ def sample( input_batch, grammar_output.structured_output_request_ids, grammar_output.grammar_bitmask, + grammar_output.num_spec_tokens, ) if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index e72e3afbcb11..6ceb8371b0a3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -73,6 +73,12 @@ def __init__( # columns and the running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" if self._align_mode: + # The physical attention page and the logical recurrent-state + # checkpoint can have different token widths. Prefix-hit requests + # must resume from the recurrent-state grid used by MambaSpec. + self._mamba_block_size = ( + self.cache_config.mamba_block_size or self.cache_config.block_size + ) self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device ) @@ -93,7 +99,7 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: if self._align_mode: # Seed the running state block from the resumed/prefilled position. self._mamba_state_idx_gpu[req_index].fill_( - (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + (new_req_data.num_computed_tokens - 1) // self._mamba_block_size ) def _get_mamba_group_info( @@ -109,6 +115,10 @@ def _get_mamba_group_info( specs.append(spec) assert specs, "no mamba layers in the model" assert all(specs[0] == s for s in specs) + assert specs[0].block_size == self._mamba_block_size, ( + "Mamba state migration and cache allocation must use the same " + "checkpoint cadence" + ) self._mamba_group_ids = group_ids self._mamba_spec = specs[0] return self._mamba_group_ids, self._mamba_spec diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index c70f169f7be6..767825145dba 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + import torch from vllm.config import VllmConfig @@ -9,12 +11,36 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): speculative_config = vllm_config.speculative_config assert speculative_config is not None if speculative_config.method == "dflash": + remote_address = os.environ.get("VLLM_K3_DRAFT_REMOTE_ADDRESS") + if remote_address: + from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import ( + RemoteK3DSparkSpeculator, + ) + + return RemoteK3DSparkSpeculator( + vllm_config, + device, + address=remote_address, + ) from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( DFlashSpeculator, ) return DFlashSpeculator(vllm_config, device) elif speculative_config.method == "dspark": + remote_address = os.environ.get( + "VLLM_K3_DRAFT_REMOTE_ADDRESS" + ) or os.environ.get("VLLM_K3_DSPARK_REMOTE_ADDRESS") + if remote_address: + from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import ( + RemoteK3DSparkSpeculator, + ) + + return RemoteK3DSparkSpeculator( + vllm_config, + device, + address=remote_address, + ) from vllm.v1.worker.gpu.spec_decode.dspark.speculator import ( DSparkSpeculator, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 256f23c42ace..2ee2f579fb5d 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -81,11 +81,17 @@ def maybe_load_mask_embedding( def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag - from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + from vllm.model_executor.models.qwen3_dflash import ( + dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, + ) speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style # Select an attention backend that supports the drafter's attention: mixing # a non-causal layer onto a causal-only backend would fail. draft_vllm_config = replace( @@ -146,4 +152,11 @@ def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo del dflash_model.lm_head dflash_model.lm_head = target_lm_head + # Opt-in rowwise-fp8 draft head (VLLM_DSPARK_FP8_DRAFT_HEAD). Runs after + # the lm_head aliasing above and before CUDA-graph capture, because the + # captured draft step must not quantize lazily. + maybe_init_fp8_draft_head = getattr(dflash_model, "maybe_init_fp8_draft_head", None) + if maybe_init_fp8_draft_head is not None: + maybe_init_fp8_draft_head() + return dflash_model diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py new file mode 100644 index 000000000000..876d6ad9eef7 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py @@ -0,0 +1,775 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Verifier-side proxy for a dedicated RTX 3090 K3 draft process.""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import dataclass +from typing import Any + +import torch +import zmq + +from vllm.config import VllmConfig +from vllm.distributed import get_tp_group +from vllm.logger import init_logger +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) +from vllm.v1.worker.gpu.spec_decode.speculator import ( + BaseSpeculator, + CUDAGraphCapturePhase, +) + +logger = init_logger(__name__) + +PROTOCOL_VERSION = 2 + + +@dataclass +class _RetainedRequestPrefix: + token_ids: torch.Tensor + committed_end: int + context_start: int + serial: int + + +def _build_valid_context_plan( + input_batch: InputBatch, + rejected_counts: list[int], +) -> tuple[list[int], list[int]]: + """Return valid row indices and per-request counts.""" + if len(rejected_counts) != input_batch.num_reqs: + raise ValueError("Rejected-token count does not match the request batch") + gather_indices: list[int] = [] + valid_counts: list[int] = [] + offset = 0 + for request_idx, (scheduled, rejected) in enumerate( + zip(input_batch.num_scheduled_tokens.tolist(), rejected_counts) + ): + valid = int(scheduled) - int(rejected) + if not 0 <= valid <= int(scheduled): + raise ValueError( + f"Invalid valid-context length for request {request_idx}: " + f"scheduled={scheduled}, rejected={rejected}" + ) + gather_indices.extend(range(offset, offset + valid)) + valid_counts.append(valid) + offset += int(scheduled) + return gather_indices, valid_counts + + +def _anchor_positions_from_context( + context_counts: list[int], context_positions: torch.Tensor +) -> list[int]: + """Return the position immediately following each request's context.""" + anchors: list[int] = [] + offset = 0 + for count in context_counts: + if count <= 0: + raise ValueError("Every remote draft request requires context rows") + offset += count + anchors.append(int(context_positions[offset - 1]) + 1) + if offset != context_positions.numel(): + raise ValueError("Remote draft context counts do not match the position tensor") + return anchors + + +def _contiguous_draft_output( + draft_tokens: torch.Tensor, + num_reqs: int, + num_speculative_tokens: int, +) -> torch.Tensor: + """Return the active TP-broadcast region with a compact row stride.""" + return draft_tokens[:num_reqs, :num_speculative_tokens].contiguous() + + +class RemoteK3DSparkSpeculator(BaseSpeculator): + """Forward target auxiliary states to a standalone greedy draft server.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + *, + address: str, + ) -> None: + self.vllm_config = vllm_config + self.device = device + self.speculative_config = vllm_config.speculative_config + assert self.speculative_config is not None + self.method = str(self.speculative_config.method) + if self.method not in ("dspark", "dflash"): + raise ValueError(f"Unsupported remote K3 draft method: {self.method}") + if self.speculative_config.draft_sample_method != "greedy": + raise ValueError("Remote K3 draft currently supports greedy drafting only") + if self.speculative_config.rejection_sample_method != "block": + raise ValueError( + "Remote K3 draft currently requires block rejection sampling" + ) + if vllm_config.model_config.dtype != torch.bfloat16: + raise ValueError("Remote K3 DSpark transport currently requires BF16") + self.num_speculative_steps = int(self.speculative_config.num_speculative_tokens) + self.max_num_reqs = int(vllm_config.scheduler_config.max_num_seqs) + self.max_num_tokens = int(vllm_config.scheduler_config.max_num_batched_tokens) + draft_hf_config = self.speculative_config.draft_model_config.hf_config + aux_layers = get_eagle3_aux_layers_from_config(self.speculative_config) + if not aux_layers: + raise ValueError( + f"Remote K3 {self.method} config does not declare auxiliary layers" + ) + self.num_aux_layers = len(aux_layers) + target_hidden_size = int( + getattr(draft_hf_config, "target_hidden_size", None) + or draft_hf_config.hidden_size + ) + self.raw_context_width = int(target_hidden_size * self.num_aux_layers) + self.address = address + self.timeout_ms = int( + os.environ.get( + "VLLM_K3_DRAFT_REMOTE_TIMEOUT_MS", + os.environ.get("VLLM_K3_DSPARK_REMOTE_TIMEOUT_MS", "30000"), + ) + ) + self.supports_mm_inputs = False + self.draft_logits: torch.Tensor | None = None + self.draft_tokens = torch.full( + (self.max_num_reqs, self.num_speculative_steps), + -1, + dtype=torch.int64, + device=device, + ) + self._known_requests: set[str] = set() + self._disabled_requests: set[str] = set() + self._active_requests: set[str] = set() + self._retained_prefixes: dict[str, _RetainedRequestPrefix] = {} + self._retained_serial = 0 + self._remote_max_requests = self.max_num_reqs + self._remote_block_size = 1 + self._remote_window_size = 0 + self._remote_prefix_cache_tokens = 0 + self._timing_log_interval = int( + os.environ.get("VLLM_K3_DRAFT_TIMING_LOG_INTERVAL", "0") + ) + if self._timing_log_interval < 0: + raise ValueError("VLLM_K3_DRAFT_TIMING_LOG_INTERVAL must be >= 0") + self._timing_count = 0 + self._timing_totals_ms: dict[str, float] = {} + + tp_group = get_tp_group() + self._tp_group = tp_group + self._tp_rank = int(tp_group.rank_in_group) + self._zmq_context: zmq.Context | None = None + self._socket: zmq.Socket | None = None + if self._tp_rank == 0: + # P2P is unavailable on the target host. Keep the mandatory D2H + # hop off pageable memory so it can be queued directly after the + # target forward on the current stream. + self._positions_staging = torch.empty( + self.max_num_tokens, + dtype=torch.int64, + pin_memory=True, + ) + self._context_staging = torch.empty( + (self.max_num_tokens, self.raw_context_width), + dtype=vllm_config.model_config.dtype, + pin_memory=True, + ) + self._rejected_staging = torch.empty( + self.max_num_reqs, + dtype=torch.int32, + pin_memory=True, + ) + self._anchor_staging = torch.empty( + self.max_num_reqs, + dtype=torch.int64, + pin_memory=True, + ) + self._zmq_context = zmq.Context() + self._connect() + response = self._rpc( + [json.dumps({"protocol": PROTOCOL_VERSION, "op": "PING"}).encode()] + ) + if response.get("op") != "PONG": + raise RuntimeError(f"Unexpected K3 draft health response: {response}") + if response.get("method") != self.method: + raise RuntimeError( + "Remote K3 draft method mismatch: " + f"target={self.method}, server={response.get('method')}" + ) + self._remote_max_requests = int( + response.get("max_requests", self.max_num_reqs) + ) + self._remote_block_size = int(response.get("block_size", 1)) + self._remote_window_size = int(response.get("window_size", 0)) + self._remote_prefix_cache_tokens = int( + response.get("prefix_cache_tokens", 0) + ) + if int(response.get("active_requests", 0)): + # A restarted verifier cannot safely identify state left by an + # older process, so establish a clean protocol epoch. + self._rpc( + [json.dumps({"protocol": PROTOCOL_VERSION, "op": "CLEAR"}).encode()] + ) + + logger.info( + "Remote K3 %s proxy initialized: address=%s, TP rank=%d, K=%d", + self.method, + address, + self._tp_rank, + self.num_speculative_steps, + ) + + def _connect(self) -> None: + assert self._zmq_context is not None + if self._socket is not None: + self._socket.close() + socket = self._zmq_context.socket(zmq.REQ) + socket.setsockopt(zmq.LINGER, 0) + socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms) + socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) + socket.connect(self.address) + self._socket = socket + + def _rpc(self, frames: list[bytes]) -> dict[str, Any]: + assert self._socket is not None + try: + self._socket.send_multipart(frames) + response = self._socket.recv_json() + except Exception: + self._connect() + raise + if not isinstance(response, dict) or not response.get("ok", False): + raise RuntimeError(f"K3 DSpark RPC failed: {response}") + if int(response.get("protocol", -1)) != PROTOCOL_VERSION: + raise RuntimeError(f"K3 DSpark protocol mismatch: {response}") + return response + + def init_cudagraph_manager(self, cudagraph_mode=None) -> None: + """The standalone drafter owns its CUDA graph lifecycle.""" + + def capture(self, *, capture_phase: CUDAGraphCapturePhase) -> None: + """The verifier has no local draft graph to capture.""" + + def _free_remote_requests(self, request_ids: set[str] | list[str]) -> None: + remote_request_ids = sorted(set(request_ids) & self._known_requests) + if not remote_request_ids: + return + self._rpc( + [ + json.dumps( + { + "protocol": PROTOCOL_VERSION, + "op": "FREE", + "request_ids": remote_request_ids, + } + ).encode() + ] + ) + self._known_requests.difference_update(remote_request_ids) + for request_id in remote_request_ids: + self._retained_prefixes.pop(request_id, None) + + def _ensure_remote_capacity(self, current_request_ids: set[str]) -> None: + while len(self._known_requests) >= self._remote_max_requests: + candidates = self._known_requests - current_request_ids + if not candidates: + raise RuntimeError( + "Remote DSpark request capacity is exhausted by active requests" + ) + request_id = min( + candidates, + key=lambda req_id: ( + self._retained_prefixes[req_id].serial + if req_id in self._retained_prefixes + else -1 + ), + ) + self._free_remote_requests({request_id}) + + @staticmethod + def _token_prefix( + input_batch: InputBatch, + request_idx: int, + prefix_end: int, + ) -> torch.Tensor | None: + token_table = input_batch.all_token_ids_cpu + if token_table is None or prefix_end < 0: + return None + state_idx = int(input_batch.idx_mapping_np[request_idx]) + if state_idx < 0 or prefix_end > token_table.shape[1]: + return None + return token_table[state_idx, :prefix_end] + + def _can_restore_prefix( + self, + retained: _RetainedRequestPrefix, + prefix_end: int, + ) -> bool: + if ( + prefix_end <= 0 + or retained.committed_end < prefix_end + or self._remote_window_size <= 0 + or self._remote_prefix_cache_tokens < self._remote_window_size + ): + return False + restore_start = max(0, prefix_end - self._remote_window_size) + restore_start = ( + restore_start // self._remote_block_size * self._remote_block_size + ) + retained_start = max( + retained.context_start, + retained.committed_end - self._remote_prefix_cache_tokens, + ) + return restore_start >= retained_start + + def _find_reconnect_source( + self, + token_prefix: torch.Tensor, + prefix_end: int, + current_request_ids: set[str], + ) -> str | None: + candidates: list[tuple[int, int, str]] = [] + for request_id in self._known_requests - current_request_ids: + retained = self._retained_prefixes.get(request_id) + if retained is None or not self._can_restore_prefix(retained, prefix_end): + continue + if torch.equal(retained.token_ids[:prefix_end], token_prefix): + candidates.append( + ( + retained.committed_end - prefix_end, + -retained.serial, + request_id, + ) + ) + return min(candidates)[2] if candidates else None + + def _reconnect_request( + self, + source_request_id: str, + request_id: str, + prefix_end: int, + token_prefix: torch.Tensor, + ) -> bool: + try: + response = self._rpc( + [ + json.dumps( + { + "protocol": PROTOCOL_VERSION, + "op": "RECONNECT", + "source_request_id": source_request_id, + "request_id": request_id, + "prefix_end": prefix_end, + } + ).encode() + ] + ) + except Exception: + logger.exception( + "Remote K3 DSpark prefix reconnect failed: source=%s, " + "request=%s, prefix_end=%d", + source_request_id, + request_id, + prefix_end, + ) + return False + + self._retained_prefixes.pop(source_request_id) + self._known_requests.discard(source_request_id) + self._known_requests.add(request_id) + self._retained_serial += 1 + self._retained_prefixes[request_id] = _RetainedRequestPrefix( + token_ids=token_prefix.clone(), + committed_end=prefix_end, + context_start=int(response.get("restored_start", 0)), + serial=self._retained_serial, + ) + logger.info( + "Remote K3 DSpark prefix reconnected: source=%s, request=%s, " + "prefix_end=%d, restored_start=%s, latency_ms=%.1f", + source_request_id, + request_id, + prefix_end, + response.get("restored_start"), + float(response.get("latency_ms", 0.0)), + ) + return True + + def _remember_prefix( + self, + input_batch: InputBatch, + request_idx: int, + request_id: str, + committed_end: int, + context_start: int | None = None, + ) -> None: + token_prefix = self._token_prefix(input_batch, request_idx, committed_end) + if token_prefix is None: + return + previous = self._retained_prefixes.get(request_id) + if context_start is None: + context_start = previous.context_start if previous is not None else 0 + self._retained_serial += 1 + self._retained_prefixes[request_id] = _RetainedRequestPrefix( + token_ids=token_prefix.clone(), + committed_end=committed_end, + context_start=context_start, + serial=self._retained_serial, + ) + + def _copy_tokens_from_response( + self, + response: dict[str, Any], + active_indices: list[int], + num_speculative_tokens: int, + ) -> None: + tokens = response.get("tokens") + expected_shape = (len(active_indices), num_speculative_tokens) + if ( + not isinstance(tokens, list) + or len(tokens) != expected_shape[0] + or any( + not isinstance(row, list) or len(row) != expected_shape[1] + for row in tokens + ) + ): + raise ValueError( + f"Remote DSpark token response has the wrong shape; " + f"expected={expected_shape}, got={tokens!r}" + ) + remote_tokens = torch.tensor(tokens, dtype=torch.int64, device=self.device) + active_gpu = torch.tensor(active_indices, dtype=torch.int64, device=self.device) + # ``draft_tokens`` is allocated at the configured maximum depth, while + # adaptive speculation and the per-batch schedule can request a + # smaller depth for an individual step. Copy into the matching width + # instead of requiring every response to have the maximum width. + self.draft_tokens[:, :num_speculative_tokens].index_copy_( + 0, active_gpu, remote_tokens + ) + + def _record_timing(self, timing_ms: dict[str, float]) -> None: + if self._timing_log_interval <= 0: + return + self._timing_count += 1 + for key, value in timing_ms.items(): + self._timing_totals_ms[key] = self._timing_totals_ms.get(key, 0.0) + value + if self._timing_count < self._timing_log_interval: + return + means = { + key: value / self._timing_count + for key, value in self._timing_totals_ms.items() + } + logger.info( + "Remote K3 %s timing over %d proposals (ms): %s", + self.method, + self._timing_count, + ", ".join(f"{key}={value:.3f}" for key, value in means.items()), + ) + self._timing_count = 0 + self._timing_totals_ms.clear() + + def _rank0_propose( + self, + input_batch: InputBatch, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + num_speculative_tokens: int, + ) -> None: + started = time.perf_counter() + if aux_hidden_states is None or len(aux_hidden_states) != self.num_aux_layers: + raise ValueError( + f"Remote K3 {self.method} requires {self.num_aux_layers} configured " + "target auxiliary hidden states" + ) + + num_reqs = input_batch.num_reqs + current_request_ids = set(input_batch.req_ids) + previous_active_requests = self._active_requests + self._disabled_requests.intersection_update(current_request_ids) + idx_mapping = input_batch.idx_mapping[:num_reqs].long() + sampled_counts = num_sampled[:num_reqs] + sampled_anchors = last_sampled[idx_mapping, 0] + prefill_anchors = next_prefill_tokens[0, idx_mapping] + anchor_tokens = torch.where( + sampled_counts > 0, + sampled_anchors, + prefill_anchors, + ).to(torch.int64) + # Queue both small D2H copies and wait once. The same stream owns the + # preceding token-table update, so this synchronization also makes the + # UVA-backed table safe for prefix matching below. + rejected_staging = self._rejected_staging[:num_reqs] + anchor_staging = self._anchor_staging[:num_reqs] + rejected_staging.copy_(num_rejected[:num_reqs], non_blocking=True) + anchor_staging.copy_(anchor_tokens, non_blocking=True) + torch.cuda.current_stream(self.device).synchronize() + rejected_counts = rejected_staging.tolist() + anchor_tokens_cpu = anchor_staging.tolist() + gather_indices, valid_counts = _build_valid_context_plan( + input_batch, rejected_counts + ) + metadata_ready = time.perf_counter() + + active_indices: list[int] = [] + requests: list[dict[str, Any]] = [] + request_context_starts: list[int | None] = [] + selected_gather_indices: list[int] = [] + gather_offset = 0 + for request_idx, request_id in enumerate(input_batch.req_ids): + valid_count = valid_counts[request_idx] + request_gather = gather_indices[gather_offset : gather_offset + valid_count] + gather_offset += valid_count + if request_id in self._disabled_requests or valid_count <= 0: + continue + first_position = int(input_batch.num_computed_tokens_np[request_idx]) + is_continuing = ( + request_id in previous_active_requests + and request_id in self._known_requests + ) + reset = False + context_start: int | None = None + if not is_continuing: + if first_position == 0: + if request_id in self._known_requests: + self._free_remote_requests({request_id}) + self._ensure_remote_capacity(current_request_ids) + reset = True + context_start = 0 + else: + token_prefix = self._token_prefix( + input_batch, + request_idx, + first_position, + ) + source_request_id: str | None = None + if token_prefix is not None: + retained = self._retained_prefixes.get(request_id) + if ( + request_id in self._known_requests + and retained is not None + and self._can_restore_prefix(retained, first_position) + and torch.equal( + retained.token_ids[:first_position], token_prefix + ) + ): + source_request_id = request_id + else: + source_request_id = self._find_reconnect_source( + token_prefix, + first_position, + current_request_ids, + ) + if source_request_id is None or token_prefix is None: + if request_id in self._known_requests: + self._free_remote_requests({request_id}) + self._ensure_remote_capacity(current_request_ids) + reset = True + context_start = first_position + logger.warning( + "Remote K3 %s cold-bootstrapping cache-restored " + "request %s at position %d from %d fresh context " + "rows; target verification preserves correctness.", + self.method, + request_id, + first_position, + valid_count, + ) + elif not self._reconnect_request( + source_request_id, + request_id, + first_position, + token_prefix, + ): + self._disabled_requests.add(request_id) + continue + requests.append( + { + "request_id": request_id, + "reset": reset, + "reset_position": first_position if reset else 0, + "context_count": valid_count, + "anchor_token_id": int(anchor_tokens_cpu[request_idx]), + } + ) + self._known_requests.add(request_id) + active_indices.append(request_idx) + request_context_starts.append(context_start) + selected_gather_indices.extend(request_gather) + + self._active_requests = self._known_requests & current_request_ids + if not active_indices: + return + requests_ready = time.perf_counter() + + indices_gpu = torch.tensor( + selected_gather_indices, dtype=torch.int64, device=self.device + ) + positions = input_batch.positions.index_select(0, indices_gpu) + context = torch.cat( + [hidden.index_select(0, indices_gpu) for hidden in aux_hidden_states], + dim=-1, + ) + num_context_rows = int(context.shape[0]) + if num_context_rows > self.max_num_tokens: + raise ValueError( + f"Remote DSpark context has {num_context_rows} rows, max is " + f"{self.max_num_tokens}" + ) + if context.shape[1] != self.raw_context_width: + raise ValueError( + f"Remote DSpark context width is {context.shape[1]}, expected " + f"{self.raw_context_width}" + ) + context_ready = time.perf_counter() + positions_staging = self._positions_staging[:num_context_rows] + context_staging = self._context_staging[:num_context_rows] + positions_staging.copy_(positions, non_blocking=True) + context_staging.copy_(context, non_blocking=True) + torch.cuda.current_stream(self.device).synchronize() + context_copied = time.perf_counter() + anchor_positions = _anchor_positions_from_context( + [int(request["context_count"]) for request in requests], + positions_staging, + ) + for request, anchor_position in zip(requests, anchor_positions): + request["anchor_position"] = anchor_position + positions_frame = positions_staging.numpy().tobytes() + # NumPy has inconsistent bfloat16 support; preserve its exact bits as u16. + context_frame = context_staging.view(torch.uint16).numpy().tobytes() + serialized = time.perf_counter() + header = { + "protocol": PROTOCOL_VERSION, + "op": "PROPOSE", + "projected": False, + "num_speculative_tokens": num_speculative_tokens, + "requests": requests, + } + response = self._rpc( + [json.dumps(header).encode(), positions_frame, context_frame] + ) + rpc_done = time.perf_counter() + self._copy_tokens_from_response( + response, + active_indices, + num_speculative_tokens, + ) + output_copied = time.perf_counter() + timing_ms = { + "metadata_d2h": (metadata_ready - started) * 1000, + "request_plan": (requests_ready - metadata_ready) * 1000, + "context_gather": (context_ready - requests_ready) * 1000, + "context_d2h": (context_copied - context_ready) * 1000, + "serialize": (serialized - context_copied) * 1000, + "rpc_roundtrip": (rpc_done - serialized) * 1000, + "tokens_h2d": (output_copied - rpc_done) * 1000, + "client_total": (output_copied - started) * 1000, + } + server_timing = response.get("timing_ms") + if isinstance(server_timing, dict): + for key, value in server_timing.items(): + if isinstance(value, int | float): + timing_ms[f"server_{key}"] = float(value) + self._record_timing(timing_ms) + for request_idx, request, anchor_position, context_start in zip( + active_indices, + requests, + anchor_positions, + request_context_starts, + ): + self._remember_prefix( + input_batch, + request_idx, + str(request["request_id"]), + anchor_position, + context_start, + ) + + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + num_speculative_tokens: int | None = None, + num_tokens_across_dp: torch.Tensor | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + del ( + attn_metadata, + slot_mappings, + last_hidden_states, + temperature, + seeds, + num_tokens_across_dp, + skip_attn_for_dummy_run, + mm_inputs, + ) + active_k = ( + int(num_speculative_tokens) + if num_speculative_tokens is not None + else self.num_speculative_steps + ) + # A scheduler can intentionally disable speculation for one step (for + # example, a request with max_tokens=1). ModelRunner treats an empty + # second dimension as a normal non-speculative step. + if active_k == 0: + return self.draft_tokens[: input_batch.num_reqs, :0].contiguous() + if not 1 <= active_k <= self.num_speculative_steps: + raise ValueError( + f"Remote DSpark depth must be in [1, " + f"{self.num_speculative_steps}], got {active_k}" + ) + output = self.draft_tokens[: input_batch.num_reqs, :active_k] + output.fill_(-1) + if self._tp_rank == 0 and not (dummy_run or is_profile): + try: + self._rank0_propose( + input_batch, + aux_hidden_states, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + active_k, + ) + except Exception: + output.fill_(-1) + # The verifier cannot know whether a timed-out request mutated + # remote KV. Fail closed for those requests until they leave + # the active batch; FREE remains safe even if the server never + # created the state. + self._disabled_requests.update(input_batch.req_ids) + logger.exception( + "Remote K3 DSpark proposal failed; drafting is disabled for " + "this step" + ) + # Slicing the active depth from the max-width persistent buffer leaves + # a larger row stride whenever adaptive K is below the configured + # maximum. NCCL broadcast requires a contiguous tensor. Materialize + # only the tiny [batch, K] result after rank 0 has populated it. + output = _contiguous_draft_output( + self.draft_tokens, + input_batch.num_reqs, + active_k, + ) + self._tp_group.broadcast(output, src=0) + return output diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index cd23ae1bb559..a61d5261dfba 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext import torch.nn as nn @@ -48,6 +48,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if hasattr(target_model, "get_language_model") else target_model ) + rope_ownership: AbstractContextManager[None] if getattr(draft_model_config.hf_config, "model_type", None) == "k3_dspark": from vllm.models.kimi_k3.nvidia.dspark_mla import ( protect_k3_compact_rope_sources, @@ -59,7 +60,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo with rope_ownership, set_model_tag("dspark_head"): draft_model = get_model( - vllm_config=draft_vllm_config, model_config=draft_model_config + vllm_config=draft_vllm_config, + model_config=draft_model_config, + load_config=speculative_config.draft_load_config, ) if get_pp_group().world_size != 1: diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 51484005a8ed..9ad96093268d 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -34,6 +34,8 @@ def limit_draft_tokens( "Speculator returned unsupported draft shape " f"{tuple(draft_tokens.shape)}; expected a 2D tensor." ) + if num_speculative_tokens == 0: + return draft_tokens[:, :0] if not 1 <= num_speculative_tokens <= max_num_speculative_tokens: raise RuntimeError( "Scheduler selected an invalid speculative-token count " diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 34f00086d2a2..3be14f97f31b 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -9,8 +9,61 @@ from vllm.v1.worker.gpu.input_batch import InputBatch +def _build_grammar_row_mapping( + req_ids: list[str], + grammar_req_ids: list[str], + grammar_num_spec_tokens: list[int], + cu_num_logits_np: np.ndarray, + num_draft_tokens_per_req: np.ndarray | None, + num_bonus_tokens: int, +) -> tuple[list[int], list[int]]: + """Map serialized grammar rows to the active compact logits layout.""" + assert len(grammar_req_ids) == len(grammar_num_spec_tokens) + assert num_bonus_tokens in (0, 1) + + req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} + source_indices: list[int] = [] + logits_indices: list[int] = [] + source_offset = 0 + + for grammar_req_id, num_source_drafts in zip( + grammar_req_ids, + grammar_num_spec_tokens, + strict=True, + ): + req_idx = req_id_to_idx[grammar_req_id] + num_active_drafts = ( + 0 + if num_draft_tokens_per_req is None + else int(num_draft_tokens_per_req[req_idx]) + ) + assert 0 <= num_active_drafts <= num_source_drafts + + logits_start = int(cu_num_logits_np[req_idx]) + num_active_logits = int( + cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx] + ) + assert num_active_logits == num_active_drafts + num_bonus_tokens + + source_indices.extend(range(source_offset, source_offset + num_active_drafts)) + logits_indices.extend(range(logits_start, logits_start + num_active_drafts)) + if num_bonus_tokens: + source_indices.append(source_offset + num_source_drafts) + logits_indices.append(logits_start + num_active_drafts) + + source_offset += num_source_drafts + num_bonus_tokens + + return source_indices, logits_indices + + class StructuredOutputsWorker: - def __init__(self, max_num_logits: int, vocab_size: int, device: torch.device): + def __init__( + self, + max_num_logits: int, + vocab_size: int, + device: torch.device, + num_bonus_tokens: int, + ): self.logits_indices = torch.zeros( max_num_logits, dtype=torch.int32, device=device ) @@ -19,6 +72,7 @@ def __init__(self, max_num_logits: int, vocab_size: int, device: torch.device): ) self.device = device self.copy_stream = torch.cuda.Stream() + self.num_bonus_tokens = num_bonus_tokens def apply_grammar_bitmask( self, @@ -26,27 +80,31 @@ def apply_grammar_bitmask( input_batch: InputBatch, grammar_req_ids: list[str], grammar_bitmask: np.ndarray, + grammar_num_spec_tokens: list[int], ) -> None: if not grammar_req_ids: return - # Asynchronously copy the bitmask to GPU. + source_indices, mapping = _build_grammar_row_mapping( + input_batch.req_ids, + grammar_req_ids, + grammar_num_spec_tokens, + input_batch.cu_num_logits_np, + input_batch.num_draft_tokens_per_req, + self.num_bonus_tokens, + ) + expected_source_rows = sum( + num_drafts + self.num_bonus_tokens for num_drafts in grammar_num_spec_tokens + ) + assert grammar_bitmask.shape[0] == expected_source_rows + grammar_bitmask = grammar_bitmask[source_indices] + + # Asynchronously copy the active bitmask rows to GPU. with torch.cuda.stream(self.copy_stream): bitmask = async_copy_to_gpu( grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]] ) - # Construct bitmask -> logits mapping - mapping: list[int] = [] - req_ids = input_batch.req_ids - cu_num_logits = input_batch.cu_num_logits_np.tolist() - req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} - for grammar_req_id in grammar_req_ids: - req_idx = req_id_to_idx[grammar_req_id] - logits_start_idx = cu_num_logits[req_idx] - logits_end_idx = cu_num_logits[req_idx + 1] - mapping.extend(range(logits_start_idx, logits_end_idx)) - # Asynchronously copy the mapping to GPU. with torch.cuda.stream(self.copy_stream): logits_indices = torch.tensor( diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index ab745d70a295..3381655e3d3e 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -367,7 +367,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32 ) grammar_output = GrammarOutput( - structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask + structured_output_request_ids=req_ids, + grammar_bitmask=grammar_bitmask, + num_spec_tokens=[0] * len(req_ids), ) worker_sample_tokens(grammar_output) @@ -581,6 +583,7 @@ def _profile_sps_curve( if sps_debug: events = model_runner._sps_debug_events model_runner._sps_debug_events = None + assert events is not None # Skip the warmup iters; report mean verify/draft GPU ms. timed = events[warmup_iters:] if timed: diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index f87a34a364f9..4add17c7e851 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -188,19 +188,46 @@ def _copy_mamba_state_block( src_block_id = tl.load(block_table_base + src_col).to(tl.int64) dim_rows = tl.load(state_dim_row_count_ptr + state_idx) row_stride = tl.load(state_dim_row_stride_ptr + state_idx) - per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size - bias_bytes = token_bias.to(tl.int64) * state_elem_size src_block_addr = state_base_addr + src_block_id * state_block_stride offsets = tl.arange(0, COPY_BLOCK_SIZE) - for d in range(0, dim_rows): - row_src = src_block_addr + d * row_stride + bias_bytes - row_dst = dst_addr + d * row_stride - for i in range(0, per_row_bytes, COPY_BLOCK_SIZE): - mask = (i + offsets) < per_row_bytes - curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8)) - curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(curr_src, mask=mask) - tl.store(curr_dst, data, mask=mask) + + # Stable row-to-lane ownership makes left shifts memmove-safe while + # exposing the dimension rows in parallel. All addresses retain + # state_elem_size alignment: tensor strides and token offsets are + # measured in whole elements before conversion to bytes. + num_dst_tokens = conv_width - token_bias + for token_idx in range(0, num_dst_tokens): + for row_base in range(0, dim_rows, COPY_BLOCK_SIZE): + rows = row_base + offsets + mask = rows < dim_rows + src_byte_addr = ( + src_block_addr + + rows * row_stride + + (token_idx + token_bias) * state_elem_size + ) + dst_byte_addr = ( + dst_addr + rows * row_stride + token_idx * state_elem_size + ) + if state_elem_size == 2: + src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16)) + dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16)) + data_u16 = tl.load(src_u16, mask=mask) + tl.store(dst_u16, data_u16, mask=mask) + elif state_elem_size == 4: + src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32)) + dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32)) + data_u32 = tl.load(src_u32, mask=mask) + tl.store(dst_u32, data_u32, mask=mask) + else: + for byte_idx in range(0, state_elem_size): + src_u8 = (src_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + dst_u8 = (dst_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + data_u8 = tl.load(src_u8, mask=mask) + tl.store(dst_u8, data_u8, mask=mask) return if is_conv_state: @@ -209,22 +236,40 @@ def _copy_mamba_state_block( # SD conv: copy # state[bt[src_col], token_bias:] -> # state[bt[dst_col], :conv_width - token_bias] - # Small per-block bytes (~60-80 KiB) make tiling degenerate, so - # conv runs as a single-CTA memcpy (NUM_TILES=1). src_block_id = tl.load(block_table_base + src_col).to(tl.int64) - src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size - src_addr = state_base_addr + src_block_id * state_block_stride + src_offset - copy_size = ( - (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size - ) - _memcpy_u64_tiled( - src_addr, - dst_addr, - copy_size, - tile_idx, - COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, - NUM_TILES=1, - ) + src_block_addr = state_base_addr + src_block_id * state_block_stride + token_bytes = state_inner_size * state_elem_size + num_dst_tokens = conv_width - token_bias + + # Distinct blocks and exact self-copies cannot have a destructive + # overlap, so retain the u64-vectorized single-CTA copy. + if src_block_id != dest_block_id or token_bias == 0: + src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes + copy_size = num_dst_tokens.to(tl.int64) * token_bytes + _memcpy_u64_tiled( + src_addr, + dst_addr, + copy_size, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) + return + + # Copy tokens from low to high. Each token-sized source and destination + # region is disjoint, so same-block left shifts are memmove-safe + # without a barrier. + for token_idx in range(0, num_dst_tokens): + src_token = src_block_addr + (token_idx + token_bias) * token_bytes + dst_token = dst_addr + token_idx * token_bytes + _memcpy_u64_tiled( + src_token, + dst_token, + token_bytes, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) return # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] @@ -521,6 +566,7 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): src_ptr = tl.load(src_ptrs + pid) dst_ptr = tl.load(dst_ptrs + pid) size = tl.load(sizes + pid) + is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr offsets = tl.arange(0, BLOCK_SIZE) for i in range(0, size, BLOCK_SIZE): @@ -530,6 +576,10 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) data = tl.load(curr_src_ptr, mask=mask) + if is_left_overlap: + # Preserve each lane's source before a lower-address lane stores + # over it. The condition is uniform within the program. + tl.debug_barrier() tl.store(curr_dst_ptr, data, mask=mask)