From ccf90ba784900324e8ff4e6495d13b47b27b941a Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 13 Apr 2026 16:01:24 -0700 Subject: [PATCH] [Model Runner V2] Add full cuda graph support for eagle prefill (#37588) Signed-off-by: Giancarlo Delfin --- vllm/v1/worker/gpu/model_runner.py | 6 - .../worker/gpu/spec_decode/eagle/cudagraph.py | 22 +- .../gpu/spec_decode/eagle/speculator.py | 322 +++++++++++------- 3 files changed, 213 insertions(+), 137 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ca52b8b66d..18f0f5015e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -464,7 +464,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states - num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp self.execute_model_state = None # dummy run the eagle speculator's propose to ensure DP/EP sync. @@ -496,7 +495,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): next_prefill_tokens=self.req_states.next_prefill_tokens, temperature=self.sampler.sampling_states.temperature.gpu, seeds=self.sampler.sampling_states.seeds.gpu, - num_tokens_across_dp=num_tokens_across_dp, dummy_run=True, skip_attn_for_dummy_run=skip_attn, mm_inputs=mm_inputs, @@ -1110,7 +1108,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): hidden_states=hidden_states, aux_hidden_states=aux_hidden_states, kv_connector_output=kv_connector_output, - num_tokens_across_dp=num_tokens_across_dp, ) if not self.is_last_pp_rank: @@ -1135,7 +1132,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): hidden_states = self.execute_model_state.hidden_states aux_hidden_states = self.execute_model_state.aux_hidden_states kv_connector_output = self.execute_model_state.kv_connector_output - num_tokens_across_dp = self.execute_model_state.num_tokens_across_dp self.execute_model_state = None if not self.is_last_pp_rank: @@ -1228,7 +1224,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): self.req_states.next_prefill_tokens, self.sampler.sampling_states.temperature.gpu, self.sampler.sampling_states.seeds.gpu, - num_tokens_across_dp=num_tokens_across_dp, mm_inputs=mm_inputs, ) self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens @@ -1336,4 +1331,3 @@ class ExecuteModelState(NamedTuple): hidden_states: torch.Tensor | None aux_hidden_states: list[torch.Tensor] | None kv_connector_output: KVConnectorOutput | None - num_tokens_across_dp: torch.Tensor | None diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py index 1e75c48966..655007109e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py @@ -19,21 +19,16 @@ from vllm.v1.worker.utils import AttentionGroup class EagleCudaGraphManager(CudaGraphManager): - """CudaGraphManager for Eagle speculative decoding (FULL mode only).""" + """CudaGraphManager for Eagle speculative decoding.""" def __init__( self, vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, - draft_tokens: torch.Tensor, + decode_query_len: int, ): - assert not cudagraph_mode.has_mode(CUDAGraphMode.PIECEWISE), ( - "EagleCudaGraphManager does not support PIECEWISE mode yet" - ) - # Eagle always uses uniform decode with query_len=1 - super().__init__(vllm_config, device, cudagraph_mode, decode_query_len=1) - self.draft_tokens = draft_tokens + super().__init__(vllm_config, device, cudagraph_mode, decode_query_len) # Use a dedicated pool for Eagle to avoid memory overlap with the main # model's cudagraph. The base class uses a shared global pool, but Eagle's @@ -44,7 +39,7 @@ class EagleCudaGraphManager(CudaGraphManager): def capture( self, - generate_fn: Callable, + forward_fn: Callable, model_state: ModelState, input_buffers: InputBuffers, block_tables: BlockTables, @@ -52,7 +47,7 @@ class EagleCudaGraphManager(CudaGraphManager): kv_cache_config: KVCacheConfig, progress_bar_desc: str = "Capturing CUDA graphs", ) -> None: - """Capture CUDA graphs for Eagle speculative decoding (FULL mode only).""" + """Capture CUDA graphs for Eagle.""" def create_forward_fn( desc: BatchExecutionDescriptor, @@ -74,7 +69,7 @@ class EagleCudaGraphManager(CudaGraphManager): kv_cache_config, ) - return lambda cg_mode: generate_fn( + return lambda cg_mode: forward_fn( num_reqs, num_tokens, attn_metadata, @@ -84,8 +79,3 @@ class EagleCudaGraphManager(CudaGraphManager): ) super().capture(create_forward_fn, progress_bar_desc) - - def run_fullgraph(self, desc: BatchExecutionDescriptor) -> torch.Tensor: - """Replay a captured FULL cudagraph and return draft tokens.""" - super().run_fullgraph(desc) - return self.draft_tokens diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index c311057dc2..881858f5d2 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -19,11 +19,16 @@ from vllm.v1.worker.gpu.attn_utils import ( init_attn_backend, ) from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.cudagraph_utils import ( + get_uniform_token_count, +) from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample -from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import EagleCudaGraphManager +from vllm.v1.worker.gpu.spec_decode.eagle.cudagraph import ( + EagleCudaGraphManager, +) from vllm.v1.worker.gpu.spec_decode.eagle.utils import load_eagle_model logger = init_logger(__name__) @@ -76,6 +81,9 @@ class EagleSpeculator: dtype=torch.int64, device=device, ) + self.last_token_indices = torch.zeros( + self.max_num_reqs, dtype=torch.int64, device=device + ) self.supports_mm_inputs = MULTIMODAL_REGISTRY.supports_multimodal_inputs( self.draft_model_config @@ -95,20 +103,30 @@ class EagleSpeculator: device=device, ) - self.cudagraph_manager: EagleCudaGraphManager | None = None + self.prefill_cudagraph_manager: EagleCudaGraphManager | None = None + self.decode_cudagraph_manager: EagleCudaGraphManager | None = None def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: - if cudagraph_mode.decode_mode() == CUDAGraphMode.FULL: - cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY - else: - cudagraph_mode = CUDAGraphMode.NONE - - self.cudagraph_manager = EagleCudaGraphManager( + cudagraph_mode = self.vllm_config.compilation_config.cudagraph_mode + # Initialize cudagraph manager for draft prefill (draft position 0). + self.prefill_cudagraph_manager = EagleCudaGraphManager( self.vllm_config, self.device, cudagraph_mode, - self.draft_tokens, + self.num_speculative_steps + 1, ) + # Initialize cudagraph manager for draft generation (draft positions > 0). + self.decode_cudagraph_manager = EagleCudaGraphManager( + self.vllm_config, + self.device, + # Only use FULL graph mode, if available, because draft decodes + # only consist of a single token. + cudagraph_mode.decode_mode(), + decode_query_len=1, + ) + # Share a single pool between prefill and decode since they never + # execute concurrently. + self.decode_cudagraph_manager.pool = self.prefill_cudagraph_manager.pool def load_model(self, target_model: nn.Module) -> None: target_attn_layer_names = get_layers_from_vllm_config( @@ -189,6 +207,47 @@ class EagleSpeculator: last_hidden_states, hidden_states = ret_hidden_states return last_hidden_states, hidden_states + def prefill( + self, + num_reqs: int, + num_tokens: int, + attn_metadata: dict[str, Any] | None, + slot_mappings: dict[str, torch.Tensor] | None, + num_tokens_across_dp: torch.Tensor | None, + cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + ) -> None: + last_token_indices = self.last_token_indices[:num_reqs] + pos = self.input_buffers.positions[last_token_indices] + idx_mapping = self.idx_mapping[:num_reqs] + + last_hidden_states, hidden_states = self.run_model( + num_tokens, + attn_metadata, + slot_mappings, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + mm_inputs=mm_inputs, + ) + sample_hidden_states = last_hidden_states[last_token_indices] + logits = self.model.compute_logits(sample_hidden_states) + + # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise + # used for draft and target sampling. + self.draft_tokens[:num_reqs, 0] = gumbel_sample( + logits, + idx_mapping, + self.temperature, + self.seeds, + pos + 1, + apply_temperature=True, + processed_logits_out=self.draft_logits[:, 0] + if self.draft_logits is not None + else None, + ) + self.hidden_states[:num_reqs] = hidden_states[last_token_indices] + self.input_buffers.positions[:num_reqs] = pos + def generate_draft( self, num_reqs: int, @@ -281,19 +340,46 @@ class EagleSpeculator: return attn_metadata def capture_model(self) -> None: - assert self.cudagraph_manager is not None + logger.info("Capturing model for Eagle speculator...") + # Reset indices to zeros to prevent stale values from prior + # dummy runs to cause out-of-bounds indexing during capture. + self.last_token_indices.zero_() + + # Capture the prefill routine (model forward + compute_logits + + # gumbel_sample). + # For FULL graphs, the entire routine is recorded as one graph. + # For PIECEWISE, only the model's compiled regions are captured + # and the rest (compute_logits, gumbel_sample) runs eagerly. + assert self.prefill_cudagraph_manager is not None + self.prefill_cudagraph_manager.capture( + self.prefill, + self.model_state, + self.input_buffers, + self.block_tables, + self.attn_groups, + self.kv_cache_config, + progress_bar_desc="Capturing eagle prefill CUDA graphs", + ) + if self.num_speculative_steps == 1: return - logger.info("Capturing model for Eagle speculator...") - self.cudagraph_manager.capture( + # Capture the decode draft generation loop (model forward + + # compute_logits + gumbel_sample + update_eagle_inputs, for + # each step). + # For FULL graphs, the entire multi-step loop is recorded as + # one graph. For PIECEWISE, only the model's compiled regions + # are captured, and the rest (compute_logits, gumbel_sample, + # update_eagle_inputs) runs eagerly. + assert self.decode_cudagraph_manager is not None + self.decode_cudagraph_manager.capture( self.generate_draft, self.model_state, self.input_buffers, self.block_tables, self.attn_groups, self.kv_cache_config, - progress_bar_desc="Capturing eagle CUDA graphs", + progress_bar_desc="Capturing eagle decode CUDA graphs", ) @torch.inference_mode() @@ -324,6 +410,10 @@ class EagleSpeculator: mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, ) -> torch.Tensor: + num_tokens = input_batch.num_tokens_after_padding + num_reqs = input_batch.num_reqs + max_query_len = input_batch.num_scheduled_tokens.max() + # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the # number of rejected tokens, we maintain the size of eagle's input_ids and # hidden_states the same as the target model's. This means, we pad each @@ -337,82 +427,88 @@ class EagleSpeculator: ) else: hidden_states = last_hidden_states - num_tokens = input_batch.num_tokens_after_padding - self.hidden_states[:num_tokens] = hidden_states + self.hidden_states[:num_tokens].copy_(hidden_states) - # Get the input ids and last token indices for the speculator. - last_token_indices = prepare_eagle_inputs( - self.input_buffers, - input_batch, - num_sampled, - num_rejected, - last_sampled, - next_prefill_tokens, - ) - - # Prefill: Run the eagle speculator with eager mode. - # TODO(woosuk): Support CUDA graph for prefill. - last_hidden_states, hidden_states = self.run_model( - num_tokens, - attn_metadata, - slot_mappings, - num_tokens_across_dp=num_tokens_across_dp, - mm_inputs=mm_inputs, - ) - sample_hidden_states = last_hidden_states[last_token_indices] - logits = self.model.compute_logits(sample_hidden_states) - - num_reqs = input_batch.num_reqs + # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. # NOTE(woosuk): For draft sampling, we only consider the temperature # and ignore the other sampling parameters such as top_k and top_p, # for simplicity and performance. # While this may slightly degrade the acceptance rate, it does not # affect the output distribution after rejection sampling. - idx_mapping = self.idx_mapping[:num_reqs] - idx_mapping.copy_(input_batch.idx_mapping) self.temperature.copy_(temperature) self.seeds.copy_(seeds) + self.idx_mapping[:num_reqs].copy_(input_batch.idx_mapping) - # Gather the values and copy them to the pre-allocated buffers. - pos = self.input_buffers.positions[:num_reqs] - torch.gather(input_batch.positions, 0, last_token_indices, out=pos) - # NOTE(woosuk): We must add 1 to the positions to match the Gumbel noise - # used for draft and target sampling. - draft_tokens = gumbel_sample( - logits, - idx_mapping, - self.temperature, - self.seeds, - pos + 1, - apply_temperature=True, - processed_logits_out=self.draft_logits[:, 0] - if self.draft_logits is not None - else None, + # Get the input ids and last token indices for the speculator. + prepare_eagle_inputs( + self.input_buffers, + input_batch, + self.last_token_indices, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + self.max_num_reqs, ) + # When all requests are decoding (no true prefills), each has + # num_speculative_steps + 1 tokens, enabling FULL graph replay. + # Mixed or prefill-only batches fall back to PIECEWISE. + prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( + self.prefill_cudagraph_manager, + num_reqs, + num_tokens, + get_uniform_token_count(num_reqs, num_tokens, max_query_len), + dp_size=self.dp_size, + dp_rank=self.dp_rank, + need_eager=is_profile, + ) + + if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: + # It is necessary to rebuild the attention metadata when + # replaying the FULL graph so that any attention metadata + # builder state is updated. + self._build_draft_attn_metadata( + num_reqs=num_reqs, + num_reqs_padded=prefill_batch_desc.num_reqs or num_reqs, + num_tokens_padded=prefill_batch_desc.num_tokens, + max_query_len=self.num_speculative_steps + 1, + ) + # Replay the full graph for draft prefill. + assert self.prefill_cudagraph_manager is not None + self.prefill_cudagraph_manager.run_fullgraph(prefill_batch_desc) + else: + # The target model's attention metadata and slot mappings + # can directly be used for draft prefill, because of the + # identical batch shape and KV cache layout. + self.prefill( + num_reqs, + prefill_batch_desc.num_tokens, + attn_metadata, + slot_mappings, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=prefill_batch_desc.cg_mode, + mm_inputs=mm_inputs, + ) + if self.num_speculative_steps == 1: # Early exit. - return draft_tokens.view(-1, 1) + return self.draft_tokens[:num_reqs, :1] - # Save the draft tokens for the first step. - self.draft_tokens[:num_reqs, 0] = draft_tokens # Prepare the inputs for the decode steps. prepare_eagle_decode( - draft_tokens, - hidden_states, - last_token_indices, + self.draft_tokens[:num_reqs, 0], input_batch.seq_lens, num_rejected, self.input_buffers, - self.hidden_states, self.max_model_len, self.max_num_reqs, ) - # Each request produces exactly 1 token per draft decode step, - # enabling FULL cudagraph. + # Each request produces exactly 1 token per draft generation step, + # enabling FULL graph replay. decode_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( - self.cudagraph_manager, + self.decode_cudagraph_manager, num_reqs, num_reqs, uniform_token_count=1, @@ -426,12 +522,12 @@ class EagleSpeculator: if not (dummy_run and skip_attn_for_dummy_run): # Build attention metadata and slot mappings for the draft # decode steps. It is necessary to rebuild the attention - # metadata even when replaying the FULL cudagraph so that - # any attention metadata builder state is updated. + # metadata even when replaying the FULL graph so that any + # attention metadata builder state is updated. slot_mappings = self.block_tables.compute_slot_mappings( - idx_mapping, + self.idx_mapping[:num_reqs], self.input_buffers.query_start_loc[: num_reqs + 1], - pos, + self.input_buffers.positions[:num_reqs], decode_batch_desc.num_tokens, ) slot_mappings_updated = build_slot_mappings_by_layer( @@ -445,8 +541,9 @@ class EagleSpeculator: ) if decode_batch_desc.cg_mode == CUDAGraphMode.FULL: - assert self.cudagraph_manager is not None - self.cudagraph_manager.run_fullgraph(decode_batch_desc) + # Replay the full graph for draft generation. + assert self.decode_cudagraph_manager is not None + self.decode_cudagraph_manager.run_fullgraph(decode_batch_desc) else: self.generate_draft( num_reqs, @@ -464,6 +561,8 @@ def _prepare_eagle_inputs_kernel( last_token_indices_ptr, eagle_input_ids_ptr, eagle_positions_ptr, + eagle_query_start_loc_ptr, + eagle_seq_lens_ptr, target_input_ids_ptr, target_positions_ptr, idx_mapping_ptr, @@ -472,20 +571,24 @@ def _prepare_eagle_inputs_kernel( num_sampled_ptr, num_rejected_ptr, query_start_loc_ptr, + seq_lens_ptr, + max_num_reqs, BLOCK_SIZE: tl.constexpr, ): - batch_idx = tl.program_id(0) - req_state_idx = tl.load(idx_mapping_ptr + batch_idx) + req_idx = tl.program_id(0) + num_reqs = tl.num_programs(0) + req_state_idx = tl.load(idx_mapping_ptr + req_idx) - query_start = tl.load(query_start_loc_ptr + batch_idx) - query_end = tl.load(query_start_loc_ptr + batch_idx + 1) + query_start = tl.load(query_start_loc_ptr + req_idx) + query_end = tl.load(query_start_loc_ptr + req_idx + 1) query_len = query_end - query_start + seq_len = tl.load(seq_lens_ptr + req_idx) # Get the true query length and next token after accounting for rejected tokens. - num_rejected = tl.load(num_rejected_ptr + batch_idx) + num_rejected = tl.load(num_rejected_ptr + req_idx) query_len -= num_rejected - num_sampled = tl.load(num_sampled_ptr + batch_idx) + num_sampled = tl.load(num_sampled_ptr + req_idx) if num_sampled > 0: next_token = tl.load(last_sampled_ptr + req_state_idx).to(tl.int32) else: @@ -501,7 +604,7 @@ def _prepare_eagle_inputs_kernel( tl.store(eagle_input_ids_ptr + query_start + block - 1, input_ids, mask=mask) last_token_index = query_start + query_len - 1 - tl.store(last_token_indices_ptr + batch_idx, last_token_index) + tl.store(last_token_indices_ptr + req_idx, last_token_index) tl.store(eagle_input_ids_ptr + last_token_index, next_token) # Copy positions. @@ -511,11 +614,29 @@ def _prepare_eagle_inputs_kernel( target_pos = tl.load(target_positions_ptr + query_start + block, mask=mask) tl.store(eagle_positions_ptr + query_start + block, target_pos, mask=mask) + # Copy query start locations. + tl.store(eagle_query_start_loc_ptr + req_idx, query_start) + # Copy sequence lengths. + tl.store(eagle_seq_lens_ptr + req_idx, seq_len) + if req_idx == (num_reqs - 1): + # Pad query_start_loc for CUDA graphs. + for i in range(num_reqs, max_num_reqs + 1, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + 1 + tl.store(eagle_query_start_loc_ptr + block, query_end, mask=mask) + # Pad seq_lens for CUDA graphs. + for i in range(num_reqs, max_num_reqs, BLOCK_SIZE): + block = i + tl.arange(0, BLOCK_SIZE) + mask = block < max_num_reqs + tl.store(eagle_seq_lens_ptr + block, 0, mask=mask) + def prepare_eagle_inputs( input_buffers: InputBuffers, input_batch: InputBatch, # [num_reqs] + last_token_indices: torch.Tensor, + # [num_reqs] num_sampled: torch.Tensor, # [num_reqs] num_rejected: torch.Tensor, @@ -523,17 +644,15 @@ def prepare_eagle_inputs( last_sampled: torch.Tensor, # [max_num_reqs] next_prefill_tokens: torch.Tensor, + max_num_reqs, ) -> torch.Tensor: num_reqs = input_batch.num_reqs - last_token_indices = torch.empty( - num_reqs, - dtype=torch.int64, - device=num_sampled.device, - ) _prepare_eagle_inputs_kernel[(num_reqs,)]( last_token_indices, input_buffers.input_ids, input_buffers.positions, + input_buffers.query_start_loc, + input_buffers.seq_lens, input_batch.input_ids, input_batch.positions, input_batch.idx_mapping, @@ -542,6 +661,8 @@ def prepare_eagle_inputs( num_sampled, num_rejected, input_batch.query_start_loc, + input_batch.seq_lens, + max_num_reqs, BLOCK_SIZE=1024, ) return last_token_indices @@ -550,18 +671,13 @@ def prepare_eagle_inputs( @triton.jit def _prepare_eagle_docode_kernel( draft_tokens_ptr, - output_hidden_states_ptr, - output_hidden_states_stride, - last_token_indices_ptr, + draft_tokens_stride, target_seq_lens_ptr, num_rejected_ptr, input_ids_ptr, positions_ptr, - input_hidden_states_ptr, - input_hidden_states_stride, query_start_loc_ptr, seq_lens_ptr, - hidden_size, max_model_len, max_num_reqs, BLOCK_SIZE: tl.constexpr, @@ -584,24 +700,9 @@ def _prepare_eagle_docode_kernel( return # draft token -> input id. - draft_token = tl.load(draft_tokens_ptr + req_idx) + draft_token = tl.load(draft_tokens_ptr + req_idx * draft_tokens_stride) tl.store(input_ids_ptr + req_idx, draft_token) - # output hidden states -> input hidden states. - src_idx = tl.load(last_token_indices_ptr + req_idx) - for i in range(0, hidden_size, BLOCK_SIZE): - block = i + tl.arange(0, BLOCK_SIZE) - mask = block < hidden_size - output_hidden_states = tl.load( - output_hidden_states_ptr + src_idx * output_hidden_states_stride + block, - mask=mask, - ) - tl.store( - input_hidden_states_ptr + req_idx * input_hidden_states_stride + block, - output_hidden_states, - mask=mask, - ) - # Compute position and seq_lens. # NOTE(woosuk): To prevent out-of-range access, we clamp these values # if they reach the max model length. @@ -618,31 +719,22 @@ def _prepare_eagle_docode_kernel( def prepare_eagle_decode( draft_tokens: torch.Tensor, - output_hidden_states: torch.Tensor, - last_token_indices: torch.Tensor, target_seq_lens: torch.Tensor, num_rejected: torch.Tensor, input_buffers: InputBuffers, - input_hidden_states: torch.Tensor, max_model_len: int, max_num_reqs: int, ): num_reqs = draft_tokens.shape[0] - hidden_size = output_hidden_states.shape[-1] _prepare_eagle_docode_kernel[(num_reqs + 1,)]( draft_tokens, - output_hidden_states, - output_hidden_states.stride(0), - last_token_indices, + draft_tokens.stride(0), target_seq_lens, num_rejected, input_buffers.input_ids, input_buffers.positions, - input_hidden_states, - input_hidden_states.stride(0), input_buffers.query_start_loc, input_buffers.seq_lens, - hidden_size, max_model_len, max_num_reqs, BLOCK_SIZE=1024,