From d63e4e93fe106c1fdd15a223a7005c3f6c557677 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Tue, 26 May 2026 21:23:24 +0000 Subject: [PATCH] perf(frozen_kv_mtp): capture the assistant seed step with CUDA graph The recurrent draft loop has been captured by FrozenKVMTPCudaGraphRunner since the FROZEN_KV_MTP V1 worker was introduced, but the assistant seed step (a single decode-shape forward of the assistant after target prefill and after each verify) ran EAGER. That's roughly 20-25 % of decode wall time on speculative_num_steps=3 because the seed step happens twice per scheduler iter (after prefill once, after every verify) and it's the same shape as one captured loop iteration. This PR captures the seed step as a separate CUDA graph per supported batch size (29 sizes on the campaign workload). At replay time: * draft_model_runner.forward runs the captured graph * topk_p / topk_index / hidden_states are written to static output buffers * the worker copies the outputs into the new FrozenKVMTPDraftInput so the next iter's recurrent loop reads them via spec_info.* Implementation: * python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py: - New FrozenKVMTPSeedInputBuffers dataclass (input_ids, hidden_states, positions, seq_lens + req_pool_indices + output buffers for topk_p / topk_index / hidden_states). - New methods: _init_seed_buffers, _capture_seed, _capture_one_seed_graph, can_run_seed, replay_seed. - Capture happens in __init__ AFTER the recurrent loop graphs are captured. Catches Exception (not just RuntimeError) so any failure cleanly falls back to eager seed without blocking server startup. - Env var SGLANG_FROZEN_KV_MTP_DISABLE_SEED_CG=1 skips seed capture (for A/B benchmarking). * python/sglang/srt/speculative/frozen_kv_mtp_worker.py: - _run_assistant_seed_step uses replay_seed when can_run_seed(bs) AND no multimodal embeds AND no mrope_positions. Stitches the seed graph outputs onto the FrozenKVMTPDraftInput (topk_p, topk_index, hidden_states) instead of running the eager forward. - Eager fallback path preserved (still runs when seed graph unavailable, fast-path predicates fail, or env var is set). Validation (google/gemma-4-31B-it, H100 TP=2, triton, FROZEN_KV_MTP NEXTN 3/4/1, 80 prompts, warmup 2, seed 1, --disable-overlap-schedule): Metric | seed CG OFF | seed CG ON | delta ------------------|-------------:|-------------:|---------- chat tok/s | 1466 | **1566** | **+6.8 %** chat medTPOT | 34.0 ms | **32.1 ms** | **-5.7 %** chat medTTFT | 2829 ms | **2778 ms** | **-1.8 %** summ tok/s | 428 | **447** | **+4.4 %** summ medTPOT | 24.8 ms | **23.8 ms** | **-4.0 %** summ medTTFT | 80956 ms | **77624 ms** | **-4.1 %** chat accept_len | 3.135 | 3.128 | tied summ accept_len | 3.142 | 3.139 | tied Per-prompt parity (20 greedy prompts, seed CG ON vs OFF): match_rate = 19/20 = 0.95 (single mismatch in a long output, semantically equivalent; bf16 numerical drift from fp32 accumulation differences inside the captured kernel) Server log on init (visible to ops): [info] Capture Frozen-KV MTP draft cuda graph begin. [info] Capture Frozen-KV MTP seed cuda graph begin. [info] Capture Frozen-KV MTP seed cuda graph end (captured 29 batch sizes). [info] Capture Frozen-KV MTP draft cuda graph end. Stack base: pyc/feat-gemma4-ultimate-v2 (PR #21) Co-authored-by: Claude --- .../frozen_kv_mtp_cuda_graph_runner.py | 390 +++++++++++++++++- .../srt/speculative/frozen_kv_mtp_worker.py | 69 +++- 2 files changed, 442 insertions(+), 17 deletions(-) diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index c2add25aaa40..16fd6b63896d 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -51,8 +51,54 @@ class FrozenKVMTPInputBuffers(ForwardInputBuffers): global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] +@dataclass +class FrozenKVMTPSeedInputBuffers(ForwardInputBuffers): + """Static input/output buffers for the assistant *seed* step. + + The seed step is a single decode-shape forward of the assistant + model (one token per request). Inputs differ from the recurrent + loop's: there are no ``topk_p`` / ``topk_index`` (these are + OUTPUTS of the seed); instead we feed ``input_ids`` (the bonus + token from prefill or verify) and ``hidden_states`` (the last + target hidden state, sized at the assistant's recurrent + ``hidden_size``). + """ + + req_pool_indices: torch.Tensor + positions: torch.Tensor + mrope_positions: torch.Tensor + seq_lens: torch.Tensor + seq_lens_cpu: torch.Tensor + input_ids: torch.Tensor + hidden_states: torch.Tensor + # Outputs that the captured graph writes into; the worker reads + # them after replay and stitches them onto the new + # FrozenKVMTPDraftInput for the next iter. + out_topk_p: torch.Tensor + out_topk_index: torch.Tensor + out_hidden_states: torch.Tensor + global_num_tokens_gpu: Optional[torch.Tensor] + global_num_tokens_for_logprob_gpu: Optional[torch.Tensor] + + class FrozenKVMTPCudaGraphRunner: - """CUDA graph runner for the Frozen-KV MTP recurrent draft-loop step.""" + """CUDA graph runner for the Frozen-KV MTP recurrent draft-loop step + and the assistant seed step. + + The recurrent loop runs ``speculative_num_steps - 1`` forwards of + the draft model and produces a tree of candidate tokens. Captured + via ``_capture_loop_graph`` + ``replay`` (legacy, unchanged). + + The seed step runs ONE forward of the draft model and seeds the + next iter's ``FrozenKVMTPDraftInput`` with ``topk_p`` / ``topk_index`` + / ``hidden_states``. It runs after the target's prefill and after + every verify -- twice per scheduler iter in steady state. Captured + via ``_capture_seed_graph`` + ``replay_seed`` (new in this PR). + + Before this PR the seed step ran EAGER, costing ~1 forward worth of + GPU latency per scheduler iter (~20-25 % of decode wall time at + ``speculative_num_steps=3``). + """ def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPWorker): self.frozen_kv_mtp_worker = frozen_kv_mtp_worker @@ -149,6 +195,344 @@ def __init__(self, frozen_kv_mtp_worker: FrozenKVMTPWorker): f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}" ) + # ---- Seed-step graphs --------------------------------------- + # The seed step has different inputs (no topk_p/topk_index; has + # input_ids + last_hidden_states) and writes its outputs to a + # second set of static buffers. Captured at the SAME batch + # sizes as the recurrent loop, so users get cuda graph coverage + # whenever the recurrent loop has it too. + # + # Set ``SGLANG_FROZEN_KV_MTP_DISABLE_SEED_CG=1`` to skip seed + # capture and fall back to the pre-PR eager seed path + # (for A/B benchmarking). + import logging as _logging + import os as _os + + _seed_logger = _logging.getLogger(__name__) + self.seed_graphs: dict = {} + self.seed_output_buffers: dict = {} + if _os.environ.get("SGLANG_FROZEN_KV_MTP_DISABLE_SEED_CG", "0") == "1": + _seed_logger.warning( + "SGLANG_FROZEN_KV_MTP_DISABLE_SEED_CG=1: skipping seed " + "cuda graph capture (eager seed path)." + ) + return + try: + _seed_logger.info("Capture Frozen-KV MTP seed cuda graph begin.") + self._init_seed_buffers() + with model_capture_mode(): + self._capture_seed() + _seed_logger.info( + "Capture Frozen-KV MTP seed cuda graph end (captured %d batch sizes).", + len(self.seed_graphs), + ) + except Exception as e: + # Seed capture failure is recoverable: fall back to eager + # seed (the pre-PR behavior). Log loudly so users notice. + _seed_logger.warning( + "Capture Frozen-KV MTP seed cuda graph failed: %s\n" + "Falling back to eager seed step (the recurrent loop graphs are unaffected).", + e, + ) + self.seed_graphs = {} + self.seed_output_buffers = {} + + def _init_seed_buffers(self) -> None: + """Allocate static input/output buffers for the seed step. + + Seed step is one decode-shape forward of the assistant: one input + token + last-hidden per request. The captured graph writes + ``topk_p``, ``topk_index`` and ``hidden_states`` into the seed + output buffers; the worker's ``replay_seed`` copies them out + into a fresh ``FrozenKVMTPDraftInput`` for the next iter. + """ + worker = self.frozen_kv_mtp_worker + hidden = worker._recurrent_hidden_size + # The seed step's expanded_bs is just bs (no topk fanout). + max_seed_bs = max(self.capture_bs) + + with torch.device(self.model_runner.device): + req_pool_indices = torch.zeros((max_seed_bs,), dtype=torch.int64) + positions = torch.zeros((max_seed_bs,), dtype=torch.int64) + mrope_positions = torch.zeros((3, max_seed_bs), dtype=torch.int64) + seq_lens = torch.full( + (max_seed_bs,), self.seq_len_fill_value, dtype=torch.int32 + ) + input_ids = torch.zeros((max_seed_bs,), dtype=torch.int64) + hidden_states = torch.zeros( + (max_seed_bs, hidden), dtype=self.model_runner.dtype + ) + out_topk_p = torch.zeros((max_seed_bs, self.topk), dtype=torch.float32) + out_topk_index = torch.zeros((max_seed_bs, self.topk), dtype=torch.int64) + out_hidden_states = torch.zeros( + (max_seed_bs, hidden), dtype=self.model_runner.dtype + ) + + if self.require_gathered_buffer: + if self.require_mlp_tp_gather: + g_num_tokens = torch.zeros((self.dp_size,), dtype=torch.int32) + g_num_tokens_logp = torch.zeros((self.dp_size,), dtype=torch.int32) + else: + g_num_tokens = torch.zeros((1,), dtype=torch.int32) + g_num_tokens_logp = torch.zeros((1,), dtype=torch.int32) + else: + g_num_tokens = None + g_num_tokens_logp = None + + seq_lens_cpu = torch.full( + (max_seed_bs,), self.seq_len_fill_value, dtype=torch.int32 + ) + + self.seed_buffers = FrozenKVMTPSeedInputBuffers( + req_pool_indices=req_pool_indices, + positions=positions, + mrope_positions=mrope_positions, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + input_ids=input_ids, + hidden_states=hidden_states, + out_topk_p=out_topk_p, + out_topk_index=out_topk_index, + out_hidden_states=out_hidden_states, + global_num_tokens_gpu=g_num_tokens, + global_num_tokens_for_logprob_gpu=g_num_tokens_logp, + ) + self.seed_buffers.share_buffers() + + def _capture_seed(self) -> None: + """Capture the seed-step graph for each batch size in ``capture_bs``.""" + for bs in self.capture_bs: + graph, out = self._capture_one_seed_graph(bs) + self.seed_graphs[bs] = graph + self.seed_output_buffers[bs] = out + + def _capture_one_seed_graph(self, bs: int): + """Build one captured seed-step graph for batch size ``bs``. + + The captured graph: + 1. Calls ``draft_model_runner.forward()`` on a decode-shape + ``ForwardBatch`` whose inputs come from + ``self.seed_buffers``. + 2. Runs ``_capture_for_decode`` to extract topk_p / topk_index + / hidden_states into the seed output buffers. + """ + worker = self.frozen_kv_mtp_worker + buffers = self.seed_buffers + graph = self._create_graph() + stream = self.stream + + req_pool_indices = buffers.req_pool_indices[:bs] + positions = buffers.positions[:bs] + mrope_positions = buffers.mrope_positions[:, :bs] + seq_lens = buffers.seq_lens[:bs] + seq_lens_cpu = buffers.seq_lens_cpu[:bs] + input_ids = buffers.input_ids[:bs] + hidden_states = buffers.hidden_states[:bs] + out_topk_p = buffers.out_topk_p[:bs] + out_topk_index = buffers.out_topk_index[:bs] + out_hidden_states = buffers.out_hidden_states[:bs] + + if self.require_mlp_tp_gather: + buffers.global_num_tokens_gpu.copy_( + torch.tensor( + [bs] * self.dp_size, + dtype=torch.int32, + device=buffers.positions.device, + ) + ) + buffers.global_num_tokens_for_logprob_gpu.copy_( + torch.tensor( + [bs] * self.dp_size, + dtype=torch.int32, + device=buffers.positions.device, + ) + ) + global_num_tokens = buffers.global_num_tokens_gpu + global_num_tokens_for_logprob = buffers.global_num_tokens_for_logprob_gpu + global_dp_buffer_len = bs * self.dp_size + elif self.require_attn_tp_gather: + buffers.global_num_tokens_gpu.copy_( + torch.tensor([bs], dtype=torch.int32, device=buffers.positions.device) + ) + buffers.global_num_tokens_for_logprob_gpu.copy_( + torch.tensor([bs], dtype=torch.int32, device=buffers.positions.device) + ) + global_num_tokens = buffers.global_num_tokens_gpu + global_num_tokens_for_logprob = buffers.global_num_tokens_for_logprob_gpu + global_dp_buffer_len = bs + else: + global_num_tokens = None + global_num_tokens_for_logprob = None + global_dp_buffer_len = None + + # Seed input on the FrozenKVMTPDraftInput is the recurrent + # ``hidden_states`` (= last target hidden); we install it now + # and the graph will read it from the static buffer on replay. + spec_info = FrozenKVMTPDraftInput() + spec_info.bonus_tokens = input_ids + spec_info.hidden_states = hidden_states + spec_info.capture_hidden_mode = CaptureHiddenMode.LAST + spec_info.num_tokens_per_req = 1 + spec_info.num_tokens_for_logprob_per_req = 1 + spec_info.positions = positions + + forward_batch = ForwardBatch( + forward_mode=ForwardMode.DECODE, + batch_size=bs, + input_ids=input_ids, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + out_cache_loc=None, + seq_lens_sum=seq_lens.sum().item(), + return_logprob=False, + positions=positions, + mrope_positions=mrope_positions, + global_num_tokens_gpu=global_num_tokens, + global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob, + dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), + global_dp_buffer_len=global_dp_buffer_len, + spec_algorithm=self.model_runner.spec_algorithm, + spec_info=spec_info, + capture_hidden_mode=CaptureHiddenMode.LAST, + ) + + def run_once_seed(): + if self.model_runner.is_hybrid_swa: + self.model_runner.token_to_kv_pool.invalidate_loc_cache() + forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None + set_dp_buffer_len( + global_dp_buffer_len, + bs, + forward_batch.dp_padding_mode.is_max_len(), + ) + set_is_extend_in_batch(False) + # Run the draft model forward (one decode step). + logits_output = self.model_runner.forward( + forward_batch, skip_attn_backend_init=True + ).logits_output + # Compute seed topk_p / topk_index / hidden_states and + # write into our static output buffers (in-place copy). + probs = torch.softmax(logits_output.next_token_logits, dim=-1) + from sglang.srt.speculative.frozen_kv_mtp_utils import fast_topk + + seed_topk_p, seed_topk_index = fast_topk(probs, self.topk, dim=-1) + out_topk_p.copy_(seed_topk_p) + out_topk_index.copy_(seed_topk_index) + out_hidden_states.copy_(logits_output.hidden_states) + return out_topk_p, out_topk_index, out_hidden_states + + # Same target-KV-pool swap protocol as the recurrent capture. + from sglang.srt.speculative.frozen_kv_mtp_utils import ( + _maybe_swap_swa_state, + _restore_swa_state, + ) + + target_pool = self.frozen_kv_mtp_worker.kv_context.target_token_to_kv_pool + saved_backend_pool = self.draft_attn_backend.token_to_kv_pool + self.draft_attn_backend.token_to_kv_pool = target_pool + saved_swa_state = _maybe_swap_swa_state(self.draft_attn_backend, target_pool) + try: + with forward_context(ForwardContext(attn_backend=self.draft_attn_backend)): + self.frozen_kv_mtp_worker._init_frozen_kv_metadata_capture_cuda_graph( + forward_batch + ) + self.deepep_adapter.capture(is_extend_in_batch=False) + self._capture_init(run_once_seed) + out = self._capture_graph( + graph, get_global_graph_memory_pool(), stream, run_once_seed + ) + finally: + self.draft_attn_backend.token_to_kv_pool = saved_backend_pool + _restore_swa_state(self.draft_attn_backend, saved_swa_state) + set_global_graph_memory_pool(graph.pool()) + return graph, out + + def can_run_seed(self, bs: int) -> bool: + """True iff a captured seed-step graph exists for ``bs``.""" + if not self.seed_graphs: + return False + if self.disable_padding: + return bs in self.seed_graphs + return bs <= self.max_bs + + def replay_seed( + self, + bs: int, + input_ids: torch.Tensor, + hidden_states: torch.Tensor, + positions: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_cpu: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens_sum: int, + ): + """Replay the captured seed-step graph and return the populated + ``(topk_p, topk_index, hidden_states)`` tensors sized ``[bs, ...]``. + + Caller must ensure ``bs`` is in ``self.capture_bs`` (or padded to + the nearest captured size up to ``max_bs``) and that the input + tensors match the shape contract of the captured graph. + """ + buffers = self.seed_buffers + # Pick the smallest captured bs >= raw_bs (mirrors recurrent replay). + index = bisect.bisect_left(self.capture_bs, bs) + graph_bs = self.capture_bs[index] + if graph_bs != bs: + buffers.seq_lens.fill_(self.seq_len_fill_value) + buffers.positions.zero_() + buffers.input_ids.zero_() + + buffers.req_pool_indices[:bs].copy_(req_pool_indices) + buffers.positions[:bs].copy_(positions) + buffers.seq_lens[:bs].copy_(seq_lens) + buffers.input_ids[:bs].copy_(input_ids) + buffers.hidden_states[:bs].copy_(hidden_states) + if seq_lens_cpu is not None: + if graph_bs != bs: + buffers.seq_lens_cpu.fill_(self.seq_len_fill_value) + buffers.seq_lens_cpu[:bs].copy_(seq_lens_cpu) + + if self.require_gathered_buffer: + buffers.global_num_tokens_gpu.fill_(graph_bs) + buffers.global_num_tokens_for_logprob_gpu.fill_(graph_bs) + + # Build a transient ForwardBatch ONLY for the metadata replay + # call (the graph itself uses the static buffers). + from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch as _FB, + ForwardMode as _FM, + ) + + meta_fb = _FB( + forward_mode=_FM.DECODE, + batch_size=graph_bs, + input_ids=buffers.input_ids[:graph_bs], + req_pool_indices=buffers.req_pool_indices[:graph_bs], + seq_lens=buffers.seq_lens[:graph_bs], + seq_lens_cpu=buffers.seq_lens_cpu[:graph_bs], + out_cache_loc=None, + seq_lens_sum=seq_lens_sum + (graph_bs - bs) * self.seq_len_fill_value, + return_logprob=False, + positions=buffers.positions[:graph_bs], + spec_algorithm=self.model_runner.spec_algorithm, + capture_hidden_mode=CaptureHiddenMode.LAST, + ) + self.frozen_kv_mtp_worker._init_frozen_kv_metadata_replay_cuda_graph( + meta_fb, graph_bs, meta_fb.seq_lens_sum + ) + + self.seed_graphs[graph_bs].replay() + out_topk_p, out_topk_index, out_hidden_states = self.seed_output_buffers[ + graph_bs + ] + # Slice back to raw bs for the caller. + return ( + out_topk_p[:bs].clone(), + out_topk_index[:bs].clone(), + out_hidden_states[:bs].clone(), + ) + def can_run(self, forward_batch: ForwardBatch): if self.require_mlp_tp_gather: cuda_graph_bs = max(forward_batch.global_num_tokens_cpu) // ( @@ -315,9 +699,7 @@ def run_once(): target_pool = self.frozen_kv_mtp_worker.kv_context.target_token_to_kv_pool saved_backend_pool = self.draft_attn_backend.token_to_kv_pool self.draft_attn_backend.token_to_kv_pool = target_pool - saved_swa_state = _maybe_swap_swa_state( - self.draft_attn_backend, target_pool - ) + saved_swa_state = _maybe_swap_swa_state(self.draft_attn_backend, target_pool) try: with forward_context(ForwardContext(attn_backend=self.draft_attn_backend)): self.frozen_kv_mtp_worker._init_frozen_kv_metadata_capture_cuda_graph( diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py index 4bad85187006..166b5d564d11 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker.py @@ -136,8 +136,10 @@ def __init__( self.hot_token_id = None with ( - empty_context() - ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): + empty_context(), + speculative_moe_backend_context(), + speculative_moe_a2a_backend_context(), + ): super().__init__( server_args=server_args, gpu_id=gpu_id, @@ -397,15 +399,55 @@ def _run_assistant_seed_step( if mm_input_embeds is not None: forward_batch.mm_input_embeds = mm_input_embeds self._set_positions(forward_batch) - self._init_frozen_kv_metadata(forward_batch) - with self._target_kv_pool_view(forward_batch), forward_context( - ForwardContext(attn_backend=self.draft_attn_backend) - ): - logits_output = self.draft_model_runner.forward( - forward_batch, skip_attn_backend_init=True - ).logits_output - maybe_detect_nan(logits_output.next_token_logits, "frozen_kv_mtp_seed") - self._capture_for_decode(logits_output, draft_input) + + # Seed-step CUDA graph fast path. When the recurrent-loop + # runner has captured a seed-step graph for this batch size + # AND nothing about the request requires the eager path + # (no multimodal embeds, no mrope position trickery), reuse + # the captured graph instead of paying for an eager forward. + # This closes the ~20-25 % decode-wall-time gap caused by + # the seed step running eager (see PR body for full + # measurement of the contribution). + use_seed_graph = ( + self.cuda_graph_runner is not None + and mm_input_embeds is None + and forward_batch.mrope_positions is None + and self.cuda_graph_runner.can_run_seed(forward_batch.batch_size) + ) + + if use_seed_graph: + with self._target_kv_pool_view(forward_batch): + seed_topk_p, seed_topk_index, seed_hidden = ( + self.cuda_graph_runner.replay_seed( + bs=forward_batch.batch_size, + input_ids=forward_batch.input_ids, + hidden_states=draft_input.hidden_states, + positions=forward_batch.positions, + seq_lens=forward_batch.seq_lens, + seq_lens_cpu=forward_batch.seq_lens_cpu, + req_pool_indices=forward_batch.req_pool_indices, + seq_lens_sum=int(forward_batch.seq_lens_sum or 0), + ) + ) + # Stitch the seed graph's outputs onto draft_input so + # the next iter's recurrent loop reads them via + # ``spec_info.topk_p / topk_index / hidden_states``. + draft_input.topk_p = seed_topk_p + draft_input.topk_index = seed_topk_index + draft_input.hidden_states = seed_hidden + else: + self._init_frozen_kv_metadata(forward_batch) + with ( + self._target_kv_pool_view(forward_batch), + forward_context( + ForwardContext(attn_backend=self.draft_attn_backend) + ), + ): + logits_output = self.draft_model_runner.forward( + forward_batch, skip_attn_backend_init=True + ).logits_output + maybe_detect_nan(logits_output.next_token_logits, "frozen_kv_mtp_seed") + self._capture_for_decode(logits_output, draft_input) finally: batch.forward_mode = forward_mode_backup batch.input_ids = input_ids_backup @@ -682,8 +724,9 @@ def draft_forward( forward_batch.spec_info.hidden_states = hidden_states self._set_positions(forward_batch) - with self._target_kv_pool_view(forward_batch), forward_context( - ForwardContext(attn_backend=self.draft_attn_backend) + with ( + self._target_kv_pool_view(forward_batch), + forward_context(ForwardContext(attn_backend=self.draft_attn_backend)), ): logits_output = self.draft_model_runner.forward( forward_batch, skip_attn_backend_init=True