| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import Any |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | from vllm.model_executor.warmup.jit_warmup import ( |
| 8 | WarmupChoices, |
| 9 | WarmupIntRange, |
| 10 | ) |
| 11 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 12 | DispatchSpec, |
| 13 | TritonWarmupTensor, |
| 14 | triton_kernel_dispatcher_with_warmup, |
| 15 | ) |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.triton_utils import tl, triton |
| 18 | from vllm.v1.attention.backends.utils import ( |
| 19 | CommonAttentionMetadata, |
| 20 | ) |
| 21 | |
| 22 | PADDING_SLOT_ID = -1 |
| 23 | |
| 24 | |
| 25 | def next_power_of_2(n: int) -> int: |
| 26 | """Return the smallest power of 2 >= n.""" |
| 27 | if n <= 0: |
| 28 | return 1 |
| 29 | n -= 1 |
| 30 | n |= n >> 1 |
| 31 | n |= n >> 2 |
| 32 | n |= n >> 4 |
| 33 | n |= n >> 8 |
| 34 | n |= n >> 16 |
| 35 | n |= n >> 32 |
| 36 | return n + 1 |
| 37 | |
| 38 | |
| 39 | @triton.jit |
| 40 | def eagle_step_slot_mapping_metadata_kernel( |
| 41 | positions_ptr, # [batch_size] - current positions (1D view for M-RoPE) |
| 42 | block_table_ptr, # [batch_size, n_blocks_per_req] |
| 43 | block_table_stride, # stride for block_table dim 1 |
| 44 | seq_lens_ptr, # [batch_size] - read and write |
| 45 | out_clamped_positions_ptr, # [batch_size] (output) |
| 46 | out_slot_mapping_ptr, # [input_batch_size] (output) |
| 47 | block_size: tl.constexpr, |
| 48 | max_model_len: tl.constexpr, |
| 49 | n_blocks_per_req: tl.constexpr, |
| 50 | PAD_ID: tl.constexpr, |
| 51 | batch_size, |
| 52 | ): |
| 53 | """Fused kernel for EAGLE autoregressive step: updates positions, slot mapping, |
| 54 | and sequence lengths in a single kernel to reduce launch overhead. |
| 55 | |
| 56 | Launched with input_batch_size threads. Threads with req_idx >= batch_size |
| 57 | are cudagraph padding slots and only write PADDING_SLOT_ID. |
| 58 | |
| 59 | Each real thread handles one request in the batch. Computes: |
| 60 | - new_position = position + 1, clamped if exceeds max_model_len |
| 61 | - slot_mapping from block table lookup |
| 62 | - seq_lens += 1, or 1 if position exceeds max |
| 63 | """ |
| 64 | req_idx = tl.program_id(0) |
| 65 | |
| 66 | if req_idx >= batch_size: |
| 67 | tl.store(out_slot_mapping_ptr + req_idx, PAD_ID) |
| 68 | return |
| 69 | |
| 70 | # Load current position and increment |
| 71 | position = tl.load(positions_ptr + req_idx) |
| 72 | new_position = position + 1 |
| 73 | |
| 74 | # Check bounds and compute clamped position |
| 75 | exceeds_max = new_position >= max_model_len |
| 76 | clamped_position = tl.where(exceeds_max, 0, new_position) |
| 77 | |
| 78 | # Block table lookup: block_number = position // block_size |
| 79 | # Clamp block_number to avoid OOB when position is at max |
| 80 | block_number = clamped_position // block_size |
| 81 | block_number = tl.minimum(block_number, n_blocks_per_req - 1) |
| 82 | |
| 83 | block_id = tl.load(block_table_ptr + req_idx * block_table_stride + block_number) |
| 84 | slot_id = block_id * block_size + (clamped_position % block_size) |
| 85 | slot_id = tl.where(exceeds_max, PAD_ID, slot_id) |
| 86 | |
| 87 | # Update seq_lens: +1 normally, or 1 if exceeded |
| 88 | seq_len = tl.load(seq_lens_ptr + req_idx) |
| 89 | new_seq_len = tl.where(exceeds_max, 1, seq_len + 1) |
| 90 | new_seq_len = tl.minimum(new_seq_len, max_model_len) |
| 91 | |
| 92 | # Store outputs |
| 93 | tl.store(out_clamped_positions_ptr + req_idx, clamped_position) |
| 94 | tl.store(out_slot_mapping_ptr + req_idx, slot_id) |
| 95 | tl.store(seq_lens_ptr + req_idx, new_seq_len) |
| 96 | |
| 97 | |
| 98 | def _eagle_step_warmup_inputs(*, max_model_len: int, max_batch_size: int): |
| 99 | block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2) |
| 100 | n_blocks_per_req = triton.cdiv(max_model_len, block_size) |
| 101 | return dict( |
| 102 | positions=TritonWarmupTensor(torch.int64), |
| 103 | block_table=TritonWarmupTensor( |
| 104 | torch.int32, shape=(1, n_blocks_per_req), strides=(n_blocks_per_req, 1) |
| 105 | ), |
| 106 | block_table_stride=n_blocks_per_req, |
| 107 | seq_lens=TritonWarmupTensor(torch.int32), |
| 108 | out_clamped_positions=TritonWarmupTensor(torch.int64), |
| 109 | out_slot_mapping=TritonWarmupTensor(torch.int64), |
| 110 | block_size=block_size, |
| 111 | max_model_len=max_model_len, |
| 112 | n_blocks_per_req=n_blocks_per_req, |
| 113 | PAD_ID=PADDING_SLOT_ID, |
| 114 | batch_size=WarmupIntRange(1, max_batch_size + 1), |
| 115 | input_batch_size=1, |
| 116 | ) |
| 117 | |
| 118 | |
| 119 | @triton_kernel_dispatcher_with_warmup( |
| 120 | kernel=eagle_step_slot_mapping_metadata_kernel, |
| 121 | warmup_inputs=_eagle_step_warmup_inputs, |
| 122 | ) |
| 123 | def _eagle_step_slot_mapping_metadata( |
| 124 | positions: torch.Tensor, |
| 125 | block_table: torch.Tensor, |
| 126 | block_table_stride: int, |
| 127 | seq_lens: torch.Tensor, |
| 128 | out_clamped_positions: torch.Tensor, |
| 129 | out_slot_mapping: torch.Tensor, |
| 130 | *, |
| 131 | block_size: int, |
| 132 | max_model_len: int, |
| 133 | n_blocks_per_req: int, |
| 134 | PAD_ID: int, |
| 135 | batch_size: int, |
| 136 | input_batch_size: int, |
| 137 | ) -> DispatchSpec: |
| 138 | return (input_batch_size,), {} |
| 139 | |
| 140 | |
| 141 | def eagle_step_update_slot_mapping_and_metadata( |
| 142 | positions_1d: torch.Tensor, |
| 143 | block_table_tensor: torch.Tensor, |
| 144 | seq_lens: torch.Tensor, |
| 145 | block_size: int, |
| 146 | max_model_len: int, |
| 147 | out_clamped_positions: torch.Tensor, |
| 148 | out_slot_mapping: torch.Tensor, |
| 149 | input_batch_size: int | None = None, |
| 150 | ) -> None: |
| 151 | """Fused update of slot mapping and metadata for one EAGLE autoregressive step. |
| 152 | Updates seq_lens in place. Writes to out_clamped_positions and out_slot_mapping. |
| 153 | |
| 154 | When input_batch_size > batch_size, threads beyond batch_size write |
| 155 | PADDING_SLOT_ID to out_slot_mapping for cudagraph padding. |
| 156 | |
| 157 | Args: |
| 158 | positions_1d: [batch_size] current positions (use positions[0] for M-RoPE) |
| 159 | block_table_tensor: [batch_size, n_blocks_per_req] |
| 160 | seq_lens: [batch_size] updated in place |
| 161 | block_size: KV cache block size |
| 162 | max_model_len: max model length for clamping |
| 163 | out_clamped_positions: [batch_size] output buffer for clamped positions |
| 164 | out_slot_mapping: [input_batch_size] output buffer for slot mapping |
| 165 | input_batch_size: total batch size including cudagraph padding; |
| 166 | defaults to batch_size (no padding) |
| 167 | |
| 168 | """ |
| 169 | batch_size = positions_1d.shape[0] |
| 170 | if input_batch_size is None: |
| 171 | input_batch_size = batch_size |
| 172 | |
| 173 | n_blocks_per_req = block_table_tensor.shape[1] |
| 174 | _eagle_step_slot_mapping_metadata( |
| 175 | positions_1d, |
| 176 | block_table_tensor, |
| 177 | block_table_tensor.stride(0), |
| 178 | seq_lens, |
| 179 | out_clamped_positions, |
| 180 | out_slot_mapping, |
| 181 | block_size=block_size, |
| 182 | max_model_len=max_model_len, |
| 183 | n_blocks_per_req=n_blocks_per_req, |
| 184 | PAD_ID=PADDING_SLOT_ID, |
| 185 | batch_size=batch_size, |
| 186 | input_batch_size=input_batch_size, |
| 187 | ) |
| 188 | |
| 189 | |
| 190 | @triton.jit |
| 191 | def eagle_prepare_inputs_padded_kernel( |
| 192 | cu_num_draft_tokens_ptr, # [num_reqs] |
| 193 | valid_sampled_tokens_count_ptr, # [num_reqs] |
| 194 | query_start_loc_gpu_ptr, # [num_reqs + 1] |
| 195 | token_indices_to_sample_ptr, # [num_reqs] (output) |
| 196 | num_rejected_tokens_gpu_ptr, # [num_reqs] (output) |
| 197 | num_reqs, # tl.int32 |
| 198 | ): |
| 199 | """Fused kernel for Eagle prepare_input_padded. This kernel computes the |
| 200 | token index to sample for each request, taking into account the number |
| 201 | of draft tokens and the number of valid sampled tokens (which is one more than |
| 202 | the number of accepted tokens). |
| 203 | """ |
| 204 | req_idx = tl.program_id(axis=0) |
| 205 | if req_idx >= num_reqs: |
| 206 | return |
| 207 | |
| 208 | # Calculate num_draft_tokens from cu_num_draft_tokens, which is an inclusive |
| 209 | # cumulative sum (first entry is the first value, not zero). |
| 210 | cu_draft_curr = tl.load(cu_num_draft_tokens_ptr + req_idx) |
| 211 | |
| 212 | if req_idx == 0: |
| 213 | num_draft_tokens = cu_draft_curr |
| 214 | else: |
| 215 | cu_draft_prev = tl.load(cu_num_draft_tokens_ptr + req_idx - 1) |
| 216 | num_draft_tokens = cu_draft_curr - cu_draft_prev |
| 217 | |
| 218 | valid_count = tl.load(valid_sampled_tokens_count_ptr + req_idx) |
| 219 | num_rejected_tokens = num_draft_tokens + 1 - valid_count |
| 220 | num_rejected_tokens = tl.where(num_draft_tokens > 0, num_rejected_tokens, 0) |
| 221 | |
| 222 | # query_start_loc[req_idx + 1] is the start position of the next request, |
| 223 | # which is one past the last token of this request. |
| 224 | q_last_tok_idx = tl.load(query_start_loc_gpu_ptr + req_idx + 1) - 1 |
| 225 | |
| 226 | index_to_sample = q_last_tok_idx - num_rejected_tokens |
| 227 | tl.store(token_indices_to_sample_ptr + req_idx, index_to_sample) |
| 228 | tl.store(num_rejected_tokens_gpu_ptr + req_idx, num_rejected_tokens) |
| 229 | |
| 230 | |
| 231 | def _eagle_prepare_inputs_warmup_inputs(*, max_batch_size: int): |
| 232 | int32 = TritonWarmupTensor(torch.int32) |
| 233 | return dict( |
| 234 | cu_num_draft_tokens=int32, |
| 235 | valid_sampled_tokens_count=int32, |
| 236 | query_start_loc_gpu=int32, |
| 237 | token_indices_to_sample=int32, |
| 238 | num_rejected_tokens_gpu=int32, |
| 239 | num_reqs=WarmupIntRange(1, max_batch_size + 1), |
| 240 | ) |
| 241 | |
| 242 | |
| 243 | @triton_kernel_dispatcher_with_warmup( |
| 244 | kernel=eagle_prepare_inputs_padded_kernel, |
| 245 | warmup_inputs=_eagle_prepare_inputs_warmup_inputs, |
| 246 | ) |
| 247 | def _eagle_prepare_inputs_padded( |
| 248 | cu_num_draft_tokens: torch.Tensor, |
| 249 | valid_sampled_tokens_count: torch.Tensor, |
| 250 | query_start_loc_gpu: torch.Tensor, |
| 251 | token_indices_to_sample: torch.Tensor, |
| 252 | num_rejected_tokens_gpu: torch.Tensor, |
| 253 | num_reqs: int, |
| 254 | ) -> DispatchSpec: |
| 255 | return (num_reqs,), {} |
| 256 | |
| 257 | |
| 258 | @triton.jit |
| 259 | def eagle_prepare_next_token_padded_kernel( |
| 260 | sampled_token_ids_ptr, # [num_reqs, num_sampled_tokens_per_req] |
| 261 | discard_request_mask_ptr, # [num_reqs] |
| 262 | backup_next_token_ids_ptr, # [num_reqs] |
| 263 | next_token_ids_ptr, # [num_reqs] (output) |
| 264 | valid_sampled_tokens_count_ptr, # [num_reqs] (output) |
| 265 | vocab_size, # tl.int32 |
| 266 | num_sampled_tokens_per_req, # tl.int32 (num_spec_tokens + 1) |
| 267 | num_reqs, # tl.int32 |
| 268 | stride_sampled_token_ids, # tl.int32 (stride for dim 0) |
| 269 | BLOCK_SIZE_TOKENS: tl.constexpr, # Power-of-2 >= num_sampled_tokens_per_req |
| 270 | ): |
| 271 | """Fused kernel for Eagle prepare_next_token_ids_padded. This kernel computes the |
| 272 | number of valid (1 + accepted) tokens for each request, and the corresponding |
| 273 | "next" token id to sample from during speculative decoding. This is the |
| 274 | "last accepted token" from the sampled tokens, or the backup token if no |
| 275 | tokens were accepted or if the request is marked as discarded. |
| 276 | """ |
| 277 | req_idx = tl.program_id(axis=0) |
| 278 | if req_idx >= num_reqs: |
| 279 | return |
| 280 | |
| 281 | # Check if this request is discarded. |
| 282 | is_discarded = tl.load(discard_request_mask_ptr + req_idx) |
| 283 | |
| 284 | if is_discarded: |
| 285 | backup_token = tl.load(backup_next_token_ids_ptr + req_idx) |
| 286 | valid_count = tl.full((), 0, dtype=tl.uint32) |
| 287 | tl.store(next_token_ids_ptr + req_idx, backup_token) |
| 288 | tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count) |
| 289 | else: |
| 290 | # Count the number of valid tokens among the sampled tokens. |
| 291 | token_offs = tl.arange(0, BLOCK_SIZE_TOKENS) |
| 292 | token_mask = token_offs < num_sampled_tokens_per_req |
| 293 | |
| 294 | row_ptr = sampled_token_ids_ptr + req_idx * stride_sampled_token_ids |
| 295 | token_ids = tl.load(row_ptr + token_offs, mask=token_mask, other=-1) |
| 296 | |
| 297 | # Rejected tokens are -1, valid tokens are in [0, vocab_size) |
| 298 | is_valid_mask = (token_ids != -1) & (token_ids < vocab_size) & token_mask |
| 299 | valid_count = tl.sum(is_valid_mask) |
| 300 | |
| 301 | if valid_count > 0: |
| 302 | # Guaranteed to be well-defined since |
| 303 | # valid_count > 0 implies is_valid_mask is not empty |
| 304 | last_valid_index = tl.max(tl.where(is_valid_mask, token_offs, -1)) |
| 305 | |
| 306 | # Select the token at that index, using a sum trick since |
| 307 | # we don't want to load again to access token_ids[last_valid_index]. |
| 308 | last_valid_token = tl.sum( |
| 309 | tl.where(token_offs == last_valid_index, token_ids, 0) |
| 310 | ) |
| 311 | tl.store(next_token_ids_ptr + req_idx, last_valid_token) |
| 312 | else: |
| 313 | # No valid tokens found, use backup token |
| 314 | backup_token = tl.load(backup_next_token_ids_ptr + req_idx) |
| 315 | tl.store(next_token_ids_ptr + req_idx, backup_token) |
| 316 | |
| 317 | tl.store(valid_sampled_tokens_count_ptr + req_idx, valid_count) |
| 318 | |
| 319 | |
| 320 | def _eagle_prepare_next_token_warmup_inputs( |
| 321 | *, vocab_size: int, num_sampled_tokens_per_req: int, max_batch_size: int |
| 322 | ): |
| 323 | int32 = TritonWarmupTensor(torch.int32) |
| 324 | sampled_width: Any = WarmupIntRange(1, num_sampled_tokens_per_req + 1) |
| 325 | return dict( |
| 326 | sampled_token_ids=TritonWarmupTensor( |
| 327 | torch.int32, |
| 328 | shape=(1, sampled_width), |
| 329 | strides=(sampled_width, 1), |
| 330 | ), |
| 331 | discard_request_mask=TritonWarmupTensor(torch.bool), |
| 332 | backup_next_token_ids=int32, |
| 333 | next_token_ids=int32, |
| 334 | valid_sampled_tokens_count=int32, |
| 335 | vocab_size=vocab_size, |
| 336 | num_sampled_tokens_per_req=sampled_width, |
| 337 | num_reqs=WarmupIntRange(1, max_batch_size + 1), |
| 338 | stride_sampled_token_ids=sampled_width, |
| 339 | BLOCK_SIZE_TOKENS=next_power_of_2(sampled_width), |
| 340 | ) |
| 341 | |
| 342 | |
| 343 | @triton_kernel_dispatcher_with_warmup( |
| 344 | kernel=eagle_prepare_next_token_padded_kernel, |
| 345 | warmup_inputs=_eagle_prepare_next_token_warmup_inputs, |
| 346 | ) |
| 347 | def _eagle_prepare_next_token_padded( |
| 348 | sampled_token_ids: torch.Tensor, |
| 349 | discard_request_mask: torch.Tensor, |
| 350 | backup_next_token_ids: torch.Tensor, |
| 351 | next_token_ids: torch.Tensor, |
| 352 | valid_sampled_tokens_count: torch.Tensor, |
| 353 | vocab_size: int, |
| 354 | num_sampled_tokens_per_req: int, |
| 355 | num_reqs: int, |
| 356 | stride_sampled_token_ids: int, |
| 357 | *, |
| 358 | BLOCK_SIZE_TOKENS: int, |
| 359 | ) -> DispatchSpec: |
| 360 | return (num_reqs,), {} |
| 361 | |
| 362 | |
| 363 | def compute_new_slot_mapping( |
| 364 | cad: CommonAttentionMetadata, |
| 365 | new_positions: torch.Tensor, |
| 366 | is_rejected_token_mask: torch.Tensor, |
| 367 | block_size: int, |
| 368 | num_new_tokens: int, |
| 369 | max_model_len: int, |
| 370 | ): |
| 371 | batch_size, n_blocks_per_req = cad.block_table_tensor.shape |
| 372 | req_indices = torch.arange(batch_size, device=cad.query_start_loc.device) |
| 373 | req_indices = torch.repeat_interleave( |
| 374 | req_indices, |
| 375 | cad.naive_query_lens() + num_new_tokens, |
| 376 | output_size=len(new_positions), |
| 377 | ) |
| 378 | # Clamp the positions to prevent an out-of-bounds error when indexing |
| 379 | # into block_table_tensor. |
| 380 | clamped_positions = torch.clamp(new_positions, max=max_model_len - 1) |
| 381 | block_table_indices = ( |
| 382 | req_indices * n_blocks_per_req + clamped_positions // block_size |
| 383 | ) |
| 384 | block_nums = cad.block_table_tensor.view(-1)[block_table_indices] |
| 385 | block_offsets = clamped_positions % block_size |
| 386 | new_slot_mapping = block_nums * block_size + block_offsets |
| 387 | # Mask out the position ids that exceed the max model length. |
| 388 | exceeds_max_model_len = new_positions >= max_model_len |
| 389 | new_slot_mapping.masked_fill_(exceeds_max_model_len, PADDING_SLOT_ID) |
| 390 | # Mask out rejected tokens to prevent saves to the KV cache. |
| 391 | new_slot_mapping.masked_fill_(is_rejected_token_mask, PADDING_SLOT_ID) |
| 392 | return new_slot_mapping |
| 393 | |
| 394 | |
| 395 | def extend_all_queries_by_N( |
| 396 | common_attn_metadata: CommonAttentionMetadata, |
| 397 | N: int, |
| 398 | arange: torch.Tensor, |
| 399 | new_slot_mapping: torch.Tensor, |
| 400 | ) -> CommonAttentionMetadata: |
| 401 | """Creates a new CommonAttentionMetadata with all query lengths increased by N. |
| 402 | Also all seq lens are increased by N. |
| 403 | This is useful e.g. in speculative decoding with parallel drafting, where we |
| 404 | extend each sequence by N tokens and predict all tokens in one pass. |
| 405 | The slot mapping is computed externally, as it requires more information. |
| 406 | """ |
| 407 | cad = common_attn_metadata |
| 408 | # query start loc must be increased by [+0, +N, +2N, ..., +batch_size * N] |
| 409 | new_query_start_loc = cad.query_start_loc + N * arange[: len(cad.query_start_loc)] |
| 410 | new_query_start_loc_cpu = cad.query_start_loc_cpu + N * torch.arange( |
| 411 | len(cad.query_start_loc_cpu), dtype=torch.int32 |
| 412 | ) |
| 413 | new_cad = cad.replace( |
| 414 | query_start_loc=new_query_start_loc, |
| 415 | query_start_loc_cpu=new_query_start_loc_cpu, |
| 416 | seq_lens=cad.seq_lens + N, |
| 417 | # each request is extended by N tokens -> batch_size * N tokens are added |
| 418 | num_actual_tokens=cad.num_actual_tokens + cad.batch_size() * N, |
| 419 | # All query lens increase by N, so max query len increases by N |
| 420 | max_query_len=cad.max_query_len + N, |
| 421 | max_seq_len=cad.max_seq_len + N, |
| 422 | slot_mapping=new_slot_mapping, |
| 423 | ) |
| 424 | return new_cad |
| 425 | |
| 426 | |
| 427 | # Unified copy/expand kernel |
| 428 | @triton.jit |
| 429 | def copy_and_expand_eagle_inputs_kernel( |
| 430 | # (Padded) Inputs from the target model |
| 431 | target_token_ids_ptr, # [total_tokens_in_batch] |
| 432 | target_positions_ptr, # [total_tokens_in_batch] |
| 433 | next_token_ids_ptr, # [num_reqs] |
| 434 | # Outputs to the drafting buffers |
| 435 | out_input_ids_ptr, # [total_draft_tokens_in_batch] (output) |
| 436 | out_positions_ptr, # [total_draft_tokens_in_batch] (output) |
| 437 | out_is_rejected_token_mask_ptr, # [total_draft_tokens_in_batch] (output) |
| 438 | out_is_masked_token_mask_ptr, # [total_draft_tokens_in_batch] (output) |
| 439 | out_new_token_indices_ptr, # [num_padding_slots_per_request * num_reqs] (output) |
| 440 | out_hidden_state_mapping_ptr, # [total_tokens_in_batch] |
| 441 | # Input metadata |
| 442 | query_start_loc_ptr, # [num_reqs + 1], last value is the total num input tokens |
| 443 | query_end_loc_ptr, # [num_reqs] |
| 444 | padding_token_id, # tl.int32 |
| 445 | parallel_drafting_token_id, # tl.int32 |
| 446 | # Sizing info |
| 447 | total_input_tokens, # tl.int32 |
| 448 | num_padding_slots_per_request, # tl.int32 |
| 449 | shift_input_ids, # tl.bool |
| 450 | BLOCK_SIZE_TOKENS: tl.constexpr, # Blocks along token dim to handle prefills |
| 451 | ): |
| 452 | """Copy and expand inputs from the target model to the drafting buffers for Eagle |
| 453 | speculative decoding. This kernel handles padding slots and parallel drafting |
| 454 | tokens, if enabled. |
| 455 | """ |
| 456 | request_idx = tl.program_id(axis=0) |
| 457 | token_batch_idx = tl.program_id(axis=1) |
| 458 | |
| 459 | # Load query locations |
| 460 | query_start_loc = tl.load(query_start_loc_ptr + request_idx) |
| 461 | next_query_start_loc = tl.load(query_start_loc_ptr + request_idx + 1) |
| 462 | query_end_loc = tl.load(query_end_loc_ptr + request_idx) |
| 463 | |
| 464 | # Calculate number of valid tokens to copy and input offset |
| 465 | # With shift_input_ids=True, we skip the first token |
| 466 | # Output layout: each request gets (input_len + num_padding_slots_per_request) slots |
| 467 | # But with shift, we lose one token per request |
| 468 | if shift_input_ids: |
| 469 | num_valid_tokens = query_end_loc - query_start_loc |
| 470 | input_offset = 1 |
| 471 | output_start = query_start_loc + request_idx * ( |
| 472 | num_padding_slots_per_request - 1 |
| 473 | ) |
| 474 | else: |
| 475 | num_valid_tokens = query_end_loc - query_start_loc + 1 |
| 476 | input_offset = 0 |
| 477 | output_start = query_start_loc + request_idx * num_padding_slots_per_request |
| 478 | |
| 479 | # Number of rejected tokens from previous speculation |
| 480 | num_rejected = next_query_start_loc - query_end_loc - 1 |
| 481 | |
| 482 | # Total output tokens for this request |
| 483 | total_output_tokens = ( |
| 484 | num_valid_tokens + num_padding_slots_per_request + num_rejected |
| 485 | ) |
| 486 | |
| 487 | # Process tokens in this block |
| 488 | j = token_batch_idx * BLOCK_SIZE_TOKENS + tl.arange(0, BLOCK_SIZE_TOKENS) |
| 489 | |
| 490 | # Compute masks for different output regions: |
| 491 | # [0, num_valid_tokens): valid tokens copied from input |
| 492 | # [num_valid_tokens]: bonus token from next_token_ids |
| 493 | # (num_valid_tokens, num_valid_tokens + num_padding_slots_per_request): |
| 494 | # parallel drafting slots |
| 495 | # [num_valid_tokens + num_padding_slots_per_request, total_output_tokens): |
| 496 | # rejected slots |
| 497 | in_bounds = j < total_output_tokens |
| 498 | is_valid_region = j < num_valid_tokens |
| 499 | is_bonus_region = j == num_valid_tokens |
| 500 | is_parallel_draft_region = (j > num_valid_tokens) & ( |
| 501 | j < num_valid_tokens + num_padding_slots_per_request |
| 502 | ) |
| 503 | is_rejected_region = j >= num_valid_tokens + num_padding_slots_per_request |
| 504 | |
| 505 | # Compute output indices |
| 506 | out_idx = output_start + j |
| 507 | |
| 508 | # For valid tokens, compute input index |
| 509 | in_idx = query_start_loc + input_offset + j |
| 510 | # Clamp to avoid out-of-bounds access (masked loads still need valid addresses) |
| 511 | in_idx_clamped = tl.minimum(in_idx, total_input_tokens - 1) |
| 512 | |
| 513 | # Load input tokens (masked to valid region) |
| 514 | token_ids = tl.load( |
| 515 | target_token_ids_ptr + in_idx_clamped, mask=is_valid_region & in_bounds, other=0 |
| 516 | ) |
| 517 | |
| 518 | # Load the starting position for this request (first position in the sequence) |
| 519 | start_pos = tl.load(target_positions_ptr + query_start_loc) |
| 520 | |
| 521 | # Load bonus token for this request |
| 522 | bonus_token = tl.load(next_token_ids_ptr + request_idx) |
| 523 | |
| 524 | # Build final token_ids based on region |
| 525 | token_ids = tl.where(is_bonus_region, bonus_token, token_ids) |
| 526 | token_ids = tl.where( |
| 527 | is_parallel_draft_region, parallel_drafting_token_id, token_ids |
| 528 | ) |
| 529 | token_ids = tl.where(is_rejected_region, padding_token_id, token_ids) |
| 530 | |
| 531 | # Build final positions: |
| 532 | # Positions are NOT shifted - they start from the first input position and increment |
| 533 | # Output position j gets start_pos + j |
| 534 | # (e.g., input positions [5,6,7] -> output [5,6,7,8,9,...]) |
| 535 | positions = start_pos + j |
| 536 | # Rejected positions are don't-care, set to 0 |
| 537 | positions = tl.where(is_rejected_region, 0, positions) |
| 538 | |
| 539 | # Compute output masks |
| 540 | is_rejected_out = is_rejected_region & in_bounds |
| 541 | is_masked_out = is_parallel_draft_region & in_bounds |
| 542 | |
| 543 | # Compute indices of new tokens (bonus + parallel drafting) for sampling |
| 544 | # New tokens are at positions |
| 545 | # [num_valid_tokens, num_valid_tokens + num_padding_slots_per_request) |
| 546 | is_new_token_region = (j >= num_valid_tokens) & ( |
| 547 | j < num_valid_tokens + num_padding_slots_per_request |
| 548 | ) |
| 549 | new_token_local_idx = ( |
| 550 | j - num_valid_tokens |
| 551 | ) # 0 for bonus, 1, 2, ... for parallel drafting |
| 552 | new_token_out_idx = ( |
| 553 | request_idx * num_padding_slots_per_request + new_token_local_idx |
| 554 | ) |
| 555 | |
| 556 | # Compute hidden state mapping (source index -> destination index) |
| 557 | # This maps each input position to its corresponding output position |
| 558 | # Hidden states don't get shifted, so we map all input tokens (including rejected) |
| 559 | if shift_input_ids: |
| 560 | num_input_tokens_this_request = next_query_start_loc - query_start_loc |
| 561 | is_input_region = j < num_input_tokens_this_request |
| 562 | src_idx = query_start_loc + j |
| 563 | tl.store(out_hidden_state_mapping_ptr + src_idx, out_idx, mask=is_input_region) |
| 564 | |
| 565 | # Store outputs |
| 566 | tl.store(out_input_ids_ptr + out_idx, token_ids, mask=in_bounds) |
| 567 | tl.store(out_positions_ptr + out_idx, positions, mask=in_bounds) |
| 568 | tl.store(out_is_rejected_token_mask_ptr + out_idx, is_rejected_out, mask=in_bounds) |
| 569 | tl.store(out_is_masked_token_mask_ptr + out_idx, is_masked_out, mask=in_bounds) |
| 570 | tl.store( |
| 571 | out_new_token_indices_ptr + new_token_out_idx, |
| 572 | out_idx, |
| 573 | mask=is_new_token_region & in_bounds, |
| 574 | ) |
| 575 | |
| 576 | |
| 577 | def _copy_and_expand_eagle_warmup_inputs( |
| 578 | *, |
| 579 | parallel_drafting_token_id: int, |
| 580 | num_padding_slots_per_request: int, |
| 581 | shift_input_ids: bool, |
| 582 | max_num_tokens: int, |
| 583 | ): |
| 584 | int32 = TritonWarmupTensor(torch.int32) |
| 585 | int64 = TritonWarmupTensor(torch.int64) |
| 586 | boolean = TritonWarmupTensor(torch.bool) |
| 587 | return dict( |
| 588 | target_token_ids_ptr=int32, |
| 589 | target_positions_ptr=int64, |
| 590 | next_token_ids_ptr=int32, |
| 591 | out_input_ids_ptr=int32, |
| 592 | out_positions_ptr=int64, |
| 593 | out_is_rejected_token_mask_ptr=boolean, |
| 594 | out_is_masked_token_mask_ptr=boolean, |
| 595 | out_new_token_indices_ptr=int32, |
| 596 | out_hidden_state_mapping_ptr=int32, |
| 597 | query_start_loc_ptr=int32, |
| 598 | query_end_loc_ptr=int32, |
| 599 | padding_token_id=0, |
| 600 | parallel_drafting_token_id=parallel_drafting_token_id, |
| 601 | total_input_tokens=WarmupIntRange(1, max_num_tokens + 1), |
| 602 | num_padding_slots_per_request=num_padding_slots_per_request, |
| 603 | shift_input_ids=shift_input_ids, |
| 604 | BLOCK_SIZE_TOKENS=WarmupIntRange(1, 257, advance=lambda value: value * 2), |
| 605 | batch_size=1, |
| 606 | num_blocks=1, |
| 607 | ) |
| 608 | |
| 609 | |
| 610 | @triton_kernel_dispatcher_with_warmup( |
| 611 | kernel=copy_and_expand_eagle_inputs_kernel, |
| 612 | warmup_inputs=_copy_and_expand_eagle_warmup_inputs, |
| 613 | ) |
| 614 | def _copy_and_expand_eagle_inputs( |
| 615 | target_token_ids_ptr: torch.Tensor, |
| 616 | target_positions_ptr: torch.Tensor, |
| 617 | next_token_ids_ptr: torch.Tensor, |
| 618 | out_input_ids_ptr: torch.Tensor, |
| 619 | out_positions_ptr: torch.Tensor, |
| 620 | out_is_rejected_token_mask_ptr: torch.Tensor, |
| 621 | out_is_masked_token_mask_ptr: torch.Tensor, |
| 622 | out_new_token_indices_ptr: torch.Tensor, |
| 623 | out_hidden_state_mapping_ptr: torch.Tensor, |
| 624 | query_start_loc_ptr: torch.Tensor, |
| 625 | query_end_loc_ptr: torch.Tensor, |
| 626 | padding_token_id: int, |
| 627 | parallel_drafting_token_id: int, |
| 628 | total_input_tokens: int, |
| 629 | num_padding_slots_per_request: int, |
| 630 | shift_input_ids: bool, |
| 631 | *, |
| 632 | BLOCK_SIZE_TOKENS: int, |
| 633 | batch_size: int, |
| 634 | num_blocks: int, |
| 635 | ) -> DispatchSpec: |
| 636 | return (batch_size, num_blocks), {} |
| 637 | |
| 638 | |
| 639 | @triton.jit |
| 640 | def copy_and_expand_dflash_inputs_kernel( |
| 641 | # Inputs |
| 642 | next_token_ids_ptr, # [num_reqs] |
| 643 | target_positions_ptr, # [num_context] |
| 644 | # Outputs |
| 645 | out_input_ids_ptr, # [num_query_total] (output) |
| 646 | out_context_positions_ptr, # [num_context] (output) |
| 647 | out_query_positions_ptr, # [num_query_total] (output) |
| 648 | out_context_slot_mapping_ptr, # [num_context] (output) |
| 649 | out_query_slot_mapping_ptr, # [num_query_total] (output) |
| 650 | out_token_indices_ptr, # [num_reqs * num_speculative_tokens] (output) |
| 651 | # Block table |
| 652 | block_table_ptr, # [max_reqs, max_blocks] |
| 653 | block_table_stride, # stride of block_table dim 0 (in elements) |
| 654 | # Metadata |
| 655 | query_start_loc_ptr, # [num_reqs + 1] |
| 656 | num_rejected_tokens_ptr, # [num_reqs] or null (0) when not padded |
| 657 | # Scalars |
| 658 | parallel_drafting_token_id, # tl.int32 |
| 659 | block_size, # tl.int32 |
| 660 | num_query_per_req, # tl.int32 |
| 661 | num_speculative_tokens, # tl.int32 |
| 662 | total_input_tokens, # tl.int32 |
| 663 | BLOCK_SIZE: tl.constexpr, |
| 664 | HAS_NUM_REJECTED: tl.constexpr = False, |
| 665 | ): |
| 666 | """Fused kernel for DFlash first-pass input setup. |
| 667 | |
| 668 | Per request, this kernel: |
| 669 | 1. Copies context positions from target_positions to |
| 670 | out_context_positions. |
| 671 | 2. Computes query positions (last_target_pos + 1 + offset) and writes |
| 672 | them to out_query_positions. |
| 673 | 3. Writes input_ids for query tokens: [next_token, mask, mask, ...]. |
| 674 | 4. Computes slot_mapping for context and query positions into separate |
| 675 | buffers via block_table lookup. |
| 676 | 5. Writes token_indices_to_sample for the mask (speculative) tokens. |
| 677 | """ |
| 678 | req_idx = tl.program_id(axis=0) |
| 679 | block_idx = tl.program_id(axis=1) |
| 680 | |
| 681 | # Load context token range for this request |
| 682 | ctx_start = tl.load(query_start_loc_ptr + req_idx) |
| 683 | ctx_end = tl.load(query_start_loc_ptr + req_idx + 1) |
| 684 | num_ctx = ctx_end - ctx_start |
| 685 | total_tokens = num_ctx + num_query_per_req |
| 686 | |
| 687 | j = block_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 688 | in_bounds = j < total_tokens |
| 689 | is_ctx = j < num_ctx |
| 690 | is_query = (~is_ctx) & in_bounds |
| 691 | query_off = j - num_ctx # offset within query portion (0-indexed) |
| 692 | |
| 693 | # --- Positions --- |
| 694 | # Context: load from target_positions |
| 695 | ctx_pos_idx = tl.minimum(ctx_start + j, total_input_tokens - 1) |
| 696 | ctx_pos = tl.load(target_positions_ptr + ctx_pos_idx, mask=is_ctx, other=0) |
| 697 | |
| 698 | # Query: last_valid_pos + 1 + query_off |
| 699 | # In padded mode, ctx_end includes rejected tokens; use valid_ctx_end |
| 700 | # to find the last accepted context position. |
| 701 | if HAS_NUM_REJECTED: |
| 702 | num_rejected = tl.load(num_rejected_tokens_ptr + req_idx) |
| 703 | valid_ctx_end = ctx_end - num_rejected |
| 704 | else: |
| 705 | valid_ctx_end = ctx_end |
| 706 | last_pos = tl.load(target_positions_ptr + valid_ctx_end - 1) |
| 707 | query_pos = last_pos + 1 + query_off |
| 708 | |
| 709 | positions = tl.where(is_ctx, ctx_pos, query_pos) |
| 710 | |
| 711 | # Context and query positions go to separate buffers. |
| 712 | ctx_pos_out = ctx_start + j |
| 713 | tl.store(out_context_positions_ptr + ctx_pos_out, ctx_pos, mask=is_ctx) |
| 714 | query_out = req_idx * num_query_per_req + query_off |
| 715 | tl.store(out_query_positions_ptr + query_out, query_pos, mask=is_query) |
| 716 | |
| 717 | # --- Slot mapping (block_table lookup for all positions) --- |
| 718 | block_num = positions // block_size |
| 719 | # # Clamp block_number to avoid OOB when position is at max |
| 720 | block_num = tl.minimum(block_num, block_table_stride - 1) |
| 721 | block_id = tl.load( |
| 722 | block_table_ptr + req_idx * block_table_stride + block_num, |
| 723 | mask=in_bounds, |
| 724 | other=0, |
| 725 | ).to(tl.int64) |
| 726 | slot = block_id * block_size + (positions % block_size) |
| 727 | tl.store(out_context_slot_mapping_ptr + ctx_pos_out, slot, mask=is_ctx) |
| 728 | tl.store(out_query_slot_mapping_ptr + query_out, slot, mask=is_query) |
| 729 | |
| 730 | # --- Input IDs (query tokens only) --- |
| 731 | bonus_token = tl.load(next_token_ids_ptr + req_idx) |
| 732 | is_bonus = is_query & (query_off == 0) |
| 733 | input_id = tl.where(is_bonus, bonus_token, parallel_drafting_token_id) |
| 734 | tl.store(out_input_ids_ptr + query_out, input_id, mask=is_query) |
| 735 | |
| 736 | # --- Token indices to sample (mask tokens, skip the bonus token) --- |
| 737 | is_sample = is_query & (query_off > 0) |
| 738 | sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1) |
| 739 | tl.store( |
| 740 | out_token_indices_ptr + sample_out_idx, |
| 741 | query_out, |
| 742 | mask=is_sample, |
| 743 | ) |
| 744 | |
| 745 | |
| 746 | def _copy_and_expand_dflash_warmup_inputs( |
| 747 | *, |
| 748 | block_table_stride: int, |
| 749 | parallel_drafting_token_id: int, |
| 750 | block_size: int, |
| 751 | num_speculative_tokens: int, |
| 752 | ): |
| 753 | int32 = TritonWarmupTensor(torch.int32) |
| 754 | int64 = TritonWarmupTensor(torch.int64) |
| 755 | triton_block_size = WarmupIntRange(1, 257, advance=lambda value: value * 2) |
| 756 | has_num_rejected = WarmupChoices(False, True) |
| 757 | return dict( |
| 758 | next_token_ids_ptr=int32, |
| 759 | target_positions_ptr=int64, |
| 760 | out_input_ids_ptr=int32, |
| 761 | out_context_positions_ptr=int64, |
| 762 | out_query_positions_ptr=int64, |
| 763 | out_context_slot_mapping_ptr=int64, |
| 764 | out_query_slot_mapping_ptr=int64, |
| 765 | out_token_indices_ptr=int32, |
| 766 | block_table_ptr=int32, |
| 767 | query_start_loc_ptr=int32, |
| 768 | num_rejected_tokens_ptr=(int32 if has_num_rejected else 0), |
| 769 | block_table_stride=block_table_stride, |
| 770 | parallel_drafting_token_id=parallel_drafting_token_id, |
| 771 | block_size=block_size, |
| 772 | num_query_per_req=1 + num_speculative_tokens, |
| 773 | num_speculative_tokens=num_speculative_tokens, |
| 774 | total_input_tokens=WarmupChoices(1, 2, 16), |
| 775 | max_tokens_per_req=triton_block_size, |
| 776 | triton_block_size=triton_block_size, |
| 777 | has_num_rejected=has_num_rejected, |
| 778 | ) |
| 779 | |
| 780 | |
| 781 | @triton_kernel_dispatcher_with_warmup( |
| 782 | kernel=copy_and_expand_dflash_inputs_kernel, |
| 783 | warmup_inputs=_copy_and_expand_dflash_warmup_inputs, |
| 784 | ) |
| 785 | def _copy_and_expand_dflash_inputs( |
| 786 | *, |
| 787 | next_token_ids_ptr: torch.Tensor, |
| 788 | target_positions_ptr: torch.Tensor, |
| 789 | out_input_ids_ptr: torch.Tensor, |
| 790 | out_context_positions_ptr: torch.Tensor, |
| 791 | out_query_positions_ptr: torch.Tensor, |
| 792 | out_context_slot_mapping_ptr: torch.Tensor, |
| 793 | out_query_slot_mapping_ptr: torch.Tensor, |
| 794 | out_token_indices_ptr: torch.Tensor, |
| 795 | block_table_ptr: torch.Tensor, |
| 796 | query_start_loc_ptr: torch.Tensor, |
| 797 | num_rejected_tokens_ptr: torch.Tensor | int, |
| 798 | block_table_stride: int, |
| 799 | parallel_drafting_token_id: int, |
| 800 | block_size: int, |
| 801 | num_query_per_req: int, |
| 802 | num_speculative_tokens: int, |
| 803 | total_input_tokens: int, |
| 804 | max_tokens_per_req: int, |
| 805 | triton_block_size: int, |
| 806 | has_num_rejected: bool, |
| 807 | ) -> DispatchSpec: |
| 808 | num_reqs = query_start_loc_ptr.shape[0] - 1 |
| 809 | num_blocks = triton.cdiv(max_tokens_per_req, triton_block_size) |
| 810 | return (num_reqs, num_blocks), dict( |
| 811 | BLOCK_SIZE=triton_block_size, |
| 812 | ) |
| 813 | |
| 814 | |
| 815 | @torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) |
| 816 | def update_num_computed_tokens_for_batch_change( |
| 817 | num_computed_tokens: torch.Tensor, |
| 818 | num_accepted_tokens: torch.Tensor, |
| 819 | prev_positions: torch.Tensor, |
| 820 | valid_sampled_token_count: torch.Tensor, |
| 821 | prev_num_draft_tokens: torch.Tensor, |
| 822 | cpu_num_computed_tokens: torch.Tensor, |
| 823 | ) -> None: |
| 824 | """Correct num_computed_tokens for async spec decode drift. |
| 825 | |
| 826 | Requests that had drafts: corrected = prev_gpu + valid_count. |
| 827 | New requests or non-draft (e.g. prefills): use CPU value directly. |
| 828 | """ |
| 829 | # Clamp because prev_positions can be -1 for new requests |
| 830 | gather_indices = prev_positions.clamp(min=0) |
| 831 | |
| 832 | valid_counts = valid_sampled_token_count[gather_indices] |
| 833 | prev_computed = num_computed_tokens[gather_indices] |
| 834 | prev_drafts = prev_num_draft_tokens[gather_indices] |
| 835 | |
| 836 | participating = (prev_positions >= 0) & (prev_drafts > 0) |
| 837 | corrected = prev_computed + valid_counts.int() |
| 838 | |
| 839 | n = prev_positions.shape[0] |
| 840 | num_computed_tokens[:n].copy_( |
| 841 | torch.where(participating, corrected, cpu_num_computed_tokens) |
| 842 | ) |
| 843 | num_accepted_tokens.copy_( |
| 844 | torch.where(participating, valid_counts, num_accepted_tokens) |
| 845 | ) |
| 846 | |
| 847 | |
| 848 | def unconditional_to_conditional_rates(rates: list[float]) -> list[float]: |
| 849 | """Convert per-position unconditional rates to per-position conditional |
| 850 | rates for the early-terminating rejection loop (c_i = p_i / p_{i-1}).""" |
| 851 | return [p / q if q > 0.0 else 0.0 for p, q in zip(rates, [1.0, *rates[:-1]])] |