diff --git a/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py b/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py new file mode 100644 index 000000000..553ec66f6 --- /dev/null +++ b/python/sglang/kernels/ops/mamba/mamba_state_indices_triton.py @@ -0,0 +1,75 @@ +"""Fused replay-prep state-indices kernel for the mamba cuda-graph path. + +``MambaAttnBackendBase._replay_metadata`` refreshes the captured per-bs +``state_indices_list`` buffer before every cuda-graph replay. The reference +form is a chain of dispatched aten ops whose host cost shows up in the bs=1 +MTP inter-phase seam: + + req_pool_indices[valid_bs:] = 0 # zero padded rows (side effect) + mamba_indices = mapping[req_pool_indices] # get_mamba_indices gather + mamba_indices = translate(mamba_indices) # identity for the static pool + mamba_indices[valid_bs:] = -1 # padding sentinel + state_indices[:total_bs].copy_(mamba_indices) + +This module fuses that chain into a single launch. +""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _fused_replay_state_indices_kernel( + req_pool_indices_ptr, # (total_bs,) int64 — static replay buffer + mamba_map_ptr, # (req_pool_size,) int32 — req_index_to_mamba_index_mapping + out_ptr, # (total_bs,) int32 — state_indices_list[bs - 1] + valid_bs, + total_bs, + BS_UPPER: tl.constexpr, +): + offs = tl.arange(0, BS_UPPER) + in_range = offs < total_bs + valid = offs < valid_bs + req = tl.load(req_pool_indices_ptr + offs, mask=valid, other=0) + idx = tl.load(mamba_map_ptr + req, mask=valid, other=0) + out_val = tl.where(valid, idx.to(tl.int32), -1) + tl.store(out_ptr + offs, out_val, mask=in_range) + # Preserve the reference chain's side effect: padded rows of the static + # req_pool_indices buffer are zeroed so captured kernels that gather + # with them stay in-bounds. + zeros = tl.zeros([BS_UPPER], dtype=req.dtype) + tl.store(req_pool_indices_ptr + offs, zeros, mask=in_range & (~valid)) + + +def fused_replay_state_indices( + *, + req_pool_indices: torch.Tensor, + mamba_index_mapping: torch.Tensor, + out_state_indices: torch.Tensor, + valid_bs: int, + total_bs: int, +) -> torch.Tensor: + """Fill the captured replay state-indices buffer in one launch. + + Mapping gather + padding sentinel + store into ``out_state_indices``, plus + the reference chain's side effect of zeroing the padded rows of + ``req_pool_indices``. Rows ``[valid_bs, total_bs)`` are padding: they get + the ``-1`` sentinel (mamba kernels skip ``state_idx < 0``) and their + ``req_pool_indices`` entries are zeroed. + + Callers must supply an identity virtual->physical mapping (the static + hybrid pool); the unified pool's allocator translate is not a flat table + gather and has to take the reference chain. + + Returns the filled ``out_state_indices[:total_bs]`` view. + """ + _fused_replay_state_indices_kernel[(1,)]( + req_pool_indices, + mamba_index_mapping, + out_state_indices, + valid_bs, + total_bs, + BS_UPPER=triton.next_power_of_2(total_bs), + ) + return out_state_indices[:total_bs] diff --git a/python/sglang/kernels/ops/speculative/cache_locs.py b/python/sglang/kernels/ops/speculative/cache_locs.py index 01ee6eae9..ec1953d66 100644 --- a/python/sglang/kernels/ops/speculative/cache_locs.py +++ b/python/sglang/kernels/ops/speculative/cache_locs.py @@ -361,6 +361,74 @@ def assign_extend_cache_locs( save_offset += BLOCK_SIZE +@triton.jit +def assign_extend_cache_locs_uniform( + req_pool_indices, + req_to_token, + start_offset, + out_cache_loc, + pool_len: tl.constexpr, + draft_token_num: tl.constexpr, +): + """Uniform-length variant of assign_extend_cache_locs: every row extends + exactly draft_token_num tokens, so the end offset is start + + draft_token_num (computed here, no end_offset tensor) and the output + offset is pid * draft_token_num (no cross-row prefix-sum loads).""" + BLOCK_SIZE: tl.constexpr = 64 + pid = tl.program_id(axis=0) + kv_start = tl.load(start_offset + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + out_cache_ptr = out_cache_loc + pid * draft_token_num + + offs = tl.arange(0, BLOCK_SIZE) + num_loop = tl.cdiv(draft_token_num, BLOCK_SIZE) + for i in range(num_loop): + o = offs + i * BLOCK_SIZE + mask = o < draft_token_num + data = tl.load(token_pool + kv_start + o, mask=mask) + tl.store(out_cache_ptr + o, data, mask=mask) + + +def assign_extend_cache_locs_uniform_func( + req_pool_indices: torch.Tensor, + req_to_token: torch.Tensor, + start_offset: torch.Tensor, + batch_size: int, + draft_token_num: int, + device, +) -> torch.Tensor: + """assign_extend_cache_locs for the uniform case (all rows extend exactly + draft_token_num tokens, e.g. spec target-verify prep). Computes end + offsets inside the kernel, removing the eager `seq_lens + draft_token_num` + add from the host critical path.""" + if _is_cuda or _is_hip or _is_musa or _is_xpu: + out_cache_loc = torch.empty( + (batch_size * draft_token_num,), + dtype=torch.int64, + device=device, + ) + assign_extend_cache_locs_uniform[(batch_size,)]( + req_pool_indices, + req_to_token, + start_offset, + out_cache_loc, + req_to_token.shape[1], + draft_token_num, + ) + return out_cache_loc + + # NPU / CPU platforms: fall back to the end_offset-tensor path. + return assign_extend_cache_locs_func( + req_pool_indices=req_pool_indices, + req_to_token=req_to_token, + start_offset=start_offset, + end_offset=start_offset + draft_token_num, + batch_size=batch_size, + draft_token_num=draft_token_num, + device=device, + ) + + def assign_extend_cache_locs_func( req_pool_indices: torch.Tensor, req_to_token: torch.Tensor, diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 6e9e57adf..61a45f42e 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -4,6 +4,9 @@ from typing import Optional, Union import torch from sglang.kernels.ops.mamba.causal_conv1d_triton import PAD_SLOT_ID +from sglang.kernels.ops.mamba.mamba_state_indices_triton import ( + fused_replay_state_indices, +) from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( scatter_mamba_states_after_mtp_verify, track_mamba_states_if_needed, @@ -36,6 +39,16 @@ class MambaAttnBackendBase(AttentionBackend): self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool self.token_to_kv_pool = model_runner.token_to_kv_pool self.enable_unified_memory = model_runner.server_args.enable_unified_memory + # Fused replay-prep state-indices fast path (fused_replay_state_indices): + # requires the static hybrid pool whose v2p translate is the identity — + # the unified pool overrides translate_mamba_indices with an allocator + # lookup that is not a flat table gather. + self._fused_state_indices_ok = ( + str(self.device).startswith("cuda") + and isinstance(self.req_to_token_pool, HybridReqToTokenPool) + and type(self.req_to_token_pool).translate_mamba_indices + is HybridReqToTokenPool.translate_mamba_indices + ) self.forward_metadata: ForwardMetadata = None self.state_indices_list = [] # Static (max_bs,) track-dest buffer captured by pointer, refreshed in-place @@ -242,6 +255,30 @@ class MambaAttnBackendBase(AttentionBackend): def init_forward_metadata(self, forward_batch: ForwardBatch): self.forward_metadata = self._forward_metadata(forward_batch) + def update_verify_buffers_to_fill_after_draft( + self, spec_info: SpecInput, cuda_graph_bs: Optional[int] + ): + # Plan-stream fixup: slot indices / static query_start_loc are + # draft-independent, but tree verify (topk > 1) copies the + # draft-produced tree links into the captured buffers on the plan + # stream, racing the draft — re-copy after the stream join. Eager + # verify reads the spec_info tensors directly; parent links are + # derived from these buffers at execution time. + if self.topk <= 1 or cuda_graph_bs is None: + return + if ( + not isinstance(spec_info, EagleVerifyInput) + or spec_info.retrieve_next_token is None # dummy / capture runs + ): + return + bs_without_pad = spec_info.retrieve_next_token.shape[0] + self.retrieve_next_token_list[cuda_graph_bs - 1][:bs_without_pad].copy_( + spec_info.retrieve_next_token + ) + self.retrieve_next_sibling_list[cuda_graph_bs - 1][:bs_without_pad].copy_( + spec_info.retrieve_next_sibling + ) + def _init_track_conv_indices( self, query_start_loc: torch.Tensor, forward_batch: ForwardBatch ): @@ -507,14 +544,26 @@ class MambaAttnBackendBase(AttentionBackend): num_padding = torch.count_nonzero( seq_lens_cpu == self.get_cuda_graph_seq_len_fill_value() ) - # Make sure forward metadata is correctly handled for padding reqs - req_pool_indices[bs - num_padding :] = 0 - mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) - # Translate using the LIVE v2p table BEFORE the padding sentinel below; - # captured Mamba kernels read state_indices_list as PHYSICAL ids. - mamba_indices = self._translate_mamba_indices(mamba_indices) - mamba_indices[bs - num_padding :] = -1 - self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) + if self._fused_state_indices_ok and self.replayssm_write_pos_list is None: + # Single-launch fast path: mapping gather + padding sentinel + store + # into the static buffer, plus zeroing padded req_pool_indices rows — + # bit-identical to the reference chain below. + mamba_indices = fused_replay_state_indices( + req_pool_indices=req_pool_indices, + mamba_index_mapping=self.req_to_token_pool.req_index_to_mamba_index_mapping, + out_state_indices=self.state_indices_list[bs - 1], + valid_bs=bs - int(num_padding), + total_bs=bs, + ) + else: + # Make sure forward metadata is correctly handled for padding reqs + req_pool_indices[bs - num_padding :] = 0 + mamba_indices = self.req_to_token_pool.get_mamba_indices(req_pool_indices) + # Translate using the LIVE v2p table BEFORE the padding sentinel below; + # captured Mamba kernels read state_indices_list as PHYSICAL ids. + mamba_indices = self._translate_mamba_indices(mamba_indices) + mamba_indices[bs - num_padding :] = -1 + self.state_indices_list[bs - 1][: len(mamba_indices)].copy_(mamba_indices) # Refresh the static track-dest buffer in-place (translated); the captured # track-save reads it, leaving the handed-in InputBuffer slot read-only. track_buf = None @@ -866,6 +915,24 @@ class HybridLinearAttnBackend(AttentionBackend): for attn_backend in self.attn_backend_list: attn_backend.on_after_cuda_graph_warmup() + def get_verify_buffers_to_fill_after_draft(self): + # Verify tree-mask / position buffers live on the full-attn child (the + # linear side consumes no mask). Handing them out lets the draft stage + # write straight into the captured verify buffers instead of allocating + # a fresh mask every step. + return self.full_attn_backend.get_verify_buffers_to_fill_after_draft() + + def update_verify_buffers_to_fill_after_draft( + self, spec_info: SpecInput, cuda_graph_bs: Optional[int] + ): + # Plan-stream fixup after draft completes: forward to both children. + # Sub-backends that cannot run under the plan stream keep the fail-loud + # NotImplementedError base behavior. + for attn_backend in self.attn_backend_list: + attn_backend.update_verify_buffers_to_fill_after_draft( + spec_info=spec_info, cuda_graph_bs=cuda_graph_bs + ) + def init_forward_metadata(self, forward_batch: ForwardBatch): if forward_batch.forward_mode.is_draft_extend_v2(): # DRAFT_EXTEND_V2 runs only full-attn layers in the draft model; skip diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 9c91d0584..86bc22d69 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -252,6 +252,13 @@ class KDAKernelDispatcher: class KDAAttnBackend(MambaAttnBackendBase): """Attention backend for KDA (Kimi Delta Attention) linear attention.""" + # Same GPU-only contract as GDNAttnBackend / Mamba2AttnBackend: KDA metadata + # never reads the spec-v2 seq_lens_cpu mirror (replay padding comes from + # forward_batch.num_padding, and the replayssm track-flush mask paths are + # gated `not is_kda`), so don't force FutureMap's blocking per-step + # seq_lens D2H (~0.5 ms/step host stall in bs=1 MTP decode). + needs_cpu_seq_lens: bool = False + def __init__(self, model_runner: ModelRunner): super().__init__(model_runner) # mamba_cache.conv is [..., kernel-1, dim] while conv_states_shape expects the window length (kernel-1) at shape[-1], hence the transpose. diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 780261541..ff09df9d2 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/python/sglang/srt/speculative/eagle_utils.py @@ -491,7 +491,7 @@ def eagle_prepare_for_verify( target_worker: TpModelWorker, ): from sglang.kernels.ops.speculative.cache_locs import ( - assign_extend_cache_locs_func, + assign_extend_cache_locs_uniform_func, ) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -511,11 +511,13 @@ def eagle_prepare_for_verify( "v2 prepare_for_verify input_ids", ) device = batch.device - batch.out_cache_loc = assign_extend_cache_locs_func( + # Uniform variant: end offsets (= start + draft_token_num) are computed + # inside the kernel, keeping the eager `seq_lens + N` add off the host + # critical path (bs=1 MTP inter-phase seam). + batch.out_cache_loc = assign_extend_cache_locs_uniform_func( req_pool_indices=batch.req_pool_indices, req_to_token=req_to_token_pool.req_to_token, start_offset=batch.seq_lens, - end_offset=batch.seq_lens + verify_input.draft_token_num, batch_size=bs, draft_token_num=verify_input.draft_token_num, device=device, diff --git a/test/registered/attention/unittests/hybrid_linear/test_verify_buffer_fixup_hook.py b/test/registered/attention/unittests/hybrid_linear/test_verify_buffer_fixup_hook.py new file mode 100644 index 000000000..7f4472f04 --- /dev/null +++ b/test/registered/attention/unittests/hybrid_linear/test_verify_buffer_fixup_hook.py @@ -0,0 +1,131 @@ +"""The mamba plan-stream verify fixup hook must refresh draft-produced tree links. + +Under ``SGLANG_ENABLE_OVERLAP_PLAN_STREAM``, ``_replay_metadata`` copies +``spec_info.retrieve_next_token`` / ``retrieve_next_sibling`` into the captured +per-bs buffers on the plan stream, racing the draft's ``build_tree`` on the +compute stream. ``update_verify_buffers_to_fill_after_draft`` is the post-join +fixup that re-copies them into the ``cuda_graph_bs`` buffers; chain mode +(``topk == 1``), the eager path (``cuda_graph_bs is None``), and dummy/capture +runs (``retrieve_next_token is None``) must stay no-ops. + +Pure host-side buffer logic — CPU tensors, no CUDA required. +""" + +import unittest + +import torch + +from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( + MambaAttnBackendBase, +) +from sglang.srt.speculative.eagle_info import EagleVerifyInput +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=2, stage="base-b", runner_config="1-gpu-large") + +_STALE = -555 +_MAX_BS = 4 +_DRAFT_TOKEN_NUM = 8 + + +def _make_backend(topk: int) -> MambaAttnBackendBase: + """Bare backend with only the fields the hook reads (no ModelRunner).""" + backend = object.__new__(MambaAttnBackendBase) + backend.topk = topk + # Mirror init_cuda_graph_state's per-bs buffer shapes: (bs, draft_token_num). + backend.retrieve_next_token_list = [ + torch.full((bs, _DRAFT_TOKEN_NUM), _STALE, dtype=torch.int32) + for bs in range(1, _MAX_BS + 1) + ] + backend.retrieve_next_sibling_list = [ + torch.full((bs, _DRAFT_TOKEN_NUM), _STALE, dtype=torch.int32) + for bs in range(1, _MAX_BS + 1) + ] + return backend + + +def _make_verify_input(bs_without_pad: int, base: int) -> EagleVerifyInput: + """EagleVerifyInput carrying just the tree-link tensors the hook consumes.""" + spec_info = object.__new__(EagleVerifyInput) + numel = bs_without_pad * _DRAFT_TOKEN_NUM + spec_info.retrieve_next_token = ( + torch.arange(base, base + numel, dtype=torch.int32) + ).reshape(bs_without_pad, _DRAFT_TOKEN_NUM) + spec_info.retrieve_next_sibling = ( + torch.arange(base + numel, base + 2 * numel, dtype=torch.int32) + ).reshape(bs_without_pad, _DRAFT_TOKEN_NUM) + return spec_info + + +class TestVerifyBufferFixupHook(CustomTestCase): + def test_refresh_overwrites_stale_links(self): + backend = _make_backend(topk=2) + cuda_graph_bs = _MAX_BS + bs_without_pad = _MAX_BS - 1 # one padded row + spec_info = _make_verify_input(bs_without_pad=bs_without_pad, base=100) + + backend.update_verify_buffers_to_fill_after_draft( + spec_info=spec_info, cuda_graph_bs=cuda_graph_bs + ) + + for buf_list, fresh in ( + (backend.retrieve_next_token_list, spec_info.retrieve_next_token), + (backend.retrieve_next_sibling_list, spec_info.retrieve_next_sibling), + ): + buf = buf_list[cuda_graph_bs - 1] + self.assertTrue( + torch.equal(buf[:bs_without_pad], fresh), + "fresh tree links not copied into the captured buffer", + ) + # The padded tail row is not covered by the copy. + self.assertTrue( + (buf[bs_without_pad:] == _STALE).all(), + "rows beyond bs_without_pad must not be written", + ) + # Buffers of other captured batch sizes stay untouched. + for other_bs in range(1, _MAX_BS): + self.assertTrue( + (buf_list[other_bs - 1] == _STALE).all(), + f"buffer for bs={other_bs} must stay untouched", + ) + + def test_chain_topk1_is_noop(self): + backend = _make_backend(topk=1) + spec_info = _make_verify_input(bs_without_pad=2, base=100) + backend.update_verify_buffers_to_fill_after_draft( + spec_info=spec_info, cuda_graph_bs=2 + ) + self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all()) + self.assertTrue((backend.retrieve_next_sibling_list[1] == _STALE).all()) + + def test_eager_path_is_noop(self): + backend = _make_backend(topk=2) + spec_info = _make_verify_input(bs_without_pad=2, base=100) + backend.update_verify_buffers_to_fill_after_draft( + spec_info=spec_info, cuda_graph_bs=None + ) + self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all()) + self.assertTrue((backend.retrieve_next_sibling_list[1] == _STALE).all()) + + def test_dummy_run_none_links_is_noop(self): + backend = _make_backend(topk=2) + spec_info = object.__new__(EagleVerifyInput) + spec_info.retrieve_next_token = None # dummy / capture run + spec_info.retrieve_next_sibling = None + backend.update_verify_buffers_to_fill_after_draft( + spec_info=spec_info, cuda_graph_bs=2 + ) + self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all()) + self.assertTrue((backend.retrieve_next_sibling_list[1] == _STALE).all()) + + def test_non_eagle_spec_input_is_noop(self): + backend = _make_backend(topk=2) + backend.update_verify_buffers_to_fill_after_draft( + spec_info=None, cuda_graph_bs=2 + ) + self.assertTrue((backend.retrieve_next_token_list[1] == _STALE).all()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py b/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py new file mode 100644 index 000000000..93d059283 --- /dev/null +++ b/test/registered/kernels/ops/mamba/test_fused_replay_state_indices.py @@ -0,0 +1,192 @@ +"""fused_replay_state_indices must be bit-identical to the unfused prep. + +The unfused reference is the exact op sequence ``_replay_metadata`` used to +launch for the static hybrid pool: + + req_pool_indices[valid_bs:total_bs] = 0 # zero padded rows (side effect) + mamba_indices = mapping[req_pool_indices] # get_mamba_indices gather + # identity v2p translate (static pool) + mamba_indices[valid_bs:] = -1 # padding sentinel + state_indices[:total_bs].copy_(mamba_indices) + +The two paths must agree bit-for-bit, INCLUDING the side effect of zeroing the +padded rows of the static ``req_pool_indices`` replay buffer — captured kernels +gather with that buffer, so a non-zeroed padded row is a delayed illegal memory +access, not a visible diff. Both paths run on guard-padded buffers across a +bs x num_padding matrix (non-power-of-two sizes exercise the BS_UPPER masking): + +1. the produced state indices are identical over the whole ``[0, total_bs)`` + range (padding sentinel rows included); +2. the ``req_pool_indices`` buffer ends up identical (padded rows zeroed); +3. neither buffer is written beyond ``total_bs`` (guard tails stay intact). +""" + +import unittest + +import torch + +from sglang.kernels.ops.mamba.mamba_state_indices_triton import ( + fused_replay_state_indices, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=5, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +# Guard tail appended to every buffer; must stay untouched by both paths. +_GUARD = 8 +_GUARD_SENTINEL = -7777 +# Poison for the out buffer so unwritten cells inside [0, total_bs) are caught. +_OUT_POISON = -12345 + +_REQ_POOL_SIZE = 160 +_MAMBA_POOL_SIZE = 4096 + + +def _reference_chain( + req_pool_indices: torch.Tensor, + mapping: torch.Tensor, + out_buf: torch.Tensor, + valid_bs: int, + total_bs: int, +) -> None: + """Replicates the _replay_metadata reference ops, in order, in place.""" + req_pool_indices[valid_bs:total_bs] = 0 + mamba_indices = mapping[req_pool_indices[:total_bs]] + # static pool: _translate_mamba_indices is the identity + mamba_indices[valid_bs:] = -1 + out_buf[: len(mamba_indices)].copy_(mamba_indices) + + +@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA (triton kernel)") +class TestFusedReplayStateIndices(CustomTestCase): + def _run_case(self, total_bs: int, num_padding: int, seed: int) -> None: + device = torch.device("cuda") + gen = torch.Generator(device="cpu").manual_seed(seed) + valid_bs = total_bs - num_padding + + # Production dtypes: req_pool_indices int64 (static replay buffer), + # req_index_to_mamba_index_mapping int32, state_indices_list int32. + req_pool = torch.randint( + 0, _REQ_POOL_SIZE, (total_bs + _GUARD,), generator=gen, dtype=torch.int64 + ) + req_pool[total_bs:] = _GUARD_SENTINEL + mapping = torch.randint( + 0, _MAMBA_POOL_SIZE, (_REQ_POOL_SIZE,), generator=gen, dtype=torch.int32 + ) + out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32) + + req_pool_ref = req_pool.to(device) + req_pool_fused = req_pool.to(device) + mapping_dev = mapping.to(device) + out_ref = out.to(device) + out_fused = out.to(device) + + _reference_chain( + req_pool_indices=req_pool_ref, + mapping=mapping_dev, + out_buf=out_ref, + valid_bs=valid_bs, + total_bs=total_bs, + ) + returned = fused_replay_state_indices( + req_pool_indices=req_pool_fused, + mamba_index_mapping=mapping_dev, + out_state_indices=out_fused, + valid_bs=valid_bs, + total_bs=total_bs, + ) + torch.cuda.synchronize() + + case = f"{total_bs=} {num_padding=} {seed=}" + # 1. state indices bit-identical over [0, total_bs), sentinels included + self.assertTrue( + torch.equal(out_ref[:total_bs], out_fused[:total_bs]), + f"state indices mismatch ({case}):\n" + f" ref {out_ref[:total_bs].tolist()}\n" + f" fused {out_fused[:total_bs].tolist()}", + ) + # The returned view is what _replay_metadata forwards downstream. + self.assertTrue( + torch.equal(returned, out_fused[:total_bs]), + f"returned view is not the filled buffer ({case})", + ) + # 2. req_pool_indices side effect bit-identical (padded rows zeroed) + self.assertTrue( + torch.equal(req_pool_ref[:total_bs], req_pool_fused[:total_bs]), + f"req_pool_indices mismatch ({case}):\n" + f" ref {req_pool_ref[:total_bs].tolist()}\n" + f" fused {req_pool_fused[:total_bs].tolist()}", + ) + # Explicit re-statement of the contract, independent of the reference: + self.assertTrue( + (req_pool_fused[valid_bs:total_bs] == 0).all(), + f"padded req_pool_indices rows not zeroed ({case})", + ) + self.assertTrue( + (out_fused[valid_bs:total_bs] == -1).all(), + f"padding sentinel rows not -1 ({case})", + ) + self.assertFalse( + (out_fused[:total_bs] == _OUT_POISON).any(), + f"unwritten cells inside [0, total_bs) ({case})", + ) + # 3. no out-of-range writes past total_bs (BS_UPPER > total_bs masking) + for name, buf in (("req_pool", req_pool_fused), ("out", out_fused)): + expected = _GUARD_SENTINEL if name == "req_pool" else _OUT_POISON + self.assertTrue( + (buf[total_bs:] == expected).all(), + f"{name} guard tail clobbered ({case}): {buf[total_bs:].tolist()}", + ) + + def test_matrix(self): + # Non-power-of-two sizes (7, 33) exercise the BS_UPPER in_range mask; + # num_padding sweeps none / one / half / all-but-one padded rows. + for total_bs in (1, 2, 7, 32, 33): + paddings = sorted( + {0, 1, total_bs // 2, total_bs - 1} & set(range(total_bs)) + ) + for num_padding in paddings: + for seed in (0, 1, 2): + with self.subTest( + total_bs=total_bs, num_padding=num_padding, seed=seed + ): + self._run_case( + total_bs=total_bs, num_padding=num_padding, seed=seed + ) + + def test_shared_mamba_slots(self): + # Multiple requests mapping to the same mamba slot (mapping is not + # injective in general) must gather identically on both paths. + mapping_const = torch.full((_REQ_POOL_SIZE,), 3, dtype=torch.int32) + device = torch.device("cuda") + total_bs, num_padding = 7, 2 + valid_bs = total_bs - num_padding + req_pool = torch.arange(total_bs + _GUARD, dtype=torch.int64) + out = torch.full((total_bs + _GUARD,), _OUT_POISON, dtype=torch.int32) + + req_ref, req_fused = req_pool.to(device), req_pool.to(device) + out_ref, out_fused = out.to(device), out.to(device) + mapping_dev = mapping_const.to(device) + + _reference_chain( + req_pool_indices=req_ref, + mapping=mapping_dev, + out_buf=out_ref, + valid_bs=valid_bs, + total_bs=total_bs, + ) + fused_replay_state_indices( + req_pool_indices=req_fused, + mamba_index_mapping=mapping_dev, + out_state_indices=out_fused, + valid_bs=valid_bs, + total_bs=total_bs, + ) + torch.cuda.synchronize() + self.assertTrue(torch.equal(out_ref, out_fused)) + self.assertTrue(torch.equal(req_ref, req_fused)) + + +if __name__ == "__main__": + unittest.main()