diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py index a38af66e4..2b83b7df5 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py @@ -15,32 +15,23 @@ from __future__ import annotations import logging -import time -from typing import TYPE_CHECKING, List, Optional - -import torch +from typing import TYPE_CHECKING from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( MultiLayerEagleDraftExtendCudaGraphRunner, MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, ) -from sglang.srt.utils import get_available_gpu_memory logger = logging.getLogger(__name__) if TYPE_CHECKING: - from sglang.srt.speculative.multi_layer_eagle_worker import ( - MultiLayerEagleDraftWorker, - ) + pass class MultiLayerEagleDraftExtendNpuGraphRunner( MultiLayerEagleDraftExtendCudaGraphRunner ): - def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int): - super().__init__(eagle_worker, step) - def _replay_graph(self, shape_key, forward_batch): seq_lens = self.buffers.seq_lens_cpu[: self.raw_bs].tolist() + [0] * ( self.bs - self.raw_bs @@ -56,86 +47,8 @@ class MultiLayerEagleDraftExtendNpuGraphRunner( class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner( MultiLayerEagleMultiStepDraftExtendCudaGraphRunner ): - def __init__(self, eagle_worker: MultiLayerEagleDraftWorker): - super().__init__(eagle_worker) + def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendNpuGraphRunner: + return MultiLayerEagleDraftExtendNpuGraphRunner(self.eagle_worker, step) - def _init_and_capture(self): - if cuda_graph_fully_disabled(): - self.runners = [None] * self.speculative_num_steps - return - - self.runners: List[Optional[MultiLayerEagleDraftExtendNpuGraphRunner]] = [] - buffer_len_list: List[int] = [] - - for step in range(self.speculative_num_steps): - if self.draft_extend_attn_backend_list[step]: - runner = MultiLayerEagleDraftExtendNpuGraphRunner( - self.eagle_worker, step - ) - self.runners.append(runner) - - self.seq_len_fill_value = runner.seq_len_fill_value - self.max_bs = runner.max_bs - buffer_len_list.append(runner.max_num_token) - self.offsets.append(self.offsets[-1] + runner.max_num_token) - else: - self.runners.append(None) - - self.cuda_graph_buffers["seq_lens_cpu"] = torch.full( - (self.max_bs,), - self.seq_len_fill_value, - dtype=torch.int32, - ) - - with torch.device(self.device): - self.cuda_graph_buffers["input_ids"] = torch.zeros( - (self.offsets[-1],), dtype=torch.int64 - ) - self.cuda_graph_buffers["out_cache_loc"] = torch.ones( - (self.offsets[-1],), dtype=torch.int64 - ) - self.cuda_graph_buffers["positions"] = torch.zeros( - (self.offsets[-1],), dtype=torch.int64 - ) - - self.cuda_graph_buffers["seq_lens"] = torch.full( - (self.max_bs,), - self.seq_len_fill_value, - dtype=torch.int32, - ) - self.cuda_graph_buffers["req_pool_indices"] = torch.zeros( - (self.max_bs,), dtype=torch.int64 - ) - self.cuda_graph_buffers["num_correct_drafts"] = torch.full( - (self.max_bs,), 1, dtype=torch.int32 - ) - self.cuda_graph_buffers["num_accept_tokens"] = torch.full( - (self.max_bs,), 1, dtype=torch.int32 - ) - - for step in range(self.speculative_num_steps - 1, -1, -1): - if self.runners[step] is not None: - tic = time.perf_counter() - before_mem = get_available_gpu_memory(self.device, self.gpu_id) - logger.info( - f"Capture draft extend CUDA graph begin. step={step}, " - f"avail mem={before_mem:.2f} GB" - ) - - self.runners[step].init_buffers_and_capture( - self.cuda_graph_buffers, - self.offsets[step], - ( - self.runners[step + 1] - if step + 1 < self.speculative_num_steps - else None - ), - ) - - after_mem = get_available_gpu_memory(self.device, self.gpu_id) - logger.info( - "Capture draft extend CUDA graph end. " - f"step={step}, elapsed={time.perf_counter() - tic:.2f} s, " - f"mem usage={(before_mem - after_mem):.2f} GB, " - f"avail mem={after_mem:.2f} GB." - ) + def _cuda_graph_disabled(self) -> bool: + return cuda_graph_fully_disabled() diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 07f873b00..31ce5c826 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -14,9 +14,11 @@ from __future__ import annotations +import contextlib import logging import time from dataclasses import dataclass +from types import SimpleNamespace from typing import TYPE_CHECKING, Callable, List, Optional import torch @@ -41,7 +43,6 @@ from sglang.srt.model_executor.forward_batch_info import ( from sglang.srt.model_executor.forward_context import ( ForwardContext, forward_context, - get_req_to_token_pool, ) from sglang.srt.model_executor.input_buffers import ForwardInputBuffers from sglang.srt.model_executor.runner import ( @@ -56,7 +57,6 @@ from sglang.srt.model_executor.runner_backend_utils import ( CUDA_GRAPH_CAPTURE_FAILED_MSG, ) from sglang.srt.speculative.eagle_info import EagleDraftExtendInput -from sglang.srt.speculative.multi_layer_eagle_utils import assign_new_state_triton from sglang.srt.speculative.spec_utils import fast_topk from sglang.srt.utils import ( get_available_gpu_memory, @@ -77,19 +77,22 @@ logger = logging.getLogger(__name__) @dataclass class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers): - # Sliced from shared parent buffers + """A single persistent buffer set shared by every MTP draft step.""" + input_ids: torch.Tensor out_cache_loc: torch.Tensor positions: torch.Tensor - # Shared from parent seq_lens: torch.Tensor seq_lens_cpu: torch.Tensor req_pool_indices: torch.Tensor num_correct_drafts: torch.Tensor num_accept_tokens: torch.Tensor - # Per-step buffers extend_seq_lens: torch.Tensor extend_start_loc: torch.Tensor + # Flat index (into the token dimension) of each request's last accepted + # token. Used both by the in-graph top-k gather and by the worker's + # per-step input_ids rotation. + select_index: torch.Tensor mrope_positions: torch.Tensor hidden_states: torch.Tensor next_token_logits_buffer: torch.Tensor @@ -100,12 +103,11 @@ class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers): class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): """Per-step multi-layer EAGLE draft-extend runner. - Subclasses DecodeCudaGraphRunner. Shares buffers across steps - via the composite MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, - so initialization is split: __init__ does basic field setup, - init_buffers_and_capture (called by the composite once shared - buffers exist) finishes by allocating per-step buffers and running - capture. + Subclasses DecodeCudaGraphRunner. All steps share a single buffer set + owned by the composite MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, + so initialization is split: __init__ does basic field setup, and + init_buffers_and_capture (called by the composite once the shared buffers + exist) attaches them and runs capture. """ def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int): @@ -151,10 +153,12 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): self.capture_bs, _ = get_batch_sizes_to_capture(model_runner) self.padded_static_len = -1 - # For Attention Backend - self.num_tokens_per_bs = self.speculative_num_steps + 1 + step + # Fixed window: every step extends each request by the same number of + # tokens, which lets all steps share one buffer set. + self.num_tokens_per_bs = self.speculative_num_draft_tokens self.max_bs = max(self.capture_bs) self.max_num_token = self.max_bs * self.num_tokens_per_bs + self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs self.eagle_worker.draft_extend_attn_backend_list[ self.step @@ -163,116 +167,13 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): self.step ].get_cuda_graph_seq_len_fill_value() - def init_buffers_and_capture( - self, - cuda_graph_buffers, - offset, - next_cuda_graph_runner, - ): - self.next_cuda_graph_runner = next_cuda_graph_runner - seq_lens_cpu = cuda_graph_buffers["seq_lens_cpu"] - self.extend_seq_lens_cpu = [self.num_tokens_per_bs] * self.max_bs + def init_buffers_and_capture(self, buffers: MultiLayerEagleDraftExtendInputBuffers): + """Attach the shared buffer set and capture this step's graphs.""" + self.buffers = buffers if self.enable_torch_compile: set_torch_compile_config() - with torch.device(self.model_runner.device): - input_ids = cuda_graph_buffers["input_ids"][ - offset : offset + self.max_num_token - ] - out_cache_loc = cuda_graph_buffers["out_cache_loc"][ - offset : offset + self.max_num_token - ] - positions = cuda_graph_buffers["positions"][ - offset : offset + self.max_num_token - ] - - seq_lens = cuda_graph_buffers["seq_lens"] - req_pool_indices = cuda_graph_buffers["req_pool_indices"] - num_correct_drafts = cuda_graph_buffers["num_correct_drafts"] - num_accept_tokens = cuda_graph_buffers["num_accept_tokens"] - - extend_seq_lens = torch.full( - (self.max_bs,), - self.num_tokens_per_bs, - dtype=torch.int32, - ) - extend_start_loc = torch.arange( - 0, - self.max_bs * self.num_tokens_per_bs, - step=self.num_tokens_per_bs, - dtype=torch.int32, - ) - - mrope_positions = torch.zeros((3, self.max_num_token), dtype=torch.int64) - - hidden_states = torch.zeros( - ( - self.max_num_token, - EagleDraftExtendInput.hidden_size_for(self.eagle_worker), - ), - dtype=EagleDraftExtendInput.dtype_for(self.eagle_worker), - ) - - if self.require_gathered_buffer: - if self.require_mlp_tp_gather: - global_num_tokens_gpu = torch.zeros( - (self.dp_size,), dtype=torch.int32 - ) - global_num_tokens_for_logprob_gpu = torch.zeros( - (self.dp_size,), dtype=torch.int32 - ) - else: - assert self.require_attn_tp_gather - global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32) - global_num_tokens_for_logprob_gpu = torch.zeros( - (1,), dtype=torch.int32 - ) - else: - global_num_tokens_gpu = None - global_num_tokens_for_logprob_gpu = None - - if hasattr( - self.model_runner.model_config.hf_config, "draft_vocab_size" - ): # llama_eagle - vocab_size = self.model_runner.model_config.hf_config.draft_vocab_size - elif hasattr( - self.model_runner.model_config.hf_config, "hot_vocab_size" - ): # llama_eagle3 - vocab_size = self.model_runner.model_config.hf_config.hot_vocab_size - else: - vocab_size = self.model_runner.model_config.vocab_size - - next_token_logits_buffer = torch.zeros( - ( - ( - self.max_bs * self.num_tokens_per_bs - if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2 - else self.max_bs - ), - vocab_size, - ), - dtype=torch.float, - ) - - self.buffers = MultiLayerEagleDraftExtendInputBuffers( - input_ids=input_ids, - out_cache_loc=out_cache_loc, - positions=positions, - seq_lens=seq_lens, - seq_lens_cpu=seq_lens_cpu, - req_pool_indices=req_pool_indices, - num_correct_drafts=num_correct_drafts, - num_accept_tokens=num_accept_tokens, - extend_seq_lens=extend_seq_lens, - extend_start_loc=extend_start_loc, - mrope_positions=mrope_positions, - hidden_states=hidden_states, - next_token_logits_buffer=next_token_logits_buffer, - global_num_tokens_gpu=global_num_tokens_gpu, - global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, - ) - self.backend = resolve_decode_backend(self) try: @@ -403,6 +304,19 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ) return forward_batch + def _postprocess_forward_batch(self, forward_batch: ForwardBatch, bs: int): + """Hook for subclasses to mutate the captured forward batch.""" + return forward_batch + + def _compute_topk(self, ret, bs: int): + """Compute top-k on the last accepted token's logits and attach it to + ``ret``. The gather index lives in a persistent buffer, so the captured + graph reads the right rows on each replay. Overridable so distributed + (vocab-sharded) builds can plug in an all-reduce-aware sampler.""" + buffers = self.buffers + probs = torch.softmax(ret.next_token_logits[buffers.select_index[:bs]], dim=-1) + ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1) + def capture_one_shape( self, size: int, @@ -415,6 +329,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): num_tokens = bs * self.num_tokens_per_bs forward_batch = self.get_forward_batch(bs) + forward_batch = self._postprocess_forward_batch(forward_batch, bs) attn_backend = self.eagle_worker.draft_extend_attn_backend_list[self.step] def run_once(): @@ -443,47 +358,8 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ): buffers.hidden_states[:num_tokens].copy_(ret.hidden_states[:num_tokens]) - # num_correct_drafts is drafts-only; the last accepted draft sits at index - # `num_correct_drafts` within the (current_token + drafts) slot range. - select_index = ( - torch.arange(bs, device=self.model_runner.device) - * (self.speculative_num_draft_tokens + self.step) - + buffers.num_correct_drafts[:bs] - + self.step - ) + self._compute_topk(ret, bs) - probs = torch.softmax(ret.next_token_logits[select_index], dim=-1) - ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1) - - if self.next_cuda_graph_runner is not None: - next_buffers = self.next_cuda_graph_runner.buffers - # rejected drafts = proposed drafts - accepted drafts. - # speculative_num_draft_tokens includes the current-token slot, so -1. - padding_lens = ( - self.speculative_num_draft_tokens - 1 - ) - buffers.num_correct_drafts[:bs] - assign_new_state_triton( - ret.topk_index, - buffers.input_ids, - buffers.positions, - buffers.hidden_states, - buffers.out_cache_loc, - buffers.extend_seq_lens, - buffers.extend_start_loc, - next_buffers.input_ids, - next_buffers.positions, - next_buffers.hidden_states, - next_buffers.out_cache_loc, - next_buffers.extend_seq_lens, - next_buffers.extend_start_loc, - next_buffers.seq_lens, - padding_lens, - forward_batch.batch_size, - self.step, - forward_batch.req_pool_indices, - get_req_to_token_pool().req_to_token, - self.eagle_worker.req_to_hidden_states_pool, - ) forward_batch.out_cache_loc = output_cache_loc_backup forward_batch.spec_info.hidden_states = hidden_states_backup return ret @@ -501,123 +377,53 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ), ) - def init_replay_state( - self, forward_batch: ForwardBatch, bs: int, raw_bs: int, num_tokens: int - ): - buffers = self.buffers - buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids) - buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) - if forward_batch.extend_seq_lens is not None: - buffers.extend_seq_lens[:raw_bs].copy_(forward_batch.extend_seq_lens) - buffers.extend_start_loc[:raw_bs].copy_(forward_batch.extend_start_loc) - buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc) - buffers.positions[:num_tokens].copy_(forward_batch.positions) - if ( - forward_batch.spec_info.hidden_states.shape[1] - == buffers.hidden_states.shape[1] - ): - buffers.hidden_states[:num_tokens].copy_( - forward_batch.spec_info.hidden_states - ) - if forward_batch.spec_info.num_correct_drafts is not None: - buffers.num_correct_drafts[:raw_bs].copy_( - forward_batch.spec_info.num_correct_drafts - ) - buffers.num_accept_tokens[:raw_bs].copy_( - forward_batch.spec_info.num_accept_tokens - ) - buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) - - if forward_batch.seq_lens_cpu is not None: - if bs != raw_bs: - buffers.seq_lens_cpu.fill_(self.seq_len_fill_value) - buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu) - - if forward_batch.extend_seq_lens_cpu is not None: - self.extend_seq_lens_cpu[:raw_bs] = forward_batch.extend_seq_lens_cpu - - def execute(self, forward_batch: ForwardBatch, init_state: bool = True): - assert forward_batch.out_cache_loc is not None + def replay(self, bs: int, seq_lens_sum: int, spec_info: EagleDraftExtendInput): + """Init this step's attention metadata for the prepared bucket and + replay its graph. Buffers must already be populated by the composite + runner's ``prepare`` (step 0) or by the previous step's in-graph chain + write + worker-side rotation (steps > 0).""" self.deepep_adapter.replay() buffers = self.buffers - - raw_bs = forward_batch.batch_size - num_tokens = raw_bs * self.num_tokens_per_bs - if self.require_mlp_tp_gather: - max_batch_size = max(forward_batch.original_global_num_tokens_cpu) - bs = self._pad_to_bucket(int(max_batch_size), self.capture_bs) - else: - bs = self._pad_to_bucket(raw_bs, self.capture_bs) - - if init_state: - self.init_replay_state(forward_batch, bs, raw_bs, num_tokens) + num_tokens = bs * self.num_tokens_per_bs if self.require_gathered_buffer: - buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs) - buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs) - - forward_batch.spec_info.hidden_states = buffers.hidden_states[:num_tokens] - forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[:bs] - forward_batch.spec_info.num_accept_tokens = buffers.num_accept_tokens[:bs] - forward_batch.spec_info.num_tokens_per_req = self.num_tokens_per_bs - forward_batch.spec_info.num_tokens_for_logprob_per_req = 1 - forward_batch.spec_info.positions = buffers.positions[:num_tokens] - forward_batch.spec_info.extend_seq_lens_tensor = buffers.extend_seq_lens[:bs] - - from types import SimpleNamespace + buffers.global_num_tokens_gpu.fill_(num_tokens) + buffers.global_num_tokens_for_logprob_gpu.fill_(num_tokens) fb_view = SimpleNamespace( batch_size=bs, forward_mode=self.forward_mode, - input_ids=getattr(forward_batch, "input_ids", None), + input_ids=buffers.input_ids[:num_tokens], req_pool_indices=buffers.req_pool_indices, seq_lens=buffers.seq_lens, - seq_lens_sum=forward_batch.seq_lens_sum - + (bs - raw_bs) * self.seq_len_fill_value, + seq_lens_sum=seq_lens_sum, seq_lens_cpu=buffers.seq_lens_cpu, encoder_lens=None, - # per-step write target (advanced in-graph by assign_new_state); - # forward_batch.out_cache_loc is frozen at step 0. + # per-step write target; out_cache_loc is frozen at prepare() time. out_cache_loc=buffers.out_cache_loc[:num_tokens], - spec_info=forward_batch.spec_info, + spec_info=spec_info, ) self.eagle_worker.draft_extend_attn_backend_list[ self.step ].init_forward_metadata_out_graph(fb_view) - self.raw_bs = raw_bs self.bs = bs shape_key = self._make_graph_key(bs) - out = self._replay_graph(shape_key, forward_batch) - - if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2: - unpadding_bs = num_tokens - elif bs != raw_bs: - forward_batch.spec_info.num_correct_drafts = buffers.num_correct_drafts[ - :raw_bs - ] - forward_batch.spec_info.num_accept_tokens = buffers.num_accept_tokens[ - :raw_bs - ] - unpadding_bs = raw_bs - else: - unpadding_bs = None - - if unpadding_bs is not None: - out_copy = out - out = LogitsProcessorOutput( - next_token_logits=out.next_token_logits[:unpadding_bs], - hidden_states=out.hidden_states[:unpadding_bs], - ) - out.topk_p = out_copy.topk_p[:raw_bs] - out.topk_index = out_copy.topk_index[:raw_bs] - return out + return self._replay_graph(shape_key, fb_view) class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: - """Composite orchestrator that owns speculative_num_steps per-step - runners with shared input buffers. Not itself a - DecodeCudaGraphRunner — it only routes work to the per-step + """Owns one shared buffer set and the per-step runners. + + Usage from the worker:: + + runner.prepare(forward_batch) + for step in range(num_steps): + _, topk_p, topk_index = runner.replay(step) + if step < num_steps - 1: + rotate_input_ids_triton(...) # advance the draft chain + + Not itself a DecodeCudaGraphRunner -- it only routes work to the per-step runners. """ @@ -630,72 +436,57 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: eagle_worker.draft_extend_attn_backend_list ) - self.runners = [] - self.cuda_graph_buffers = {} + self.runners: List[Optional[MultiLayerEagleDraftExtendCudaGraphRunner]] = [] self.seq_len_fill_value = 1 self.max_bs = 1 - self.offsets = [0] + self.num_tokens_per_bs = 1 self._init_and_capture() + def _create_runner(self, step: int) -> MultiLayerEagleDraftExtendCudaGraphRunner: + return MultiLayerEagleDraftExtendCudaGraphRunner(self.eagle_worker, step) + + def _capture_context(self, step: int): + """Context manager active while capturing ``step``'s graphs. Subclasses + can use it e.g. to temporarily expose a sharded local vocab size.""" + return contextlib.nullcontext() + + def _on_runners_created(self): + """Hook called after all per-step runners exist but before buffers are + allocated/captured (e.g. to allocate shared sconv buffers).""" + + def _cuda_graph_disabled(self) -> bool: + return check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED) + def _init_and_capture(self): - if check_cuda_graph_backend(Phase.DECODE, Backend.DISABLED): + if self._cuda_graph_disabled(): self.runners = [None] * self.speculative_num_steps return - self.runners: List[Optional[MultiLayerEagleDraftExtendCudaGraphRunner]] = [] - buffer_len_list: List[int] = [] + self.runners = [] - # 1. Construct per-step runners (cheap setup only). + # 1. Construct per-step runners (each initializes its own attn cuda + # graph state). They share the same fixed window size. for step in range(self.speculative_num_steps): if self.draft_extend_attn_backend_list[step]: - runner = MultiLayerEagleDraftExtendCudaGraphRunner( - self.eagle_worker, step - ) + runner = self._create_runner(step) self.runners.append(runner) - self.seq_len_fill_value = runner.seq_len_fill_value self.max_bs = runner.max_bs - buffer_len_list.append(runner.max_num_token) - self.offsets.append(self.offsets[-1] + runner.max_num_token) + self.num_tokens_per_bs = runner.num_tokens_per_bs + self.capture_bs = runner.capture_bs + self.require_gathered_buffer = runner.require_gathered_buffer + self.require_mlp_tp_gather = runner.require_mlp_tp_gather + self.require_mlp_sync = runner.require_mlp_sync + self.disable_padding = runner.disable_padding else: self.runners.append(None) - # 2. Allocate shared buffers. - self.cuda_graph_buffers["seq_lens_cpu"] = torch.full( - (self.max_bs,), - self.seq_len_fill_value, - dtype=torch.int32, - ) + self._on_runners_created() - with torch.device(self.device): - self.cuda_graph_buffers["input_ids"] = torch.zeros( - (self.offsets[-1],), dtype=torch.int64 - ) - self.cuda_graph_buffers["out_cache_loc"] = torch.ones( - (self.offsets[-1],), dtype=torch.int64 - ) - self.cuda_graph_buffers["positions"] = torch.zeros( - (self.offsets[-1],), dtype=torch.int64 - ) - - self.cuda_graph_buffers["seq_lens"] = torch.full( - (self.max_bs,), - self.seq_len_fill_value, - dtype=torch.int32, - ) - self.cuda_graph_buffers["req_pool_indices"] = torch.zeros( - (self.max_bs,), dtype=torch.int64 - ) - self.cuda_graph_buffers["num_correct_drafts"] = torch.full( - (self.max_bs,), 1, dtype=torch.int32 - ) - self.cuda_graph_buffers["num_accept_tokens"] = torch.full( - (self.max_bs,), 1, dtype=torch.int32 - ) - - # 3. Per-step capture, in reverse order so that next_cuda_graph_runner - # is already initialized when this step references it. + # 2. Allocate the single shared buffer set and capture each step in + # reverse order. + self.buffers = self._allocate_buffers() for step in range(self.speculative_num_steps - 1, -1, -1): if self.runners[step] is not None: tic = time.perf_counter() @@ -705,15 +496,8 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: f"avail mem={before_mem:.2f} GB" ) - self.runners[step].init_buffers_and_capture( - self.cuda_graph_buffers, - self.offsets[step], - ( - self.runners[step + 1] - if step + 1 < self.speculative_num_steps - else None - ), - ) + with self._capture_context(step): + self.runners[step].init_buffers_and_capture(self.buffers) after_mem = get_available_gpu_memory(self.device, self.gpu_id) logger.info( @@ -723,18 +507,190 @@ class MultiLayerEagleMultiStepDraftExtendCudaGraphRunner: f"avail mem={after_mem:.2f} GB." ) - def reset_buffers(self, forward_batch, batch_result): - self.cuda_graph_buffers["input_ids"].zero_() - self.cuda_graph_buffers["seq_lens"].fill_(self.seq_len_fill_value) - self.cuda_graph_buffers["out_cache_loc"].zero_() - self.cuda_graph_buffers["positions"].zero_() - # `batch_result.accept_lens` is drafts + bonus. - bs = forward_batch.batch_size - self.cuda_graph_buffers["num_correct_drafts"][:bs].copy_( - batch_result.accept_lens - 1 + def _vocab_size(self) -> int: + model_runner = self.eagle_worker.mtp_model_runner(0) + if hasattr(model_runner.model_config.hf_config, "draft_vocab_size"): + return model_runner.model_config.hf_config.draft_vocab_size + if hasattr(model_runner.model_config.hf_config, "hot_vocab_size"): + return model_runner.model_config.hf_config.hot_vocab_size + return model_runner.model_config.vocab_size + + def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers: + runner = next(r for r in self.runners if r is not None) + max_bs = self.max_bs + num_tokens_per_bs = self.num_tokens_per_bs + max_num_token = max_bs * num_tokens_per_bs + hidden_size = EagleDraftExtendInput.hidden_size_for(self.eagle_worker) + dtype = EagleDraftExtendInput.dtype_for(self.eagle_worker) + vocab_size = self._vocab_size() + + seq_lens_cpu = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32) + + with torch.device(self.device): + input_ids = torch.zeros((max_num_token,), dtype=torch.int64) + out_cache_loc = torch.ones((max_num_token,), dtype=torch.int64) + positions = torch.zeros((max_num_token,), dtype=torch.int64) + mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64) + hidden_states = torch.zeros((max_num_token, hidden_size), dtype=dtype) + + seq_lens = torch.full((max_bs,), self.seq_len_fill_value, dtype=torch.int32) + req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64) + num_correct_drafts = torch.full((max_bs,), 1, dtype=torch.int32) + num_accept_tokens = torch.full((max_bs,), 1, dtype=torch.int32) + + # Fixed window: every request extends by exactly num_tokens_per_bs + # tokens, and start locs are a constant arange. + extend_seq_lens = torch.full( + (max_bs,), num_tokens_per_bs, dtype=torch.int32 + ) + extend_start_loc = torch.arange( + 0, max_num_token, step=num_tokens_per_bs, dtype=torch.int32 + ) + select_index = torch.zeros((max_bs,), dtype=torch.int64) + + next_token_logits_buffer = torch.zeros( + (max_num_token, vocab_size), dtype=torch.float + ) + + if self.require_gathered_buffer: + if self.require_mlp_tp_gather: + dp_size = runner.dp_size + global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32) + global_num_tokens_for_logprob_gpu = torch.zeros( + (dp_size,), dtype=torch.int32 + ) + else: + global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32) + global_num_tokens_for_logprob_gpu = torch.zeros( + (1,), dtype=torch.int32 + ) + else: + global_num_tokens_gpu = None + global_num_tokens_for_logprob_gpu = None + + return MultiLayerEagleDraftExtendInputBuffers( + input_ids=input_ids, + out_cache_loc=out_cache_loc, + positions=positions, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + req_pool_indices=req_pool_indices, + num_correct_drafts=num_correct_drafts, + num_accept_tokens=num_accept_tokens, + extend_seq_lens=extend_seq_lens, + extend_start_loc=extend_start_loc, + select_index=select_index, + mrope_positions=mrope_positions, + hidden_states=hidden_states, + next_token_logits_buffer=next_token_logits_buffer, + global_num_tokens_gpu=global_num_tokens_gpu, + global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, ) - self.cuda_graph_buffers["num_accept_tokens"][:bs].copy_( - batch_result.accept_lens + + def _prepare_extra(self, forward_batch: ForwardBatch) -> None: + """Hook for subclasses to populate extra per-call buffers (e.g. sconv).""" + + def prepare(self, forward_batch: ForwardBatch): + """Populate the shared buffers once from ``forward_batch`` and bucketize + the batch size. Subsequent ``replay(step)`` calls reuse this state.""" + buffers = self.buffers + raw_bs = forward_batch.batch_size + num_tokens = raw_bs * self.num_tokens_per_bs + + # Bucketize to a captured batch size (padding the tail). + if self.require_mlp_tp_gather: + max_batch_size = max(forward_batch.original_global_num_tokens_cpu) + bs = self.get_runner(0)._pad_to_bucket(int(max_batch_size), self.capture_bs) + else: + bs = self.get_runner(0)._pad_to_bucket(raw_bs, self.capture_bs) + + # Reset padded slots, then copy the real values in. + buffers.input_ids.zero_() + buffers.out_cache_loc.zero_() + buffers.positions.zero_() + buffers.seq_lens.fill_(self.seq_len_fill_value) + + buffers.input_ids[:num_tokens].copy_(forward_batch.input_ids) + buffers.positions[:num_tokens].copy_(forward_batch.positions) + buffers.out_cache_loc[:num_tokens].copy_(forward_batch.out_cache_loc) + buffers.seq_lens[:raw_bs].copy_(forward_batch.seq_lens) + buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) + + if ( + forward_batch.spec_info.hidden_states.shape[1] + == buffers.hidden_states.shape[1] + ): + buffers.hidden_states[:num_tokens].copy_( + forward_batch.spec_info.hidden_states + ) + + buffers.num_correct_drafts[:raw_bs].copy_( + forward_batch.spec_info.num_correct_drafts + ) + buffers.num_accept_tokens[:raw_bs].copy_( + forward_batch.spec_info.num_accept_tokens + ) + + if forward_batch.seq_lens_cpu is not None: + if bs != raw_bs: + buffers.seq_lens_cpu.fill_(self.seq_len_fill_value) + buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu) + + # select_index[i] = i * window + num_correct_drafts[i]: the flat index + # of request i's last accepted token. Used by the in-graph top-k gather + # and by the worker's rotation. + arange = torch.arange(bs, device=self.device, dtype=torch.int64) + buffers.select_index[:bs].copy_( + arange * self.num_tokens_per_bs + buffers.num_correct_drafts[:bs] + ) + + if self.require_gathered_buffer: + buffers.global_num_tokens_gpu.fill_(bs * self.num_tokens_per_bs) + buffers.global_num_tokens_for_logprob_gpu.fill_(bs * self.num_tokens_per_bs) + + # Reusable spec_info for per-step attention metadata. + padded_num_tokens = bs * self.num_tokens_per_bs + spec_info = EagleDraftExtendInput( + hidden_states=buffers.hidden_states[:padded_num_tokens], + num_correct_drafts=buffers.num_correct_drafts[:bs], + num_accept_tokens=buffers.num_accept_tokens[:bs], + ) + spec_info.num_tokens_per_req = self.num_tokens_per_bs + spec_info.num_tokens_for_logprob_per_req = 1 + spec_info.positions = buffers.positions[:padded_num_tokens] + spec_info.extend_seq_lens_tensor = buffers.extend_seq_lens[:bs] + self._replay_spec_info = spec_info + + self.raw_bs = raw_bs + self.bs = bs + self.raw_num_tokens = num_tokens + self.seq_lens_sum = ( + forward_batch.seq_lens_sum + (bs - raw_bs) * self.seq_len_fill_value + ) + + self._prepare_extra(forward_batch) + + def replay(self, step: int): + """Replay ``step``'s graph at the prepared bucket. Returns + ``(LogitsProcessorOutput, topk_p, topk_index)`` sliced to the real + batch size.""" + runner = self.runners[step] + runner.raw_bs = self.raw_bs + out = runner.replay(self.bs, self.seq_lens_sum, self._replay_spec_info) + raw_bs = self.raw_bs + raw_num_tokens = self.raw_num_tokens + logits_output = LogitsProcessorOutput( + next_token_logits=out.next_token_logits[:raw_num_tokens], + hidden_states=( + out.hidden_states[:raw_num_tokens] + if out.hidden_states is not None + else None + ), + ) + return ( + logits_output, + out.topk_p[:raw_bs], + out.topk_index[:raw_bs], ) def get_runner(self, step): diff --git a/python/sglang/srt/speculative/multi_layer_eagle_utils.py b/python/sglang/srt/speculative/multi_layer_eagle_utils.py index f1db97e6c..d9ae4cb77 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_utils.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_utils.py @@ -13,21 +13,11 @@ # ============================================================================== from sglang.srt.speculative.triton_ops.multi_layer_eagle import ( - assign_hidden_states_pool_kernel, - assign_hidden_states_pool_torch, - assign_hidden_states_pool_triton, - assign_new_state_kernel, - assign_new_state_triton, rotate_input_ids_kernel, rotate_input_ids_triton, ) __all__ = [ - "assign_hidden_states_pool_kernel", - "assign_hidden_states_pool_torch", - "assign_hidden_states_pool_triton", - "assign_new_state_kernel", - "assign_new_state_triton", "rotate_input_ids_kernel", "rotate_input_ids_triton", ] diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 0cb0d7847..3fb0228ae 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -58,10 +58,7 @@ from sglang.srt.speculative.eagle_utils import ( from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, ) -from sglang.srt.speculative.multi_layer_eagle_utils import ( - assign_hidden_states_pool_triton, - rotate_input_ids_triton, -) +from sglang.srt.speculative.multi_layer_eagle_utils import rotate_input_ids_triton from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import ( draft_tp_context, @@ -188,18 +185,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): ) self.init_lm_head() - # KV cache reversion buffer; sized to mirror req_to_token (indexed by - # req_pool_idx). - self.req_to_hidden_states_pool = torch.empty( - ( - self.req_to_token_pool.req_to_token.shape[0], - self.speculative_num_steps - 1, - self.model_config.hidden_size, - ), - dtype=self.model_config.dtype, - device=self.device, - ) - def init_attention_backends(self): with self.draft_tp_context( self.draft_runner_list[0].tp_group @@ -255,12 +240,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(self) ) - def reset_cuda_graph_buffers(self, forward_batch, batch_result): - if self.cuda_graph_runner_for_draft_extend: - self.cuda_graph_runner_for_draft_extend.reset_buffers( - forward_batch, batch_result - ) - def draft(self, batch: ScheduleBatch): draft_input: EagleDraftInput = batch.spec_info forward_batch, can_cuda_graph = self.prepare_for_draft( @@ -492,6 +471,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): forward_batch.extend_seq_lens, topk_index, ) + next_draft_input = EagleDraftInput( topk_p=torch.cat(topk_p_list, dim=1), topk_index=torch.cat(topk_index_list, dim=1), @@ -503,17 +483,6 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): num_tokens_for_logprob_per_req=1, ) - # Update req_to_hidden_states_pool for KV Cache reversion - if forward_batch.extend_seq_lens is not None: - assign_hidden_states_pool_triton( - target_hidden_states, - forward_batch.req_pool_indices, - self.req_to_hidden_states_pool, - self.speculative_num_steps - 1, - forward_batch.batch_size, - forward_batch.extend_seq_lens, - forward_batch.extend_start_loc, - ) return next_draft_input def _draft_extend_for_decode( @@ -543,6 +512,11 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): torch.get_device_module(self.device).current_stream().wait_stream( self.plan_stream ) + # `batch_result.accept_lens` includes the bonus token, so drafts-only + # is accept_lens - 1. Stash on spec_info for the cuda-graph prepare(). + forward_batch.spec_info.num_correct_drafts = batch_result.accept_lens - 1 + forward_batch.spec_info.num_accept_tokens = batch_result.accept_lens + # Run draft extend batch in the main compute stream can_cuda_graph = ( self.cuda_graph_runner_for_draft_extend @@ -553,7 +527,24 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): next_token_ids_backup = batch_result.next_token_ids.clone() if can_cuda_graph: - self.reset_cuda_graph_buffers(forward_batch, batch_result) + cgr = self.cuda_graph_runner_for_draft_extend + # Populate the single shared buffer set once; each step replays + # against it and the chain is advanced in place between steps. + cgr.prepare(forward_batch) + for step in range(self.speculative_num_steps): + _, ret_topk_p, ret_topk_index = cgr.replay(step) + ret_topk_p_list.append(ret_topk_p.clone()) + ret_topk_index_list.append(ret_topk_index.clone()) + # Advance the draft chain by rotating the shared input_ids window + # in place; step N+1's graph then reads the rotated values. + if step < self.speculative_num_steps - 1: + rotate_input_ids_triton( + cgr.buffers.input_ids[: cgr.raw_num_tokens], + cgr.buffers.extend_start_loc[: cgr.raw_bs], + cgr.buffers.extend_seq_lens[: cgr.raw_bs], + ret_topk_index, + cgr.buffers.select_index[: cgr.raw_bs], + ) else: logger.warning_once( "can't use cuda graph for draft extend! may have correctness issue!" @@ -572,24 +563,10 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): # its own metadata in forward_extend (post-pad), otherwise # per-runner attn_backend.forward_metadata is never initialized for # draft_runner_list[1+]. - if not _is_npu or can_cuda_graph: + if not _is_npu: forward_batch.mark_forward_metadata_ready() - for step in range(self.speculative_num_steps): - # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}") - if can_cuda_graph: - draft_logits_output = ( - self.cuda_graph_runner_for_draft_extend.get_runner(step).execute( - forward_batch, init_state=(step == 0) - ) - ) - ret_topk_p, ret_topk_index = ( - draft_logits_output.topk_p, - draft_logits_output.topk_index, - ) - else: - # Skip relies on the unconditional mark above (pre-existing - # no-pre-plan behavior preserved verbatim). + for step in range(self.speculative_num_steps): draft_logits_output = self.draft_runner_list[step].forward( forward_batch ) @@ -615,42 +592,9 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): ret_topk_index, select_index, ) - ret_topk_p_list.append(ret_topk_p) - ret_topk_index_list.append(ret_topk_index) + ret_topk_p_list.append(ret_topk_p) + ret_topk_index_list.append(ret_topk_index) - # Update req_to_hidden_states_pool for KV Cache reversion - if ( - forward_batch.extend_seq_lens is not None - and self.cuda_graph_runner_for_draft_extend is not None - ): - if can_cuda_graph: - last_runner = self.cuda_graph_runner_for_draft_extend.get_last_runner() - hidden_states = last_runner.buffers.hidden_states - req_pool_indices = last_runner.buffers.req_pool_indices - extend_seq_lens = last_runner.buffers.extend_seq_lens - extend_start_loc = last_runner.buffers.extend_start_loc - else: - hidden_states = draft_logits_output.logits_output.hidden_states - req_pool_indices = forward_batch.req_pool_indices - extend_seq_lens = forward_batch.extend_seq_lens - extend_start_loc = forward_batch.extend_start_loc - assign_hidden_states_pool_triton( - hidden_states, - req_pool_indices, - self.req_to_hidden_states_pool, - self.speculative_num_steps - 1, - forward_batch.batch_size, - extend_seq_lens, - extend_start_loc, - ) - - # Reorganize the spec info for the next batch - # draft_logits_output.next_token_logits = draft_logits_output.next_token_logits[ - # select_index - # ] - # draft_logits_output.hidden_states = draft_logits_output.hidden_states[ - # select_index - # ] batch_result.next_token_ids = next_token_ids_backup # Construct the return values next_draft_input = batch_result.next_draft_input diff --git a/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py b/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py index f1ce9d4b0..a4716c9a9 100644 --- a/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py +++ b/python/sglang/srt/speculative/triton_ops/multi_layer_eagle.py @@ -12,7 +12,6 @@ # limitations under the License. # ============================================================================== -import torch import triton import triton.language as tl @@ -70,281 +69,3 @@ def rotate_input_ids_triton( BLOCK_SIZE=BLOCK_SIZE, ) return input_ids - - -@triton.jit -def assign_new_state_kernel( - # Source pointers - old_input_ids_ptr, - old_positions_ptr, - old_hidden_states_ptr, - old_out_cache_loc_ptr, - old_extend_seq_lens_ptr, - old_extend_start_loc_ptr, - # Destination pointers - input_ids_ptr, - positions_ptr, - hidden_states_ptr, - out_cache_loc_ptr, - extend_seq_lens_ptr, - extend_start_loc_ptr, - # Auxiliary data pointers - next_token_ids_ptr, - seq_lens_ptr, - padding_lens_ptr, - req_pool_indices_ptr, - req_to_token_ptr, - req_to_hidden_states_pool_ptr, - # Scalars and Strides - step, - stride_hidden_seq, - stride_hidden_dim, # hidden_states strides - stride_pool_req, - stride_pool_step, - stride_pool_dim, # pool strides - stride_req_token_0, - stride_req_token_1, # req_to_token strides - # Meta-parameters - HIDDEN_DIM: tl.constexpr, - BLOCK_SEQ: tl.constexpr, - BLOCK_HID: tl.constexpr, -): - pid = tl.program_id(0) - - seq_len: tl.tensor = tl.load(seq_lens_ptr + pid) - old_extend_len = tl.load(old_extend_seq_lens_ptr + pid) - old_start = tl.load(old_extend_start_loc_ptr + pid) - new_extend_len = old_extend_len + 1 - new_start = old_start + pid - - tl.store(extend_seq_lens_ptr + pid, new_extend_len) - tl.store(extend_start_loc_ptr + pid, new_start) - - offs_seq = tl.arange(0, BLOCK_SEQ) - mask_seq = offs_seq < old_extend_len - - old_ids = tl.load(old_input_ids_ptr + old_start + offs_seq, mask=mask_seq) - tl.store(input_ids_ptr + new_start + offs_seq, old_ids, mask=mask_seq) - padding_len = tl.load(padding_lens_ptr + pid) - tl.store( - input_ids_ptr + new_start + old_extend_len - padding_len, - tl.load(next_token_ids_ptr + pid), - ) - - old_pos = tl.load(old_positions_ptr + old_start + offs_seq, mask=mask_seq) - tl.store(positions_ptr + new_start + 1 + offs_seq, old_pos, mask=mask_seq) - tl.store( - positions_ptr + new_start, max(tl.load(old_positions_ptr + old_start) - 1, 0) - ) - - old_cache = tl.load(old_out_cache_loc_ptr + old_start + offs_seq, mask=mask_seq) - tl.store(out_cache_loc_ptr + new_start + 1 + offs_seq, old_cache, mask=mask_seq) - - req_idx = tl.load(req_pool_indices_ptr + pid) - token_idx_col = seq_len - old_extend_len - 1 - if token_idx_col >= 0: - req_token_ptr_loc = ( - req_to_token_ptr - + (req_idx * stride_req_token_0) - + (token_idx_col * stride_req_token_1) - ) - last_cache_loc = tl.load(req_token_ptr_loc) - tl.store(out_cache_loc_ptr + new_start, last_cache_loc) - - pool_vec_offset_base = ((req_idx + 1) * stride_pool_req) + ( - -(step + 1) * stride_pool_step - ) - - for off_h in range(0, HIDDEN_DIM, BLOCK_HID): - offs_h = off_h + tl.arange(0, BLOCK_HID) - mask_h = offs_h < HIDDEN_DIM - - for i in range(BLOCK_SEQ): - if i < old_extend_len: - old_h_ptr = ( - old_hidden_states_ptr - + (old_start + i) * stride_hidden_seq - + (offs_h * stride_hidden_dim) - ) - new_h_ptr = ( - hidden_states_ptr - + (new_start + 1 + i) * stride_hidden_seq - + (offs_h * stride_hidden_dim) - ) - - chunk_old = tl.load(old_h_ptr, mask=mask_h) - tl.store(new_h_ptr, chunk_old, mask=mask_h) - - pool_ptrs = ( - req_to_hidden_states_pool_ptr - + pool_vec_offset_base - + (offs_h * stride_pool_dim) - ) - pool_val = tl.load(pool_ptrs, mask=mask_h) - - new_h_start_ptrs = ( - hidden_states_ptr - + (new_start * stride_hidden_seq) - + (offs_h * stride_hidden_dim) - ) - tl.store(new_h_start_ptrs, pool_val, mask=mask_h) - - -def assign_new_state_triton( - next_token_ids: torch.Tensor, - old_input_ids: torch.Tensor, - old_positions: torch.Tensor, - old_hidden_states: torch.Tensor, - old_out_cache_loc: torch.Tensor, - old_extend_seq_lens: torch.Tensor, - old_extend_start_loc: torch.Tensor, - input_ids: torch.Tensor, - positions: torch.Tensor, - hidden_states: torch.Tensor, - out_cache_loc: torch.Tensor, - extend_seq_lens: torch.Tensor, - extend_start_loc: torch.Tensor, - seq_lens: torch.Tensor, - padding_lens: torch.Tensor, - num_seqs: int, - step: int, - req_pool_indices: torch.Tensor, - req_to_token: torch.Tensor, - req_to_hidden_states_pool: torch.Tensor, -): - """ - Wrapper function to calculate offsets and launch the Triton kernel. - """ - hidden_dim = hidden_states.shape[1] - - BLOCK_SEQ = 8 - BLOCK_HID = 64 - - grid = (num_seqs,) - - assign_new_state_kernel[grid]( - # Pointers - old_input_ids, - old_positions, - old_hidden_states, - old_out_cache_loc, - old_extend_seq_lens, - old_extend_start_loc, - input_ids, - positions, - hidden_states, - out_cache_loc, - extend_seq_lens, - extend_start_loc, - next_token_ids, - seq_lens, - padding_lens, - req_pool_indices, - req_to_token, - req_to_hidden_states_pool, - # Constants/Strides - step, - old_hidden_states.stride(0), - old_hidden_states.stride(1), - req_to_hidden_states_pool.stride(0), - req_to_hidden_states_pool.stride(1), - req_to_hidden_states_pool.stride(2), - req_to_token.stride(0), - req_to_token.stride(1), - # Meta - HIDDEN_DIM=hidden_dim, - BLOCK_SEQ=BLOCK_SEQ, - BLOCK_HID=BLOCK_HID, - ) - - -@triton.jit -def assign_hidden_states_pool_kernel( - hidden_states_ptr, - req_pool_indices_ptr, - req_to_hidden_states_pool_ptr, - extend_seq_lens_ptr, - extend_start_loc_ptr, - stride_hidden_seq, - stride_hidden_dim, - stride_pool_req, - stride_pool_step, - stride_pool_dim, - HIDDEN_DIM: tl.constexpr, - pool_size: tl.constexpr, - BLOCK_HID: tl.constexpr, -): - pid = tl.program_id(0) - - extend_len = tl.load(extend_seq_lens_ptr + pid) - start_loc = tl.load(extend_start_loc_ptr + pid) - end_loc = start_loc + extend_len - - req_idx = tl.load(req_pool_indices_ptr + pid) - pool_vec_offset_base = req_idx * stride_pool_req - - for i in range(pool_size): - for off_h in range(0, HIDDEN_DIM, BLOCK_HID): - offs_h = off_h + tl.arange(0, BLOCK_HID) - mask_h = offs_h < HIDDEN_DIM - - hid_ptr = ( - hidden_states_ptr - + (end_loc - pool_size + i) * stride_hidden_seq - + offs_h * stride_hidden_dim - ) - hid_val = tl.load(hid_ptr, mask=mask_h) - - pool_ptr = ( - req_to_hidden_states_pool_ptr - + pool_vec_offset_base - + i * stride_pool_step - + offs_h * stride_pool_dim - ) - tl.store(pool_ptr, hid_val, mask=mask_h) - - -def assign_hidden_states_pool_triton( - hidden_states: torch.Tensor, - req_pool_indices: torch.Tensor, - req_to_hidden_states_pool: torch.Tensor, - pool_size: int, - num_seqs: int, - extend_seq_lens: torch.Tensor, - extend_start_loc: torch.Tensor, -): - grid = (num_seqs,) - assign_hidden_states_pool_kernel[grid]( - hidden_states, - req_pool_indices, - req_to_hidden_states_pool, - extend_seq_lens, - extend_start_loc, - hidden_states.stride(0), - hidden_states.stride(1), - req_to_hidden_states_pool.stride(0), - req_to_hidden_states_pool.stride(1), - req_to_hidden_states_pool.stride(2), - HIDDEN_DIM=hidden_states.shape[1], - pool_size=pool_size, - BLOCK_HID=64, - ) - - -def assign_hidden_states_pool_torch( - hidden_states: torch.Tensor, - req_pool_indices: torch.Tensor, - req_to_hidden_states_pool: torch.Tensor, - pool_size: int, - num_seqs: int, - extend_seq_lens: torch.Tensor, - extend_start_loc: torch.Tensor, -): - for req in range(num_seqs): - pool_idx = req_pool_indices[req] - extend_len = extend_seq_lens[req] - start_loc = extend_start_loc[req] - end_loc = start_loc + extend_len - req_to_hidden_states_pool[pool_idx, :pool_size, :].copy_( - hidden_states[end_loc - pool_size : end_loc, :] - )