diff --git a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py b/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py index 46ae1a6a7b72..53ddd7f27dc6 100644 --- a/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py +++ b/python/sglang/srt/layers/attention/linear/inkling_sconv_backend.py @@ -85,6 +85,8 @@ class InklingShortConvMetadata(msgspec.Struct): # [B, conv_kernel - 1] input positions whose conv window feeds the prefix # cache. Extend only, and only when tracking is on. track_conv_indices: Optional[torch.Tensor] = None + # Kernel-facing checkpoint slots. ForwardBatch retains the virtual source. + track_cache_indices: Optional[torch.Tensor] = None class InklingShortConvAttnBackend(ShortConvAttnBackend): @@ -144,7 +146,10 @@ def _alloc_graph_buffers(self): self._graph_track_conv_indices = torch.zeros( (max_bs, self.conv_state_len), dtype=torch.int64, device=dev ) - # Inert track fields for graph capture: a capture warmup batch that + # Prefill can capture before init_cuda_graph_state. Keep the translated + # checkpoint destinations at one address across both graph phases. + self._graph_track_indices = torch.empty(max_bs, dtype=torch.int64, device=dev) + # Inert prefill track fields (excluding draft extend): a warmup batch that # carries no tracking metadata must still LAUNCH the track scatter # (all rows masked off), or the python-level `if` specializes the # scatter out of the captured graph. @@ -155,8 +160,6 @@ def _alloc_graph_buffers(self): self._graph_track_inert_seqlens = torch.zeros( max_bs, dtype=torch.int64, device=dev ) - # Graph-static: captured track scatters read this address. - self._track_indices_buf = torch.zeros(max_bs, dtype=torch.int64, device=dev) # Same address-stability requirement; the base only sizes this from # init_cuda_graph_state, which the prefill graph never calls. self._alloc_cache_indices_buf(max_bs) @@ -250,6 +253,7 @@ def init_forward_metadata_capture_cpu_graph(self, *args, **kwargs): def _prepare_slot_indices(self, forward_batch: ForwardBatch): self._reset_step_state() + self._prepare_track_indices(forward_batch) req_pool_indices = forward_batch.req_pool_indices n = req_pool_indices.shape[0] buf = self._cache_indices_buf @@ -271,22 +275,22 @@ def _prepare_slot_indices(self, forward_batch: ForwardBatch): return self.forward_metadata = self._forward_metadata(forward_batch) self._refresh_cache_indices() - if not self._slot_gather_recordable: - self._translate_track_indices(forward_batch) - - def _translate_track_indices(self, forward_batch: ForwardBatch): - """Translate virtual track ids to the physical slots the conv kernels write.""" - track_indices = forward_batch.mamba_track_indices - buf = self._track_indices_buf - # Prep may re-run on the same batch; never translate twice. - if track_indices is None or track_indices.data_ptr() == buf.data_ptr(): + + def _prepare_track_indices(self, forward_batch: ForwardBatch): + # Every metadata view must explicitly provide virtual IDs or None. + indices = forward_batch.mamba_track_indices + if indices is None or self._slot_gather_recordable: + self.sconv_metadata.track_cache_indices = indices return - n = track_indices.shape[0] - assert n <= buf.shape[0], ( - f"track-index buffer too small: rows={n} vs bound {buf.shape[0]}" + n = indices.shape[0] + assert n <= self._graph_track_indices.shape[0], ( + "checkpoint-index buffer too small for the forward batch" ) - buf[:n].copy_(self._translate_mamba_indices(track_indices)) - forward_batch.mamba_track_indices = buf[:n] + # Keep the batch virtual across repeated eager metadata initialization + # and across draft backends. Captured consumers read this fixed buffer. + out = self._graph_track_indices[:n] + out.copy_(self._translate_mamba_indices(indices)) + self.sconv_metadata.track_cache_indices = out def _refresh_sconv_metadata( self, forward_batch: ForwardBatch, *, on_graph_path: bool @@ -451,6 +455,11 @@ def _refresh_track_conv_indices( ``batch_size`` rows while the track lengths cover only live requests, and every row it may read must index inside *this* replay's token buffer. """ + if forward_batch.forward_mode.is_draft_extend_v2(): + # Draft extend snapshots the accepted window in its own cache update. + # Inert prefill tracking would overwrite its checkpoint destinations. + self.sconv_metadata.track_conv_indices = None + return if forward_batch.mamba_track_mask is None: if not on_graph_path: self.sconv_metadata.track_conv_indices = None @@ -460,6 +469,9 @@ def _refresh_track_conv_indices( rows = forward_batch.batch_size forward_batch.mamba_track_mask = self._graph_track_inert_mask[:rows] forward_batch.mamba_track_indices = self._graph_track_inert_indices[:rows] + self.sconv_metadata.track_cache_indices = self._graph_track_inert_indices[ + :rows + ] forward_batch.mamba_track_seqlens = self._graph_track_inert_seqlens[:rows] rows = forward_batch.batch_size query_start_loc = self.sconv_metadata.query_start_loc @@ -532,6 +544,8 @@ def commit_conv_state_after_mtp_verify( slot_ids = self._translate_mamba_indices( pool.get_mamba_indices(req_pool_indices) ) + if mamba_track_indices is not None: + mamba_track_indices = self._translate_mamba_indices(mamba_track_indices) scatter_mamba_states_after_mtp_verify( pool.get_speculative_mamba2_params_all_layers(), slot_ids, @@ -603,6 +617,7 @@ def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): and self.supports_draft_extend_metadata_staging ): self.short_conv_backend._reset_step_state() + self.short_conv_backend._prepare_track_indices(forward_batch) self.short_conv_backend._refresh_sconv_metadata( forward_batch, on_graph_path=True ) diff --git a/python/sglang/srt/models/inkling_common/attn.py b/python/sglang/srt/models/inkling_common/attn.py index bf2914fa850e..76a623f5e833 100644 --- a/python/sglang/srt/models/inkling_common/attn.py +++ b/python/sglang/srt/models/inkling_common/attn.py @@ -725,7 +725,7 @@ def _fused_attn_prologue_decode(self, q, k, v, forward_batch, log_scaling_tau=No activation=self.k_sconv.activation, use_residual=self.k_sconv.use_residual, track_mask=forward_batch.mamba_track_mask, - track_indices=forward_batch.mamba_track_indices, + track_indices=self.k_sconv._conv_state(forward_batch).track_cache_indices, do_store=do_store, mxfp8_quant=do_mxfp8_store, sfk=sfk, diff --git a/python/sglang/srt/models/inkling_common/kernels/comm.py b/python/sglang/srt/models/inkling_common/kernels/comm.py index fa0fc14c61d0..fced9b720ed0 100644 --- a/python/sglang/srt/models/inkling_common/kernels/comm.py +++ b/python/sglang/srt/models/inkling_common/kernels/comm.py @@ -359,7 +359,7 @@ def ar_sconv_norm_fused( activation=sconv.activation, use_residual=sconv.use_residual, track_mask=forward_batch.mamba_track_mask, - track_indices=forward_batch.mamba_track_indices, + track_indices=sconv._conv_state(forward_batch).track_cache_indices, shared=shared, ) res.v5_cur = 1 - cur @@ -834,7 +834,7 @@ def ar_scattered_sconv_fused( if fm.is_decode() and forward_batch.mamba_track_mask is not None: b = forward_batch.batch_size track_mask = forward_batch.mamba_track_mask[:b] - track_dst = forward_batch.mamba_track_indices[:b] + track_dst = sconv._conv_state(forward_batch).track_cache_indices[:b] track_from_cache = True esz = comm.buffer.element_size() diff --git a/python/sglang/srt/models/inkling_common/sconv.py b/python/sglang/srt/models/inkling_common/sconv.py index ebfe6986e9d8..73cf6fb33f5c 100644 --- a/python/sglang/srt/models/inkling_common/sconv.py +++ b/python/sglang/srt/models/inkling_common/sconv.py @@ -198,7 +198,7 @@ def _apply_decode_sconv_kernel( activation=self.activation, use_residual=self.use_residual, track_mask=forward_batch.mamba_track_mask, - track_indices=forward_batch.mamba_track_indices, + track_indices=self._conv_state(forward_batch).track_cache_indices, ) def _prepare_extend_sconv_cache( @@ -216,7 +216,7 @@ def _prepare_extend_sconv_cache( sconv_cache=sconv_cache, track_conv_indices=track_conv_indices, mask=forward_batch.mamba_track_mask, - dst_indices=forward_batch.mamba_track_indices, + dst_indices=self._conv_state(forward_batch).track_cache_indices, ) def _save_intermediate_conv_windows( @@ -269,7 +269,7 @@ def _update_sconv_cache_for_draft_extend( else: draft_token_num = hidden_states.shape[1] - mamba_track_indices = getattr(forward_batch, "mamba_track_indices", None) + mamba_track_indices = self._conv_state(forward_batch).track_cache_indices do_tracking = ( mamba_track_indices is not None and get_exec().mamba.enable_mamba_extra_buffer @@ -349,7 +349,7 @@ def extend_fused_ar_inputs(self, forward_batch: ForwardBatch): if meta.track_conv_indices is not None: track_rows = meta.track_conv_indices track_mask = forward_batch.mamba_track_mask - track_dst = forward_batch.mamba_track_indices + track_dst = meta.track_cache_indices else: w1 = self.kernel_size[0] - 1 track_rows = torch.empty((0, w1), dtype=torch.int64, device=dev) diff --git a/python/sglang/srt/speculative/draft_checkpoint.py b/python/sglang/srt/speculative/draft_checkpoint.py new file mode 100644 index 000000000000..3d8c1bfa3eeb --- /dev/null +++ b/python/sglang/srt/speculative/draft_checkpoint.py @@ -0,0 +1,27 @@ +from typing import Optional + +import torch + + +def track_indices(buffer: Optional[torch.Tensor], bs: int): + return None if buffer is None else buffer[:bs] + + +def refresh_track_indices(buffer, src, *, raw_bs, bs, copy_dsts=None, copy_srcs=None): + """Refresh virtual destinations without clearing rows that will be copied. + + Ordinary EAGLE joins its existing grouped copy. The multi-layer runner + uses a fused preparation kernel for other inputs and copies here instead. + """ + if buffer is None: + return + if src is None: + buffer[:bs].zero_() + return + if raw_bs < bs: + buffer[raw_bs:bs].zero_() + if copy_dsts is None: + buffer[:raw_bs].copy_(src) + else: + copy_dsts.append(buffer[:raw_bs]) + copy_srcs.append(src) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 3dc941524b8f..edb036cd9e62 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -43,6 +43,7 @@ get_parallel, get_spec, ) +from sglang.srt.speculative.draft_checkpoint import refresh_track_indices, track_indices from sglang.srt.speculative.eagle_info import EagleDraftExtendInput from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim from sglang.srt.speculative.spec_utils import resolve_num_tokens_per_req @@ -74,6 +75,7 @@ def resolve_draft_extend_seq_len_fill_value( class EagleDraftExtendInputBuffers(ForwardInputBuffers): input_ids: torch.Tensor req_pool_indices: torch.Tensor + mamba_track_indices: Optional[torch.Tensor] out_cache_loc: torch.Tensor positions: torch.Tensor mrope_positions: torch.Tensor @@ -268,6 +270,11 @@ def __init__( self.buffers = EagleDraftExtendInputBuffers( input_ids=input_ids, req_pool_indices=req_pool_indices, + mamba_track_indices=( + torch.zeros(self.max_bs, dtype=torch.int64, device=self.device) + if get_exec().mamba.enable_mamba_extra_buffer + else None + ), out_cache_loc=out_cache_loc, positions=positions, mrope_positions=mrope_positions, @@ -421,6 +428,7 @@ def capture_one_shape( batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, + mamba_track_indices=track_indices(buffers.mamba_track_indices, bs), seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu, next_token_logits_buffer=next_token_logits_buffer, @@ -572,6 +580,14 @@ def execute(self, forward_batch: ForwardBatch, select_index: torch.Tensor): copy_srcs.append(forward_batch.spec_info.num_accept_tokens) copy_dsts.append(buffers.select_index[:raw_bs]) copy_srcs.append(select_index) + refresh_track_indices( + buffers.mamba_track_indices, + forward_batch.mamba_track_indices, + raw_bs=raw_bs, + bs=bs, + copy_dsts=copy_dsts, + copy_srcs=copy_srcs, + ) _grouped_foreach_copy_(copy_dsts, copy_srcs) # hidden_states is large + contiguous: copy_() uses the cudaMemcpyAsync @@ -633,6 +649,8 @@ def execute(self, forward_batch: ForwardBatch, select_index: torch.Tensor): encoder_lens=None, out_cache_loc=buffers.out_cache_loc[:num_tokens], out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), + # Virtual input stays separate from the backend's physical buffer. + mamba_track_indices=track_indices(buffers.mamba_track_indices, bs), spec_info=forward_batch.spec_info, ) self.draft_extend_attn_backend.init_forward_metadata_out_graph(fb_view) 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 33a45e57173d..cabb774b521e 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 @@ -66,6 +66,7 @@ get_parallel, get_spec, ) +from sglang.srt.speculative.draft_checkpoint import refresh_track_indices, track_indices from sglang.srt.speculative.eagle_info import EagleDraftExtendInput from sglang.srt.speculative.eagle_utils import get_draft_input_from_target_hidden_dim from sglang.srt.speculative.multi_layer_eagle_utils import ( @@ -116,6 +117,7 @@ class MultiLayerEagleDraftExtendInputBuffers(ForwardInputBuffers): seq_lens: torch.Tensor seq_lens_cpu: torch.Tensor req_pool_indices: torch.Tensor + mamba_track_indices: Optional[torch.Tensor] num_correct_drafts: torch.Tensor num_accept_tokens: torch.Tensor extend_seq_lens: torch.Tensor @@ -335,6 +337,7 @@ def get_forward_batch(self, bs: int) -> ForwardBatch: batch_size=bs, input_ids=input_ids, req_pool_indices=req_pool_indices, + mamba_track_indices=track_indices(buffers.mamba_track_indices, bs), seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu, next_token_logits_buffer=next_token_logits_buffer, @@ -493,6 +496,8 @@ def replay( encoder_lens=None, # per-step write target; out_cache_loc is frozen at prepare() time. out_cache_loc=buffers.out_cache_loc[:num_tokens], + # Virtual input stays separate from the backend's physical buffer. + mamba_track_indices=track_indices(buffers.mamba_track_indices, bs), spec_info=spec_info, ) if ( @@ -702,6 +707,11 @@ def _allocate_buffers(self) -> MultiLayerEagleDraftExtendInputBuffers: seq_lens=seq_lens, seq_lens_cpu=seq_lens_cpu, req_pool_indices=req_pool_indices, + mamba_track_indices=( + torch.zeros(self.max_bs, dtype=torch.int64, device=self.device) + if get_exec().mamba.enable_mamba_extra_buffer + else None + ), num_correct_drafts=num_correct_drafts, num_accept_tokens=num_accept_tokens, extend_seq_lens=extend_seq_lens, @@ -755,6 +765,7 @@ def _stage_metadata(self, bs: int, raw_bs: int): seq_lens=buffers.seq_lens[:bs], extend_seq_lens=buffers.extend_seq_lens[:bs], out_cache_loc=buffers.out_cache_loc[: bs * self.captured_req_width], + mamba_track_indices=track_indices(buffers.mamba_track_indices, bs), ) for backend in backends: backend.init_forward_metadata_out_graph(batch) @@ -773,6 +784,13 @@ def prepare(self, forward_batch: ForwardBatch, *, staged: bool = False): else: bs = self.get_runner(0)._pad_to_bucket(raw_bs, self.capture_bs) + refresh_track_indices( + buffers.mamba_track_indices, + forward_batch.mamba_track_indices, + raw_bs=raw_bs, + bs=bs, + ) + fill_draft_extend_prepare_buffers( buffers.input_ids, buffers.positions, diff --git a/test/registered/kernels/ops/attention/test_inkling_checkpoint_indices.py b/test/registered/kernels/ops/attention/test_inkling_checkpoint_indices.py new file mode 100644 index 000000000000..a53395082d12 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_inkling_checkpoint_indices.py @@ -0,0 +1,997 @@ +"""Inkling checkpoint writes use live physical slots, including graph replay.""" + +import unittest +from contextlib import ExitStack, nullcontext +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.layers.attention.linear.inkling_sconv_backend import ( + InklingShortConvAttnBackend, + InklingShortConvHybridAttnBackend, +) +from sglang.srt.mem_cache.allocator.unified_sub_pool import MultiEndedAllocator +from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool +from sglang.srt.mem_cache.unified_memory_pool import ( + MambaSubPoolSpec, + MHASubPoolSpec, + UnifiedHybridReqToTokenPool, + UnifiedKVPool, + UnifiedMambaPool, + UnifiedMambaSlotAllocator, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.srt.models.inkling_common.sconv import ShortConvolution +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +_BACKEND = "sglang.srt.layers.attention.linear.inkling_sconv_backend" + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") +class TestInklingCheckpointIndices(CustomTestCase): + def make_backend( + self, *, lazy=False, static=False, slot_count=6, draft_token_num=None + ): + mamba = MambaSubPoolSpec( + name="mamba", + layer_num=2, + grow_direction="up", + conv_state_shapes=((3, 64),), + conv_dtype=torch.bfloat16, + temporal_state_shape=(0,), + temporal_dtype=torch.float32, + ) + full = MHASubPoolSpec( + name="full", + layer_num=1, + head_num=1, + head_dim=8, + store_dtype=torch.bfloat16, + grow_direction="down", + ) + shared = UnifiedKVPool( + total_bytes=64 * mamba.entry_bytes(), + sub_pool_specs=[mamba, full], + device="cuda", + enable_memory_saver=False, + ) + shared._raw.zero_() + pool = UnifiedMambaPool( + unified_buffer=shared, + sub_pool_name="mamba", + spec_state_size=0, + mamba_layer_ids=[0, 1], + ) + allocator = MultiEndedAllocator( + kvcache=pool, + unified_buffer=shared, + sub_pool_name="mamba", + device="cuda", + is_id_owner=True, + lazy_compaction=lazy, + ) + slots = allocator.alloc(slot_count) + self.assertIsNotNone(slots) + req_type = HybridReqToTokenPool if static else UnifiedHybridReqToTokenPool + req_pool = req_type.__new__(req_type) + req_pool.size = 4 + req_pool.mamba_pool = pool + req_pool.mamba_allocator = UnifiedMambaSlotAllocator( + allocator, max_size=pool.size, device="cuda" + ) + req_pool.req_index_to_mamba_index_mapping = slots[1:2].to(torch.int32).repeat(4) + backend = InklingShortConvAttnBackend.__new__(InklingShortConvAttnBackend) + backend.req_to_token_pool = req_pool + backend.device = "cuda" + backend.conv_state_len = 3 + backend._cache_indices_buf = None + backend._slot_gather_recordable = static + config = SimpleNamespace( + decode=SimpleNamespace(bs=[1, 2, 4], max_bs=4), + prefill=SimpleNamespace(bs=[8]), + ) + with ( + patch( + _BACKEND + ".get_exec", + return_value=SimpleNamespace( + graph=SimpleNamespace(cuda_graph_config=config) + ), + ), + patch( + _BACKEND + ".get_spec", + return_value=SimpleNamespace( + speculative_num_draft_tokens=draft_token_num + ), + ), + ): + backend._alloc_graph_buffers() + return backend, allocator, pool, slots + + def batch(self, ids, mode=ForwardMode.EXTEND, mask=None): + n = len(ids) + return SimpleNamespace( + forward_mode=mode, + batch_size=n, + req_pool_indices=torch.arange(n, device="cuda"), + mamba_track_indices=ids, + mamba_track_mask=torch.ones(n, dtype=torch.bool, device="cuda") + if mask is None + else mask, + ) + + def scatter(self, backend, batch, cache): + hidden = ( + torch.arange(3 * 64, device="cuda", dtype=torch.float32) + .reshape(3, 64) + .to(torch.bfloat16) + ) + rows = torch.arange(3, device="cuda").repeat(batch.batch_size, 1) + ShortConvolution._prepare_extend_sconv_cache( + SimpleNamespace(_conv_state=lambda _: backend.sconv_metadata), + batch, + cache, + hidden, + rows, + ) + return hidden + + def test_prefill_checkpoint_after_real_compaction(self): + for lazy in (False, True): + with self.subTest(lazy=lazy): + backend, allocator, pool, slots = self.make_backend(lazy=lazy) + virtual = slots[-1:].clone() + allocator.free(slots[:1].clone()) + if lazy: + allocator._flush(urgent=True) + physical = backend._translate_mamba_indices(virtual) + self.assertFalse(torch.equal(virtual, physical.to(virtual.dtype))) + ids = torch.cat([virtual, virtual.new_tensor([-1])]) + original = ids.clone() + batch = self.batch(ids, mask=torch.tensor([True, False], device="cuda")) + backend._prepare_slot_indices(batch) + cache = pool.mamba_cache.conv[0][0] + before = cache.clone() + expected = self.scatter(backend, batch, cache) + torch.testing.assert_close(cache[physical[0]], expected, rtol=0, atol=0) + before[physical[0]] = expected + torch.testing.assert_close(cache, before, rtol=0, atol=0) + torch.testing.assert_close(ids, original, rtol=0, atol=0) + self.assertEqual( + backend.sconv_metadata.track_cache_indices.dtype, torch.int64 + ) + + def test_replay_refreshes_captured_destination_after_compaction(self): + backend, allocator, pool, slots = self.make_backend(lazy=True) + virtual = slots[-1:].clone() + batch = self.batch(virtual, mode=ForwardMode.DECODE) + backend.init_forward_metadata_out_graph(batch) + pointer = backend.sconv_metadata.track_cache_indices.data_ptr() + cache = pool.mamba_cache.conv[0][0] + self.scatter(backend, batch, cache) # compile before capture + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + self.scatter(backend, batch, cache) + allocator.free(slots[:1].clone()) + allocator._flush(urgent=True) + physical = backend._translate_mamba_indices(virtual) + cache.zero_() + fresh = self.batch(virtual, mode=ForwardMode.DECODE) + backend.init_forward_metadata_out_graph(fresh) + self.assertEqual(backend.sconv_metadata.track_cache_indices.data_ptr(), pointer) + graph.replay() + expected = ( + torch.arange(3 * 64, device="cuda", dtype=torch.float32) + .reshape(3, 64) + .to(torch.bfloat16) + ) + torch.testing.assert_close(cache[physical[0]], expected, rtol=0, atol=0) + self.assertEqual( + torch.count_nonzero(cache).item(), torch.count_nonzero(expected).item() + ) + torch.testing.assert_close(virtual, slots[-1:], rtol=0, atol=0) + + def test_decode_tracks_the_physical_checkpoint(self): + backend, allocator, pool, slots = self.make_backend() + allocator.free(slots[:1].clone()) + batch = self.batch(slots[-1:].clone(), mode=ForwardMode.DECODE) + backend._prepare_slot_indices(batch) + cache = pool.mamba_cache.conv[0][0] + conv = SimpleNamespace( + activation=None, + _conv_state=lambda _: backend.sconv_metadata, + use_residual=True, + _weight_2d=lambda: torch.ones(64, 4, dtype=torch.bfloat16, device="cuda"), + ) + ShortConvolution._apply_decode_sconv_kernel( + conv, + torch.ones(1, 64, dtype=torch.bfloat16, device="cuda"), + cache, + backend._cache_indices, + {"cache_mask": torch.ones(1, 1, 1, dtype=torch.bool, device="cuda")}, + batch, + ) + physical = backend._translate_mamba_indices(slots[-1:]) + torch.testing.assert_close( + cache[physical[0]], cache[backend._cache_indices[0]], rtol=0, atol=0 + ) + self.assertGreater(torch.count_nonzero(cache[physical[0]]).item(), 0) + + def test_static_identity_path_and_disabled_tracking(self): + backend, _, _, slots = self.make_backend(static=True) + ids = slots[-1:].clone() + batch = self.batch(ids) + backend._prepare_slot_indices(batch) + self.assertIs(batch.mamba_track_indices, ids) + batch.mamba_track_indices = None + backend._prepare_slot_indices(batch) + self.assertIsNone(batch.mamba_track_indices) + backend._slot_gather_recordable = False + backend._prepare_slot_indices(batch) + self.assertIsNone(batch.mamba_track_indices) + + def test_graph_hooks_translate_once_for_each_mode(self): + backend, allocator, _, slots = self.make_backend() + allocator.free(slots[:1].clone()) + for mode in ( + ForwardMode.DECODE, + ForwardMode.TARGET_VERIFY, + ForwardMode.DRAFT_EXTEND_V2, + ): + with self.subTest(mode=mode): + ids = slots[-1:].clone() + batch = self.batch(ids, mode=mode) + with ( + patch.object(backend, "_refresh_sconv_metadata"), + patch.object( + backend, + "_translate_mamba_indices", + wraps=backend._translate_mamba_indices, + ) as translate, + ): + backend.init_forward_metadata_out_graph(batch) + calls = translate.call_count + self.assertEqual(calls, 2) # active and checkpoint slots + backend.init_forward_metadata_in_graph(batch) + self.assertEqual(translate.call_count, calls) + torch.testing.assert_close( + backend.sconv_metadata.track_cache_indices, + backend._translate_mamba_indices(ids).to(torch.int64), + ) + torch.testing.assert_close(ids, slots[-1:]) + + def test_distinct_strided_destinations_keep_static_storage(self): + backend, allocator, pool, slots = self.make_backend() + allocator.free(slots[:1].clone()) + selected = slots[2:].flip(0) + storage = torch.stack( + [selected, slots[1:2].expand_as(selected)], dim=1 + ).flatten() + original = storage.clone() + ids = storage[::2] + self.assertFalse(ids.is_contiguous()) + batch = self.batch(ids) + backend._prepare_slot_indices(batch) + pointer = backend.sconv_metadata.track_cache_indices.data_ptr() + physical = backend._translate_mamba_indices(selected).to(torch.int64) + torch.testing.assert_close( + backend.sconv_metadata.track_cache_indices, physical, rtol=0, atol=0 + ) + cache = pool.mamba_cache.conv[0][0] + expected = cache.clone() + hidden = torch.arange(12 * 64, device="cuda").reshape(12, 64).to(torch.bfloat16) + rows = torch.arange(12, device="cuda").reshape(4, 3) + ShortConvolution._prepare_extend_sconv_cache( + SimpleNamespace(_conv_state=lambda _: backend.sconv_metadata), + batch, + cache, + hidden, + rows, + ) + expected[physical] = hidden.reshape(4, 3, 64) + torch.testing.assert_close(cache, expected, rtol=0, atol=0) + smaller = self.batch(ids[:2]) + backend._prepare_slot_indices(smaller) + self.assertEqual(backend.sconv_metadata.track_cache_indices.data_ptr(), pointer) + torch.testing.assert_close(storage, original, rtol=0, atol=0) + + def test_decode_graph_refreshes_distinct_destinations_and_masks_padding(self): + backend, allocator, pool, slots = self.make_backend(lazy=True, slot_count=12) + req_pool = backend.req_to_token_pool + req_pool.req_index_to_mamba_index_mapping = torch.cat( + [slots.new_zeros(1), slots[1:5]] + ).to(torch.int32) + ids = slots[-4:].clone() + track_mask = torch.ones(4, dtype=torch.bool, device="cuda") + batch = self.batch(ids, mode=ForwardMode.DECODE, mask=track_mask) + batch.req_pool_indices = torch.arange(1, 5, device="cuda") + backend.init_forward_metadata_out_graph(batch) + pointer = backend.sconv_metadata.track_cache_indices.data_ptr() + cache = pool.mamba_cache.conv[0][0] + hidden = torch.zeros(4, 64, dtype=torch.bfloat16, device="cuda") + weight = torch.ones(64, 4, dtype=torch.bfloat16, device="cuda") + conv = SimpleNamespace( + activation=None, + use_residual=True, + _weight_2d=lambda: weight, + _conv_state=lambda _: backend.sconv_metadata, + ) + + def decode(): + backend.init_forward_metadata_in_graph(batch) + ShortConvolution._apply_decode_sconv_kernel( + conv, + hidden, + cache, + backend._cache_indices, + backend.sconv_metadata.precomputed, + batch, + ) + + decode() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + decode() + for step, live in enumerate((3, 1, 4)): + with self.subTest(live=live): + if step < 2: + hole = slots[:1] if step == 0 else slots[5:6] + allocator.free(hole.clone()) + allocator._flush(urgent=True) + ids.copy_(slots[-4:].roll(step + 1)) + ids[live:].zero_() + original = ids.clone() + track_mask.zero_() + track_mask[:live] = True + if live > 1: + track_mask[1] = False + fresh = self.batch(ids, mode=ForwardMode.DECODE, mask=track_mask) + fresh.req_pool_indices = torch.arange(1, 5, device="cuda") + fresh.req_pool_indices[live:] = 0 + backend.init_forward_metadata_out_graph(fresh) + self.assertEqual( + backend.sconv_metadata.track_cache_indices.data_ptr(), pointer + ) + physical = backend._translate_mamba_indices(ids).to(torch.int64) + active = backend._cache_indices.clone().to(torch.int64) + cache.fill_(-7) + cache[0].zero_() + hidden.copy_(torch.arange(4, device="cuda")[:, None] + step + 1) + hidden[live:].zero_() + expected = cache.clone() + for row in range(live): + window = torch.cat([cache[active[row], 1:], hidden[row : row + 1]]) + expected[active[row]] = window + if row != 1: + expected[physical[row]] = window + graph.replay() + torch.testing.assert_close(cache, expected, rtol=0, atol=0) + torch.testing.assert_close(ids, original, rtol=0, atol=0) + + def test_oversized_checkpoint_batch_fails_without_reallocating(self): + backend, _, _, slots = self.make_backend() + pointer = backend._graph_track_indices.data_ptr() + batch = self.batch(slots[-1:].repeat(5)) + with self.assertRaisesRegex( + AssertionError, "checkpoint-index buffer too small" + ): + backend._prepare_slot_indices(batch) + self.assertEqual(backend._graph_track_indices.data_ptr(), pointer) + + def test_verify_commit_translates_both_destinations(self): + backend, allocator, _, slots = self.make_backend() + allocator.free(slots[:1].clone()) + indices = torch.tensor([0], device="cuda") + for tracking in (False, True): + with self.subTest(tracking=tracking): + ids = slots[-1:].clone() if tracking else None + with ( + patch.object( + backend.req_to_token_pool, + "get_speculative_mamba2_params_all_layers", + return_value=object(), + ), + patch( + _BACKEND + ".scatter_mamba_states_after_mtp_verify" + ) as scatter, + ): + backend.commit_conv_state_after_mtp_verify( + req_pool_indices=indices, + last_correct_step_indices=indices, + mamba_track_indices=ids, + mamba_steps_to_track=indices if tracking else None, + ) + if tracking: + torch.testing.assert_close( + scatter.call_args.args[3], backend._translate_mamba_indices(ids) + ) + torch.testing.assert_close(ids, slots[-1:], rtol=0, atol=0) + else: + self.assertIsNone(scatter.call_args.args[3]) + self.assertIsNone(scatter.call_args.args[4]) + torch.testing.assert_close( + scatter.call_args.args[1], + backend._translate_mamba_indices(slots[1:2].to(torch.int32)), + ) + + def test_explicit_none_tracking_still_refreshes_active_slots(self): + backend, allocator, _, slots = self.make_backend() + allocator.free(slots[:1].clone()) + batch = SimpleNamespace( + forward_mode=ForwardMode.DRAFT_EXTEND_V2, + req_pool_indices=torch.tensor([0], device="cuda"), + mamba_track_indices=None, + ) + backend.init_forward_metadata_out_graph(batch) + self.assertIsNone(backend.sconv_metadata.track_cache_indices) + torch.testing.assert_close( + backend._cache_indices, + backend._translate_mamba_indices(slots[1:2]).to(torch.int32), + ) + + def test_missing_tracking_field_rejects_incomplete_metadata_view(self): + for static in (False, True): + with self.subTest(static=static): + backend, _, _, slots = self.make_backend(static=static) + batch = self.batch(slots[-1:].clone()) + backend._prepare_slot_indices(batch) + del batch.mamba_track_indices + with self.assertRaisesRegex(AttributeError, "mamba_track_indices"): + backend._prepare_slot_indices(batch) + + def test_multilayer_staging_provides_checkpoint_ids_or_none(self): + from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import ( + MultiLayerEagleMultiStepDraftExtendCudaGraphRunner, + ) + + _, backend, _, pool, slots, buffers = self._make_draft_checkpoint_runner( + multilayer=True, static=True, tracking=True + ) + wrapper = InklingShortConvHybridAttnBackend.__new__( + InklingShortConvHybridAttnBackend + ) + wrapper.short_conv_backend = backend + wrapper.full_attn_backend = Mock(supports_draft_extend_metadata_staging=True) + composite = MultiLayerEagleMultiStepDraftExtendCudaGraphRunner.__new__( + MultiLayerEagleMultiStepDraftExtendCudaGraphRunner + ) + composite.buffers = buffers + composite.captured_req_width = 4 + composite.draft_extend_attn_backend_list = [wrapper] + for step, tracking in enumerate((True, False, True)): + with self.subTest(tracking=tracking, step=step): + ids = slots[-4:].roll(step).clone() + original = ids.clone() + buffers.mamba_track_indices = ids if tracking else None + composite._stage_metadata(bs=3, raw_bs=2) + torch.testing.assert_close( + backend._cache_indices, + torch.cat([slots[1:3], slots.new_zeros(1)]).to(torch.int32), + ) + if tracking: + torch.testing.assert_close( + backend.sconv_metadata.track_cache_indices, ids[:3] + ) + self.assertEqual( + backend.sconv_metadata.track_cache_indices.data_ptr(), + ids.data_ptr(), + ) + batch = self.batch( + ids[:3], mask=torch.tensor([True, True, False], device="cuda") + ) + cache = pool.mamba_cache.conv[0][0] + expected = cache.clone() + hidden = self.scatter(backend, batch, cache) + expected[ids[:2].long()] = hidden + torch.testing.assert_close(cache, expected, rtol=0, atol=0) + torch.testing.assert_close(ids, original, rtol=0, atol=0) + else: + self.assertIsNone(backend.sconv_metadata.track_cache_indices) + + def test_draft_eager_preplan_and_repeated_steps_preserve_virtual_source(self): + from sglang.srt.speculative import eagle_worker_common as common + from sglang.srt.speculative.eagle_info import EagleDraftExtendInput + + for static in (False, True): + for graph_available in (False, True): + with self.subTest(static=static, graph_available=graph_available): + steps = [] + for step in range(2): + backend, allocator, pool, slots = self.make_backend( + static=static, slot_count=12, draft_token_num=4 + ) + backend.mamba_cache_chunk_size = 4 + backend.req_to_token_pool.req_index_to_mamba_index_mapping.fill_( + slots[3] + ) + if not static: + allocator.free(slots[step : step + 1].clone()) + wrapper = InklingShortConvHybridAttnBackend.__new__( + InklingShortConvHybridAttnBackend + ) + wrapper.full_attn_backend = Mock() + wrapper.short_conv_backend = backend + wrapper.attn_backend_list = [wrapper.full_attn_backend, backend] + steps.append((wrapper, backend, pool)) + ids = slots[-1:].clone() + original = ids.clone() + batch = self.batch(ids, mode=ForwardMode.DRAFT_EXTEND_V2) + batch.seq_lens = torch.tensor([2], device="cuda") + batch.extend_seq_lens = torch.tensor([4], device="cuda") + batch.extend_prefix_lens = None + batch.extend_num_tokens = 4 + batch.spec_info = EagleDraftExtendInput( + num_accept_tokens=torch.tensor([3], device="cuda") + ) + batch.mark_forward_metadata_ready = Mock() + schedule_batch = SimpleNamespace( + seq_lens=batch.seq_lens, + seq_lens_cpu=None, + forward_mode=batch.forward_mode, + model_config=SimpleNamespace(vocab_size=32), + ) + runner = SimpleNamespace( + spec_algorithm=SimpleNamespace(is_standalone=lambda: False), + attn_backend=steps[0][0], + model_config=SimpleNamespace(model_is_mrope=False), + ) + graph_runner = ( + SimpleNamespace(can_run_graph=lambda _: False) + if graph_available + else None + ) + with ( + patch.object( + common.ForwardBatch, "init_new", return_value=batch + ), + patch.object(common, "maybe_detect_oob"), + patch.object(common, "is_npu", return_value=False), + ): + planned = common.prepare_for_draft_extend( + batch.spec_info, + schedule_batch, + torch.zeros(4, dtype=torch.int64, device="cuda"), + 4, + runner, + graph_runner, + return_hidden_states_before_norm=False, + ) + self.assertIs(planned, batch) + destinations = [] + # Inspect the real pre-plan, then repeat step 0 and enter step 1. + for index in (0, 0, 1): + wrapper, backend, pool = steps[index] + if destinations: + wrapper.init_forward_metadata(batch) + self.assertIs(batch.mamba_track_indices, ids) + torch.testing.assert_close(ids, original, rtol=0, atol=0) + physical = backend._translate_mamba_indices(ids).long() + torch.testing.assert_close( + backend.sconv_metadata.track_cache_indices, + physical, + rtol=0, + atol=0, + ) + destinations.append(physical.clone()) + cache = pool.mamba_cache.conv[0][0] + cache.fill_(-7) + hidden = torch.arange(4, device="cuda")[:, None].expand(4, 64) + hidden = hidden.to(torch.bfloat16) + active = backend._cache_indices.long() + joined = torch.cat([cache[active[0]], hidden]) + expected = cache.clone() + expected[physical[0]] = joined[1:4] + expected[active[0]] = joined[3:6] + with patch( + "sglang.srt.models.inkling_common.sconv.get_exec", + return_value=SimpleNamespace( + mamba=SimpleNamespace( + enable_mamba_extra_buffer=True, + mamba_track_interval=4, + ) + ), + ): + ShortConvolution._update_sconv_cache_for_draft_extend( + SimpleNamespace( + _conv_state=lambda _: backend.sconv_metadata + ), + batch, + cache, + backend._cache_indices, + hidden, + ) + torch.testing.assert_close(cache, expected, rtol=0, atol=0) + if not static: + self.assertFalse(torch.equal(destinations[0], destinations[-1])) + + def test_static_draft_staging_keeps_checkpoint_metadata(self): + from sglang.srt.speculative.eagle_info import EagleDraftExtendInput + + backend, _, _, slots = self.make_backend(static=True, draft_token_num=4) + wrapper = InklingShortConvHybridAttnBackend.__new__( + InklingShortConvHybridAttnBackend + ) + wrapper.short_conv_backend = backend + wrapper.full_attn_backend = Mock(supports_draft_extend_metadata_staging=True) + batch = self.batch(slots[-1:].clone(), mode=ForwardMode.DRAFT_EXTEND_V2) + batch.seq_lens = torch.tensor([6], device="cuda") + batch.extend_seq_lens = torch.tensor([4], device="cuda") + batch.extend_prefix_lens = None + batch.extend_num_tokens = 4 + batch.spec_info = EagleDraftExtendInput( + num_accept_tokens=torch.tensor([3], device="cuda") + ) + wrapper.init_forward_metadata_out_graph(batch, in_capture=True) + wrapper.init_forward_metadata_in_graph(batch) + self.assertIs( + backend.sconv_metadata.track_cache_indices, batch.mamba_track_indices + ) + + def test_prefill_capture_reinitializes_inert_checkpoint_metadata(self): + backend, _, pool, _ = self.make_backend() + backend.mamba_cache_chunk_size = 4 + batch = self.batch(torch.zeros(1, dtype=torch.int64, device="cuda")) + batch.mamba_track_indices = None + batch.mamba_track_mask = None + batch.seq_lens = torch.tensor([4], device="cuda") + batch.extend_seq_lens = torch.tensor([4], device="cuda") + batch.extend_prefix_lens = None + batch.extend_num_tokens = 4 + batch.spec_info = None + cache = pool.mamba_cache.conv[0][0] + expected = cache.clone() + for _ in range(2): + backend.init_forward_metadata_out_graph(batch, in_capture=True) + torch.testing.assert_close( + backend.sconv_metadata.track_cache_indices, + torch.zeros(1, dtype=torch.int64, device="cuda"), + ) + self.scatter(backend, batch, cache) + torch.testing.assert_close(cache, expected, rtol=0, atol=0) + + def test_draft_runners_capture_and_replay_checkpoint_destinations(self): + for multilayer in (False, True): + for static in (False, True): + for tracking in (False, True): + with self.subTest( + multilayer=multilayer, static=static, tracking=tracking + ): + self._check_draft_checkpoint_replay( + multilayer, static, tracking + ) + + def _allocate_production_draft_buffers(self, multilayer, tracking): + from sglang.srt.speculative import eagle_draft_extend_cuda_graph_runner as eagle + from sglang.srt.speculative import ( + multi_layer_eagle_draft_extend_cuda_graph_runner as multi, + ) + + model_runner = SimpleNamespace( + device="cuda", + tp_size=1, + model_config=SimpleNamespace( + dtype=torch.bfloat16, vocab_size=8, hf_config=SimpleNamespace() + ), + graph_shared_output=SimpleNamespace( + get_logits_buffer=lambda vocab, rows: torch.zeros( + rows, vocab, device="cuda" + ) + ), + ) + worker = SimpleNamespace( + draft_runner=model_runner, + speculative_algorithm=SimpleNamespace(is_standalone=lambda: False), + seed_dsa_topk_from_draft_extend=False, + mtp_model_runner=lambda step: model_runner, + ) + runtime = SimpleNamespace( + mamba=SimpleNamespace(enable_mamba_extra_buffer=tracking), + graph=SimpleNamespace( + disable_cuda_graph_padding=False, enable_profile_cuda_graph=False + ), + ) + if multilayer: + composite = ( + multi.MultiLayerEagleMultiStepDraftExtendCudaGraphRunner.__new__( + multi.MultiLayerEagleMultiStepDraftExtendCudaGraphRunner + ) + ) + composite.runners = [SimpleNamespace(model_runner=model_runner)] + composite.eagle_worker = worker + composite.max_bs = 4 + composite.captured_req_width = 4 + composite.seq_len_fill_value = 4 + composite.device = "cuda" + composite.prune_draft_extend_logits = False + composite.rotates_in_graph = False + composite.require_gathered_buffer = False + with ( + patch.object(multi, "get_exec", return_value=runtime), + patch.object( + multi, "get_draft_input_from_target_hidden_dim", return_value=64 + ), + ): + return composite._allocate_buffers() + + worker.draft_extend_attn_backend = SimpleNamespace( + init_cuda_graph_state=lambda *args: None, + get_cuda_graph_seq_len_fill_value=lambda: 4, + ) + # Only model/distributed capture scaffolding is replaced. __init__ and + # its real input-buffer allocation run before the kernel replay test. + with ExitStack() as stack: + replacements = { + "get_exec": runtime, + "get_parallel": SimpleNamespace(pp_size=1), + "deployment_attn_dp_size": 1, + "get_flags": SimpleNamespace( + capture=SimpleNamespace(enable_torch_compile=False) + ), + "get_spec": SimpleNamespace(speculative_num_steps=1), + "get_batch_sizes_to_capture": ([4], []), + "resolve_num_tokens_per_req": 4, + "get_draft_input_from_target_hidden_dim": 64, + "require_gathered_buffer": False, + "require_mlp_tp_gather": False, + "require_mlp_sync": False, + "require_attn_tp_gather": False, + "resolve_decode_backend": Mock(), + } + for name, result in replacements.items(): + stack.enter_context(patch.object(eagle, name, return_value=result)) + stack.enter_context(patch.object(eagle, "model_capture_mode", nullcontext)) + stack.enter_context(patch.object(eagle, "DeepEPCudaGraphRunnerAdapter")) + stack.enter_context( + patch.object(eagle.EagleDraftExtendInputBuffers, "share_buffers") + ) + stack.enter_context( + patch.object(eagle.EAGLEDraftExtendCudaGraphRunner, "capture") + ) + return eagle.EAGLEDraftExtendCudaGraphRunner(worker).buffers + + def _make_draft_checkpoint_runner(self, multilayer, static, tracking): + from sglang.srt.speculative import eagle_draft_extend_cuda_graph_runner as eagle + from sglang.srt.speculative import ( + multi_layer_eagle_draft_extend_cuda_graph_runner as multi, + ) + + backend, allocator, pool, slots = self.make_backend( + lazy=True, static=static, slot_count=12, draft_token_num=4 + ) + backend.mamba_cache_chunk_size = 4 + req_pool = backend.req_to_token_pool + req_pool.req_index_to_mamba_index_mapping = torch.cat( + [slots.new_zeros(1), slots[1:5]] + ).to(torch.int32) + buffers = self._allocate_production_draft_buffers(multilayer, tracking) + buffers.req_pool_indices.copy_(torch.arange(1, 5, device="cuda")) + buffers.seq_lens.fill_(6) + buffers.seq_lens_cpu.fill_(6) + buffers.num_accept_tokens.fill_(3) + if tracking: + self.assertIsNotNone(buffers.mamba_track_indices) + self.assertEqual(buffers.mamba_track_indices.shape, (4,)) + self.assertEqual(buffers.mamba_track_indices.dtype, torch.int64) + self.assertEqual(buffers.mamba_track_indices.device.type, "cuda") + torch.testing.assert_close( + buffers.mamba_track_indices, + torch.zeros(4, dtype=torch.int64, device="cuda"), + ) + buffers.mamba_track_indices.copy_(slots[-4:]) + else: + self.assertIsNone(buffers.mamba_track_indices) + cls = ( + multi.MultiLayerEagleDraftExtendCudaGraphRunner + if multilayer + else eagle.EAGLEDraftExtendCudaGraphRunner + ) + runner = cls.__new__(cls) + runner.buffers = buffers + runner.captured_req_width = 4 + runner.forward_mode = ForwardMode.DRAFT_EXTEND_V2 + runner.extend_seq_lens_cpu = [4] * 4 + runner.require_mlp_tp_gather = False + runner.require_attn_tp_gather = False + runner.require_gathered_buffer = False + runner.capture_bs = [4] + runner.seq_len_fill_value = 4 + runner.num_front_tokens = 0 + runner.prune_draft_extend_logits = False + runner.metadata_captured_in_graph = False + runner.step = 0 + runner.deepep_adapter = Mock() + runner.device_module = torch.cuda + runner.attn_backend = SimpleNamespace( + supports_draft_extend_metadata_staging=False + ) + runner.draft_extend_attn_backend = backend + runner.eagle_worker = SimpleNamespace(draft_extend_attn_backend_list=[backend]) + runner.model_runner = SimpleNamespace( + model_config=SimpleNamespace(model_is_mrope=False), + spec_algorithm=SimpleNamespace(is_standalone=lambda: False), + device_timer=None, + canary_manager=None, + ) + return runner, backend, allocator, pool, slots, buffers + + def test_static_staged_runner_checkpoint_replay(self): + for tracking in (False, True): + with self.subTest(tracking=tracking): + self._check_draft_checkpoint_replay(True, True, tracking, staging=True) + + def _check_draft_checkpoint_replay( + self, multilayer, static, tracking, staging=False + ): + from sglang.srt.speculative import eagle_draft_extend_cuda_graph_runner as eagle + from sglang.srt.speculative import ( + multi_layer_eagle_draft_extend_cuda_graph_runner as multi, + ) + from sglang.srt.speculative.eagle_info import EagleDraftExtendInput + + runner, backend, allocator, pool, slots, buffers = ( + self._make_draft_checkpoint_runner(multilayer, static, tracking) + ) + hook_backend = backend + if staging: + wrapper = InklingShortConvHybridAttnBackend.__new__( + InklingShortConvHybridAttnBackend + ) + wrapper.short_conv_backend = backend + wrapper.full_attn_backend = Mock( + supports_draft_extend_metadata_staging=True + ) + wrapper.full_attn_backend.draft_extend_metadata_captured_in_graph.return_value = False + hook_backend = wrapper + runner.attn_backend = wrapper + runner.draft_extend_attn_backend = wrapper + runner.eagle_worker.draft_extend_attn_backend_list = [wrapper] + cache = pool.mamba_cache.conv[0][0] + captured = [] + graph = torch.cuda.CUDAGraph() + output = SimpleNamespace( + next_token_logits=buffers.next_token_logits_buffer, + hidden_states=buffers.hidden_states, + ) + + def model_forward(input_ids, positions, batch): + if not captured: + captured.append(batch) + ShortConvolution._update_sconv_cache_for_draft_extend( + SimpleNamespace(_conv_state=lambda _: backend.sconv_metadata), + batch, + cache, + backend._cache_indices, + buffers.hidden_states, + ) + return output + + def capture_graph(key, run_once, **kwargs): + run_once() + with torch.cuda.graph(graph): + run_once() + + runner.model_runner.model = SimpleNamespace(forward=model_forward) + runner.backend = SimpleNamespace(capture_one=capture_graph) + runner._replay_graph = lambda *args: (graph.replay(), output)[1] + sconv_module = "sglang.srt.models.inkling_common.sconv" + with ( + patch( + sconv_module + ".get_exec", + return_value=SimpleNamespace( + mamba=SimpleNamespace( + mamba_track_interval=4, enable_mamba_extra_buffer=True + ) + ), + ), + patch.object(eagle, "maybe_flashinfer_autotune_speculative_draft"), + patch.object(eagle, "set_dp_buffer_len"), + patch.object(eagle, "set_is_extend_in_batch"), + ): + if multilayer: + batch = runner.get_forward_batch(4) + hook_backend.init_forward_metadata_out_graph(batch, in_capture=True) + + def run_once(): + hook_backend.init_forward_metadata_in_graph(batch) + return model_forward(batch.input_ids, batch.positions, batch) + + capture_graph(None, run_once) + else: + runner.capture_one_shape(4, None) + pointer = ( + backend.sconv_metadata.track_cache_indices.data_ptr() + if tracking + else None + ) + self.assertIsNone(backend.sconv_metadata.track_conv_indices) + for live in (4, 2, 1, 3, 4): + if live == 2 and not static: + allocator.free(slots[:1].clone()) + allocator._flush(urgent=True) + ids = slots[-live:].flip(0).clone() + spec = EagleDraftExtendInput( + hidden_states=buffers.hidden_states[: live * 4], + num_correct_drafts=torch.ones( + live, dtype=torch.int32, device="cuda" + ), + num_accept_tokens=torch.full( + (live,), 3, dtype=torch.int32, device="cuda" + ), + ) + fresh = SimpleNamespace( + batch_size=live, + input_ids=buffers.input_ids[: live * 4], + positions=buffers.positions[: live * 4], + out_cache_loc=buffers.out_cache_loc[: live * 4], + req_pool_indices=torch.arange(1, live + 1, device="cuda"), + seq_lens=torch.full((live,), 6, dtype=torch.int64, device="cuda"), + seq_lens_sum=6 * live, + seq_lens_cpu=None, + extend_seq_lens=torch.full( + (live,), 4, dtype=torch.int32, device="cuda" + ), + extend_seq_lens_cpu=None, + mamba_track_indices=ids if tracking and live != 3 else None, + spec_info=spec, + ) + cache.fill_(-7) + cache[0].zero_() + buffers.hidden_states.copy_( + torch.arange(16, device="cuda")[:, None] + 1 + ) + original = ids.clone() + active = backend._translate_mamba_indices(slots[1 : live + 1]).long() + physical = backend._translate_mamba_indices(ids).long() + expected = cache.clone() + for row in range(live): + joined = torch.cat( + [ + cache[active[row]], + buffers.hidden_states[row * 4 : (row + 1) * 4], + ] + ) + if tracking and live != 3: + expected[physical[row]] = joined[1:4] + expected[active[row]] = joined[3:6] + if multilayer: + composite = multi.MultiLayerEagleMultiStepDraftExtendCudaGraphRunner.__new__( + multi.MultiLayerEagleMultiStepDraftExtendCudaGraphRunner + ) + composite.buffers = buffers + composite.captured_req_width = 4 + composite.require_mlp_tp_gather = False + composite.require_gathered_buffer = False + composite.capture_bs = [4] + composite.num_front_tokens = 0 + composite.seq_len_fill_value = 4 + composite.runners = [runner] + composite.draft_extend_attn_backend_list = [hook_backend] + if not staging: + composite._stage_metadata = Mock() + composite.prepare(fresh) + runner.replay( + 4, composite.seq_lens_sum, composite._replay_spec_info, None + ) + else: + runner.execute(fresh, torch.arange(live, device="cuda") * 4 + 1) + torch.testing.assert_close(cache[1:], expected[1:], rtol=0, atol=0) + torch.testing.assert_close(ids, original, rtol=0, atol=0) + if tracking: + expected_ids = torch.zeros_like(buffers.mamba_track_indices) + if live != 3: + expected_ids[:live] = original + torch.testing.assert_close( + buffers.mamba_track_indices, expected_ids, rtol=0, atol=0 + ) + self.assertEqual( + backend.sconv_metadata.track_cache_indices.data_ptr(), pointer + ) + else: + self.assertIsNone(captured[0].mamba_track_indices) + + +if __name__ == "__main__": + unittest.main()