@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 4aaaf8c8ce
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 6, 23:00 UTC

feat(spec_decode): fuse EAGLE step slot mapping and metadata updates (#33503)

Signed-off-by: sladynnunes <snunes@usc.edu>

Sladyn committed Mar 10, 2026 at 21:35 UTC 4aaaf8c8ce517dd97a1cb2610e57fc161755a3a3
3 files changed +318 -59
tests/v1/spec_decode/test_eagle_step_kernel.py new
+175
@@ -0,0 +1,175 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Unit tests for the fused EAGLE slot mapping kernel."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.v1.spec_decode.utils import (
9 + PADDING_SLOT_ID,
10 + eagle_step_update_slot_mapping_and_metadata,
11 +)
12 +
13 +# Skip if no CUDA - Triton kernel requires GPU
14 +pytest.importorskip("triton")
15 +if not torch.cuda.is_available():
16 + pytest.skip("CUDA required for EAGLE kernel tests", allow_module_level=True)
17 +
18 +
19 +def _reference_eagle_step_slot_mapping(
20 + positions_1d: torch.Tensor,
21 + block_table_tensor: torch.Tensor,
22 + seq_lens: torch.Tensor,
23 + block_size: int,
24 + max_model_len: int,
25 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
26 + """Python reference for eagle_step_update_slot_mapping_and_metadata."""
27 + new_positions = positions_1d + 1
28 + exceeds_max = new_positions >= max_model_len
29 + clamped_positions = torch.where(
30 + exceeds_max, torch.zeros_like(positions_1d), new_positions
31 + )
32 + block_numbers = (clamped_positions // block_size).clamp(
33 + max=block_table_tensor.shape[1] - 1
34 + )
35 + block_ids = block_table_tensor[
36 + torch.arange(positions_1d.shape[0], device=positions_1d.device),
37 + block_numbers.long(),
38 + ].long()
39 + slot_mapping = block_ids * block_size + (clamped_positions % block_size)
40 + slot_mapping = torch.where(
41 + exceeds_max, torch.full_like(slot_mapping, PADDING_SLOT_ID), slot_mapping
42 + )
43 + new_seq_lens = torch.where(exceeds_max, torch.ones_like(seq_lens), seq_lens + 1)
44 + new_seq_lens = new_seq_lens.clamp(max=max_model_len)
45 + return clamped_positions, slot_mapping, new_seq_lens
46 +
47 +
48 +def test_eagle_step_slot_mapping_kernel():
49 + """Test fused kernel matches Python reference for slot mapping and metadata."""
50 + device = torch.device("cuda")
51 + batch_size = 32
52 + block_size = 16
53 + max_model_len = 4096
54 + n_blocks_per_req = (max_model_len + block_size - 1) // block_size
55 +
56 + positions_1d = torch.randint(
57 + 0, max_model_len - 10, (batch_size,), dtype=torch.int64, device=device
58 + )
59 + block_table_tensor = torch.randint(
60 + 0, 1000, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device
61 + )
62 + seq_lens = torch.randint(1, 100, (batch_size,), dtype=torch.int32, device=device)
63 +
64 + ref_clamped, ref_slot, ref_seq_lens = _reference_eagle_step_slot_mapping(
65 + positions_1d.clone(),
66 + block_table_tensor,
67 + seq_lens.clone(),
68 + block_size,
69 + max_model_len,
70 + )
71 +
72 + out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device)
73 + out_slot = torch.zeros(batch_size, dtype=torch.int64, device=device)
74 + seq_lens_copy = seq_lens.clone()
75 + eagle_step_update_slot_mapping_and_metadata(
76 + positions_1d=positions_1d,
77 + block_table_tensor=block_table_tensor,
78 + seq_lens=seq_lens_copy,
79 + block_size=block_size,
80 + max_model_len=max_model_len,
81 + out_clamped_positions=out_clamped,
82 + out_slot_mapping=out_slot,
83 + )
84 +
85 + assert torch.equal(out_clamped, ref_clamped), (
86 + f"clamped: {out_clamped} vs {ref_clamped}"
87 + )
88 + assert torch.equal(out_slot, ref_slot), f"slot: {out_slot} vs {ref_slot}"
89 + assert torch.equal(seq_lens_copy, ref_seq_lens), (
90 + f"seq_lens: {seq_lens_copy} vs {ref_seq_lens}"
91 + )
92 +
93 +
94 +def test_eagle_step_slot_mapping_kernel_exceeds_max():
95 + """Test fused kernel when position exceeds max_model_len."""
96 + device = torch.device("cuda")
97 + batch_size = 4
98 + block_size = 16
99 + max_model_len = 100
100 + n_blocks_per_req = (max_model_len + block_size - 1) // block_size
101 +
102 + positions_1d = torch.tensor([50, 98, 99, 100], dtype=torch.int64, device=device)
103 + block_table_tensor = torch.randint(
104 + 0, 100, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device
105 + )
106 + seq_lens = torch.tensor([51, 99, 100, 101], dtype=torch.int32, device=device)
107 +
108 + out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device)
109 + out_slot = torch.zeros(batch_size, dtype=torch.int64, device=device)
110 + eagle_step_update_slot_mapping_and_metadata(
111 + positions_1d=positions_1d,
112 + block_table_tensor=block_table_tensor,
113 + seq_lens=seq_lens,
114 + block_size=block_size,
115 + max_model_len=max_model_len,
116 + out_clamped_positions=out_clamped,
117 + out_slot_mapping=out_slot,
118 + )
119 +
120 + assert out_clamped[0].item() == 51
121 + assert out_clamped[1].item() == 99
122 + assert out_clamped[2].item() == 0
123 + assert out_clamped[3].item() == 0
124 + assert out_slot[2].item() == PADDING_SLOT_ID
125 + assert out_slot[3].item() == PADDING_SLOT_ID
126 + assert seq_lens[2].item() == 1
127 + assert seq_lens[3].item() == 1
128 +
129 +
130 +def test_eagle_step_slot_mapping_kernel_cudagraph_padding():
131 + """Test that padding threads write PADDING_SLOT_ID when
132 + input_batch_size > batch_size (cudagraph padding)."""
133 + device = torch.device("cuda")
134 + batch_size = 4
135 + input_batch_size = 8
136 + block_size = 16
137 + max_model_len = 4096
138 + n_blocks_per_req = (max_model_len + block_size - 1) // block_size
139 +
140 + positions_1d = torch.tensor([10, 20, 30, 40], dtype=torch.int64, device=device)
141 + block_table_tensor = torch.randint(
142 + 0, 100, (batch_size, n_blocks_per_req), dtype=torch.int32, device=device
143 + )
144 + seq_lens = torch.tensor([11, 21, 31, 41], dtype=torch.int32, device=device)
145 +
146 + ref_clamped, ref_slot, ref_seq_lens = _reference_eagle_step_slot_mapping(
147 + positions_1d.clone(),
148 + block_table_tensor,
149 + seq_lens.clone(),
150 + block_size,
151 + max_model_len,
152 + )
153 +
154 + out_clamped = torch.zeros(batch_size, dtype=torch.int64, device=device)
155 + out_slot = torch.full((input_batch_size,), -999, dtype=torch.int64, device=device)
156 + seq_lens_copy = seq_lens.clone()
157 + eagle_step_update_slot_mapping_and_metadata(
158 + positions_1d=positions_1d,
159 + block_table_tensor=block_table_tensor,
160 + seq_lens=seq_lens_copy,
161 + block_size=block_size,
162 + max_model_len=max_model_len,
163 + out_clamped_positions=out_clamped,
164 + out_slot_mapping=out_slot,
165 + input_batch_size=input_batch_size,
166 + )
167 +
168 + # Real slots should match the reference
169 + assert torch.equal(out_clamped, ref_clamped)
170 + assert torch.equal(out_slot[:batch_size], ref_slot)
171 + assert torch.equal(seq_lens_copy, ref_seq_lens)
172 +
173 + # Padding slots should be PADDING_SLOT_ID
174 + for i in range(batch_size, input_batch_size):
175 + assert out_slot[i].item() == PADDING_SLOT_ID
vllm/v1/spec_decode/eagle.py
+35 -59
@@ -44,6 +44,7 @@ from vllm.v1.spec_decode.utils import (
44 copy_and_expand_eagle_inputs_kernel,
45 eagle_prepare_inputs_padded_kernel,
46 eagle_prepare_next_token_padded_kernel,
47 + eagle_step_update_slot_mapping_and_metadata,
48 extend_all_queries_by_N,
49 )
50 from vllm.v1.utils import CpuGpuBuffer
@@ -533,41 +534,46 @@ class SpecDecodeBaseProposer:
534 common_attn_metadata._seq_lens_cpu = None
535 common_attn_metadata._num_computed_tokens_cpu = None
536
537 + block_size = self.block_size
538 + assert block_size > 0, "block_size has not been initialized."
539 for token_index in range(self.num_speculative_tokens - 1):
540 # Update the inputs.
541 # cast to int32 is crucial when eagle model is compiled.
542 # tensor.argmax() returns int64 by default.
543 input_ids = draft_token_ids_list[-1].int()
544 + # Use fused kernel for slot mapping and metadata updates.
545 + # Write clamped positions directly into the positions buffer to
546 + # avoid an extra D2D copy for the common (non-mrope) case.
547 + positions_1d = positions[0] if self.uses_mrope else positions
548 if self.uses_mrope:
542 - positions += 1
543 - # NOTE(woosuk): We should handle the case where the draft model
544 - # generates tokens beyond the max model length.
545 - # Since it is complex to remove such requests from the batch,
546 - # we keep them in the batch but adjust the position ids
547 - # and slot mappings to avoid the
548 - # out-of-range access during the model execution.
549 - # The draft tokens generated with this adjustment
550 - # should be ignored.
551 - exceeds_max_model_len = positions[0] >= self.max_model_len
552 - # Mask out the position ids that exceed the max model length.
553 - # Otherwise, we may get out-of-range error in RoPE.
554 - clamped_positions = torch.where(
555 - exceeds_max_model_len.unsqueeze(0),
556 - torch.zeros_like(positions),
557 - positions,
558 - )
549 + out_pos = self.mrope_positions[0, :batch_size]
550 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
551 + out_pos = self.xdrope_positions[0, :batch_size]
552 else:
560 - positions += 1
561 - exceeds_max_model_len = positions >= self.max_model_len
562 - clamped_positions = torch.where(exceeds_max_model_len, 0, positions)
563 - # For data integrity when async scheduling, we shouldn't use in place
564 - # operations in case they are modified in next step's `prepare_input`
565 - # of main model.
566 - # Increment the sequence lengths.
567 - common_attn_metadata.seq_lens += 1
568 - # For the requests that exceed the max model length, we set the
569 - # sequence length to 1 to minimize their overheads in attention.
570 - common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len, 1)
553 + out_pos = self.positions[:batch_size]
554 + eagle_step_update_slot_mapping_and_metadata(
555 + positions_1d=positions_1d,
556 + block_table_tensor=common_attn_metadata.block_table_tensor,
557 + seq_lens=common_attn_metadata.seq_lens,
558 + block_size=block_size,
559 + max_model_len=self.max_model_len,
560 + out_clamped_positions=out_pos,
561 + out_slot_mapping=self._slot_mapping_buffer[:input_batch_size],
562 + input_batch_size=input_batch_size,
563 + )
564 + common_attn_metadata.slot_mapping = self._slot_mapping_buffer[:batch_size]
565 + if self.uses_mrope:
566 + self.mrope_positions[1:, :batch_size] = self.mrope_positions[
567 + 0, :batch_size
568 + ]
569 + positions = self.mrope_positions[:, :batch_size]
570 + elif self.uses_xdrope_dim > 0 and self.draft_uses_xdrope_dim > 0:
571 + self.xdrope_positions[1:, :batch_size] = self.xdrope_positions[
572 + 0, :batch_size
573 + ]
574 + positions = self.xdrope_positions[0, :batch_size]
575 + else:
576 + positions = self.positions[:batch_size]
577 # Increment the maximum sequence length. We increment max_seq_len
578 # unconditionally even though some seq_lens may have been capped above,
579 # as max_seq_len serves as an upper bound for sequence lengths.
@@ -582,33 +588,6 @@ class SpecDecodeBaseProposer:
588 if common_attn_metadata._num_computed_tokens_cpu is not None:
589 common_attn_metadata._num_computed_tokens_cpu += 1
590
585 - # Compute the slot mapping.
586 - block_size = self.block_size
587 - assert block_size > 0, "block_size has not been initialized."
588 - if self.uses_mrope:
589 - # all dimensions of positions are the same
590 - block_numbers = clamped_positions[0] // block_size
591 - else:
592 - block_numbers = clamped_positions // block_size
593 - block_ids = common_attn_metadata.block_table_tensor.gather(
594 - dim=1, index=block_numbers.view(-1, 1)
595 - )
596 - block_ids = block_ids.view(-1)
597 - if self.uses_mrope:
598 - common_attn_metadata.slot_mapping = (
599 - block_ids * block_size + clamped_positions[0] % block_size
600 - )
601 - else:
602 - common_attn_metadata.slot_mapping = (
603 - block_ids * block_size + clamped_positions % block_size
604 - )
605 - # Mask out the slot mappings that exceed the max model length.
606 - # Otherwise, the KV cache will be inadvertently updated with the
607 - # padding tokens.
608 - common_attn_metadata.slot_mapping.masked_fill_(
609 - exceeds_max_model_len, PADDING_SLOT_ID
610 - )
611 -
591 # Rebuild attention metadata
592 for attn_group in self.draft_attn_groups:
593 attn_metadata = attn_group.get_metadata_builder().build_for_drafting(
@@ -620,7 +599,6 @@ class SpecDecodeBaseProposer:
599
600 # copy inputs to buffer for cudagraph
601 self.input_ids[:batch_size] = input_ids
623 - self._set_positions(batch_size, clamped_positions)
602 self.hidden_states[:batch_size] = hidden_states
603 if self.supports_mm_inputs:
604 self.inputs_embeds[:batch_size] = self.model.embed_input_ids(input_ids)
@@ -646,9 +624,7 @@ class SpecDecodeBaseProposer:
624 num_tokens=input_batch_size,
625 num_tokens_across_dp=batch_size_across_dp,
626 cudagraph_runtime_mode=cudagraph_runtime_mode,
649 - slot_mapping=self._get_slot_mapping(
650 - input_batch_size, common_attn_metadata.slot_mapping
651 - ),
627 + slot_mapping=self._get_slot_mapping(input_batch_size),
628 ):
629 ret_hidden_states = self.model(**model_kwargs)
630 if not self.model_returns_tuple():
vllm/v1/spec_decode/utils.py
+108
@@ -11,6 +11,114 @@ from vllm.v1.attention.backends.utils import (
11 PADDING_SLOT_ID = -1
12
13
14 +@triton.jit
15 +def eagle_step_slot_mapping_metadata_kernel(
16 + positions_ptr, # [batch_size] - current positions (1D view for M-RoPE)
17 + block_table_ptr, # [batch_size, n_blocks_per_req]
18 + block_table_stride, # stride for block_table dim 1
19 + seq_lens_ptr, # [batch_size] - read and write
20 + out_clamped_positions_ptr, # [batch_size] (output)
21 + out_slot_mapping_ptr, # [input_batch_size] (output)
22 + block_size: tl.constexpr,
23 + max_model_len: tl.constexpr,
24 + n_blocks_per_req: tl.constexpr,
25 + PAD_ID: tl.constexpr,
26 + batch_size,
27 +):
28 + """
29 + Fused kernel for EAGLE autoregressive step: updates positions, slot mapping,
30 + and sequence lengths in a single kernel to reduce launch overhead.
31 +
32 + Launched with input_batch_size threads. Threads with req_idx >= batch_size
33 + are cudagraph padding slots and only write PADDING_SLOT_ID.
34 +
35 + Each real thread handles one request in the batch. Computes:
36 + - new_position = position + 1, clamped if exceeds max_model_len
37 + - slot_mapping from block table lookup
38 + - seq_lens += 1, or 1 if position exceeds max
39 + """
40 + req_idx = tl.program_id(0)
41 +
42 + if req_idx >= batch_size:
43 + tl.store(out_slot_mapping_ptr + req_idx, PAD_ID)
44 + return
45 +
46 + # Load current position and increment
47 + position = tl.load(positions_ptr + req_idx)
48 + new_position = position + 1
49 +
50 + # Check bounds and compute clamped position
51 + exceeds_max = new_position >= max_model_len
52 + clamped_position = tl.where(exceeds_max, 0, new_position)
53 +
54 + # Block table lookup: block_number = position // block_size
55 + # Clamp block_number to avoid OOB when position is at max
56 + block_number = clamped_position // block_size
57 + block_number = tl.minimum(block_number, n_blocks_per_req - 1)
58 +
59 + block_id = tl.load(block_table_ptr + req_idx * block_table_stride + block_number)
60 + slot_id = block_id * block_size + (clamped_position % block_size)
61 + slot_id = tl.where(exceeds_max, PAD_ID, slot_id)
62 +
63 + # Update seq_lens: +1 normally, or 1 if exceeded
64 + seq_len = tl.load(seq_lens_ptr + req_idx)
65 + new_seq_len = tl.where(exceeds_max, 1, seq_len + 1)
66 + new_seq_len = tl.minimum(new_seq_len, max_model_len)
67 +
68 + # Store outputs
69 + tl.store(out_clamped_positions_ptr + req_idx, clamped_position)
70 + tl.store(out_slot_mapping_ptr + req_idx, slot_id)
71 + tl.store(seq_lens_ptr + req_idx, new_seq_len)
72 +
73 +
74 +def eagle_step_update_slot_mapping_and_metadata(
75 + positions_1d: torch.Tensor,
76 + block_table_tensor: torch.Tensor,
77 + seq_lens: torch.Tensor,
78 + block_size: int,
79 + max_model_len: int,
80 + out_clamped_positions: torch.Tensor,
81 + out_slot_mapping: torch.Tensor,
82 + input_batch_size: int | None = None,
83 +) -> None:
84 + """
85 + Fused update of slot mapping and metadata for one EAGLE autoregressive step.
86 + Updates seq_lens in place. Writes to out_clamped_positions and out_slot_mapping.
87 +
88 + When input_batch_size > batch_size, threads beyond batch_size write
89 + PADDING_SLOT_ID to out_slot_mapping for cudagraph padding.
90 +
91 + Args:
92 + positions_1d: [batch_size] current positions (use positions[0] for M-RoPE)
93 + block_table_tensor: [batch_size, n_blocks_per_req]
94 + seq_lens: [batch_size] updated in place
95 + block_size: KV cache block size
96 + max_model_len: max model length for clamping
97 + out_clamped_positions: [batch_size] output buffer for clamped positions
98 + out_slot_mapping: [input_batch_size] output buffer for slot mapping
99 + input_batch_size: total batch size including cudagraph padding;
100 + defaults to batch_size (no padding)
101 + """
102 + batch_size = positions_1d.shape[0]
103 + if input_batch_size is None:
104 + input_batch_size = batch_size
105 + n_blocks_per_req = block_table_tensor.shape[1]
106 +
107 + eagle_step_slot_mapping_metadata_kernel[(input_batch_size,)](
108 + positions_1d,
109 + block_table_tensor,
110 + block_table_tensor.stride(0),
111 + seq_lens,
112 + out_clamped_positions,
113 + out_slot_mapping,
114 + block_size=block_size,
115 + max_model_len=max_model_len,
116 + n_blocks_per_req=n_blocks_per_req,
117 + PAD_ID=PADDING_SLOT_ID,
118 + batch_size=batch_size,
119 + )
120 +
121 +
122 @triton.jit
123 def eagle_prepare_inputs_padded_kernel(
124 cu_num_draft_tokens_ptr, # [num_reqs]