Refactor / simplify MultiLayerEagleDraftExtendCudaGraphRunner to use rotation (#28492)
This commit is contained in:
Tarushii Goel
+6
-93
@@ -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()
|
||||
|
||||
+276
-320
@@ -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):
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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, :]
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user