diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py index e789a23236ea..c5903f9bbb42 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.py @@ -38,6 +38,7 @@ from .._compat import KvCacheConfig, nvtx_range, prefer_pinned, str_dtype_to_torch from ..utils.logger import ad_logger from ..utils.node_utils import extract_op_args, get_op_schema +from .mamba.replay_metadata import REPLAY_WORK_ITEM_WIDTH Constant = Union[int, float, str, None] @@ -2204,10 +2205,9 @@ def from_base( class ReplayOldXHandler(SpeculativeOnly, StateResourceHandler): - """Per-layer old_x cache for the replay SSM kernel (single-buffered, bf16). + """Per-layer old_x cache for the replay SSM kernel (double-buffered, bf16). - Shape: (max_batch, T, num_heads, head_dim) — T is determined by the manager's - spec_config (max_draft_len + 1), not by this handler. Acts as a type marker. + Shape: (max_batch, 2, replay_history_size, num_heads, head_dim). Routes to MambaHybridCacheManager via get_replay_old_x(layer_idx). """ @@ -2235,7 +2235,7 @@ def __eq__(self, other) -> bool: class ReplayOldBHandler(SpeculativeOnly, StateResourceHandler): """Per-layer old_B cache for the replay SSM kernel (double-buffered, bf16). - Shape: (max_batch, 2, T, n_groups, d_state) — T from manager. + Shape: (max_batch, 2, replay_history_size, n_groups, d_state). Routes to MambaHybridCacheManager via get_replay_old_B(layer_idx). """ @@ -2263,7 +2263,7 @@ def __eq__(self, other) -> bool: class ReplayOldDtHandler(SpeculativeOnly, StateResourceHandler): """Per-layer old_dt cache for the replay SSM kernel (double-buffered, fp32). - Shape: (max_batch, 2, num_heads, T) — T from manager. + Shape: (max_batch, 2, num_heads, replay_history_size). Routes to MambaHybridCacheManager via get_replay_old_dt(layer_idx). """ @@ -2285,7 +2285,7 @@ def __eq__(self, other) -> bool: class ReplayOldDAcumsumHandler(SpeculativeOnly, StateResourceHandler): """Per-layer old_dA_cumsum cache for the replay SSM kernel (double-buffered, fp32). - Shape: (max_batch, 2, num_heads, T) — T from manager. + Shape: (max_batch, 2, num_heads, replay_history_size). Routes to MambaHybridCacheManager via get_replay_old_dA_cumsum(layer_idx). """ @@ -2346,6 +2346,35 @@ def __eq__(self, other) -> bool: return isinstance(other, ReplayPrevNumAcceptedHandler) +class ReplayWorkItemsHandler(ResourceHandler): + """Shared per-forward replay work items for the checkpoint replay SSM kernel. + + Shape: (max_batch, REPLAY_WORK_ITEM_WIDTH) int32. Each row is + (position_in_decode_batch, cache_slot, prev_num_accepted_tokens, cache_buf_idx). + """ + + def allocate(self, sequence_info) -> torch.Tensor: + return torch.empty( + sequence_info.max_num_state_slots, + REPLAY_WORK_ITEM_WIDTH, + device=sequence_info.device, + dtype=torch.int32, + ) + + def __eq__(self, other) -> bool: + return isinstance(other, ReplayWorkItemsHandler) + + +class ReplayNWritesHandler(ResourceHandler): + """Shared single-element device tensor holding the replay write-count.""" + + def allocate(self, sequence_info) -> torch.Tensor: + return torch.empty(1, device=sequence_info.device, dtype=torch.int32) + + def __eq__(self, other) -> bool: + return isinstance(other, ReplayNWritesHandler) + + class IntermediateConvStateHandler(SpeculativeOnly, StateResourceHandler): """Intermediate conv state cache descriptor for speculative decoding. diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py index 21e57e545ad5..a70996eaabaa 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/flashinfer_backend_mamba.py @@ -13,7 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import List, Optional +from typing import List, Optional, cast import torch from flashinfer.mamba import selective_state_update as _flashinfer_ssm_update @@ -31,11 +31,13 @@ IntermediateSSMStateHandler, MHACallable, ReplayCacheBufIdxHandler, + ReplayNWritesHandler, ReplayOldBHandler, ReplayOldDAcumsumHandler, ReplayOldDtHandler, ReplayOldXHandler, ReplayPrevNumAcceptedHandler, + ReplayWorkItemsHandler, ResourceHandlerDict, ) from .mamba_backend_common import ( @@ -64,7 +66,7 @@ def _fi_align(t: torch.Tensor) -> torch.Tensor: "ssm_state_cache", "intermediate_ssm_state_cache", # replay buffers: written in-place by the precompute and main kernels - # (double-buffered B/dt/dA_cumsum, single-buffered x); None in non-replay mode + # (double-buffered x/B/dt/dA_cumsum); None in non-replay mode "replay_old_x", "replay_old_b", "replay_old_dt", @@ -95,14 +97,22 @@ def _flashinfer_cached_ssm( intermediate_ssm_state_cache: Optional[ torch.Tensor ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode - replay_old_x: Optional[torch.Tensor], # [max_batch, T, nheads, head_dim]; None in non-replay - replay_old_b: Optional[torch.Tensor], # [max_batch, 2, T, ngroups, dstate]; None in non-replay - replay_old_dt: Optional[torch.Tensor], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_old_x: Optional[ + torch.Tensor + ], # [max_batch, 2, history, nheads, head_dim]; None in non-replay + replay_old_b: Optional[ + torch.Tensor + ], # [max_batch, 2, history, ngroups, dstate]; None in non-replay + replay_old_dt: Optional[ + torch.Tensor + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_old_da_cumsum: Optional[ torch.Tensor - ], # [max_batch, 2, nheads, T] fp32; None in non-replay + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_cache_buf_idx: Optional[torch.Tensor], # [max_batch] int32; None in non-replay replay_prev_num_accepted: Optional[torch.Tensor], # [max_batch] int32; None in non-replay + replay_work_items: Optional[torch.Tensor], # [max_batch, 4] int32; None in non-replay + replay_n_writes: Optional[torch.Tensor], # [1] int32; None in non-replay # CONSTANTS time_step_limit: List[float], chunk_size: int, @@ -198,6 +208,34 @@ def _flashinfer_cached_ssm( use_replay = batch_info.is_use_replay() if use_replay: + missing_replay_tensors = [ + name + for name, tensor in ( + ("replay_old_x", replay_old_x), + ("replay_old_b", replay_old_b), + ("replay_old_dt", replay_old_dt), + ("replay_old_da_cumsum", replay_old_da_cumsum), + ("replay_cache_buf_idx", replay_cache_buf_idx), + ("replay_prev_num_accepted", replay_prev_num_accepted), + ("replay_work_items", replay_work_items), + ("replay_n_writes", replay_n_writes), + ) + if tensor is None + ] + if missing_replay_tensors: + raise RuntimeError( + "flashinfer_cached_ssm replay path missing required tensors: " + f"{', '.join(missing_replay_tensors)}" + ) + replay_old_x = cast(torch.Tensor, replay_old_x) + replay_old_b = cast(torch.Tensor, replay_old_b) + replay_old_dt = cast(torch.Tensor, replay_old_dt) + replay_old_da_cumsum = cast(torch.Tensor, replay_old_da_cumsum) + replay_cache_buf_idx = cast(torch.Tensor, replay_cache_buf_idx) + replay_prev_num_accepted = cast(torch.Tensor, replay_prev_num_accepted) + replay_work_items = cast(torch.Tensor, replay_work_items) + replay_n_writes = cast(torch.Tensor, replay_n_writes) + # Replay path: fast-forward SSM state via tl.dot on cached values. # State is updated in-place; no disable_state_update needed. # x_extend/B_extend/C_extend are non-contiguous views from the CUDA graph's @@ -218,6 +256,8 @@ def _flashinfer_cached_ssm( B_extend, C_extend, out=preallocated_ssm_out_e, + n_writes=replay_n_writes, + replay_work_items=replay_work_items[:num_extend], D=D_full, dt_bias=dt_bias_hp, dt_softplus=True, @@ -225,6 +265,11 @@ def _flashinfer_cached_ssm( launch_with_pdl=True, # PDL chain: triton_causal_conv extend → precompute → main ) else: + if intermediate_ssm_state_cache is None: + raise RuntimeError( + "flashinfer_cached_ssm non-replay extend branch requires " + "intermediate_ssm_state_cache" + ) if intermediate_ssm_state_cache.size(1) < tokens_per_extend: raise RuntimeError( "flashinfer_cached_ssm: intermediate_ssm_state_cache is too small " @@ -342,14 +387,22 @@ def _flashinfer_cached_ssm_fake( intermediate_ssm_state_cache: Optional[ torch.Tensor ], # [spec_state_size, max_draft_len+1, num_heads, head_dim, d_state]; None in replay mode - replay_old_x: Optional[torch.Tensor], # [max_batch, T, nheads, head_dim]; None in non-replay - replay_old_b: Optional[torch.Tensor], # [max_batch, 2, T, ngroups, dstate]; None in non-replay - replay_old_dt: Optional[torch.Tensor], # [max_batch, 2, nheads, T] fp32; None in non-replay + replay_old_x: Optional[ + torch.Tensor + ], # [max_batch, 2, history, nheads, head_dim]; None in non-replay + replay_old_b: Optional[ + torch.Tensor + ], # [max_batch, 2, history, ngroups, dstate]; None in non-replay + replay_old_dt: Optional[ + torch.Tensor + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_old_da_cumsum: Optional[ torch.Tensor - ], # [max_batch, 2, nheads, T] fp32; None in non-replay + ], # [max_batch, 2, nheads, history] fp32; None in non-replay replay_cache_buf_idx: Optional[torch.Tensor], # [max_batch] int32; None in non-replay replay_prev_num_accepted: Optional[torch.Tensor], # [max_batch] int32; None in non-replay + replay_work_items: Optional[torch.Tensor], # [max_batch, 4] int32; None in non-replay + replay_n_writes: Optional[torch.Tensor], # [1] int32; None in non-replay # CONSTANTS time_step_limit: List[float], chunk_size: int, @@ -400,7 +453,7 @@ def get_cache_initializers( ssm_h = ret["ssm_state_cache"] - # All 7 optional caches are always registered positionally (None = unused in this mode). + # Optional replay/spec caches are registered positionally (None = unused in this mode). # intermediate_ssm_state_cache: real in non-replay, None in replay. # replay_old_*: real in replay mode (SM80+), None otherwise. if use_replay: @@ -418,6 +471,8 @@ def get_cache_initializers( ret["replay_old_da_cumsum"] = ReplayOldDAcumsumHandler(num_heads=ssm_h.num_heads) ret["replay_cache_buf_idx"] = ReplayCacheBufIdxHandler() ret["replay_prev_num_accepted"] = ReplayPrevNumAcceptedHandler() + ret["replay_work_items"] = ReplayWorkItemsHandler() + ret["replay_n_writes"] = ReplayNWritesHandler() else: ret["intermediate_ssm_state_cache"] = IntermediateSSMStateHandler.from_base(ssm_h) ret["replay_old_x"] = None @@ -426,4 +481,6 @@ def get_cache_initializers( ret["replay_old_da_cumsum"] = None ret["replay_cache_buf_idx"] = None ret["replay_prev_num_accepted"] = None + ret["replay_work_items"] = None + ret["replay_n_writes"] = None return ret diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py new file mode 100644 index 000000000000..ce924b2215af --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/mamba/replay_metadata.py @@ -0,0 +1,22 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Replay metadata layout shared by AutoDeploy Mamba descriptors.""" + +REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 +REPLAY_WORK_CACHE_SLOT = 1 +REPLAY_WORK_PNAT = 2 +REPLAY_WORK_CACHE_BUF_IDX = 3 +REPLAY_WORK_ITEM_WIDTH = 4 diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index b272740deab4..9e3387d26cb3 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py @@ -993,6 +993,7 @@ def _prepare_inputs( _ungathered_new_lens=new_tokens_lens, **extra_args, ) + self.cache_seq_interface.prepare_replay_metadata() self.iter_states["num_ctx_requests"] = num_prefill self.iter_states["num_ctx_tokens"] = num_prefill_tokens diff --git a/tensorrt_llm/_torch/auto_deploy/shim/interface.py b/tensorrt_llm/_torch/auto_deploy/shim/interface.py index 76451e46d70c..30ec0b8ae616 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/interface.py +++ b/tensorrt_llm/_torch/auto_deploy/shim/interface.py @@ -56,17 +56,26 @@ IntermediateSSMStateHandler, KVPagedResourceHandler, ReplayCacheBufIdxHandler, + ReplayNWritesHandler, ReplayOldBHandler, ReplayOldDAcumsumHandler, ReplayOldDtHandler, ReplayOldXHandler, ReplayPrevNumAcceptedHandler, + ReplayWorkItemsHandler, ResourceHandler, ResourceHandlerDict, SequenceInfo, SSMResourceHandler, StateResourceHandler, ) +from ..custom_ops.mamba.replay_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from ..distributed.common import all_gather_object, get_world_size from ..distributed.common import is_initialized as is_distributed_initialized from ..utils.cuda_mem_tracker import bytes_to, get_mem_info @@ -160,6 +169,8 @@ def __init__( self._kernel_handles_cyclic_swa: bool = False # lookup of unmanaged resources self._unmanaged_resources: List[str] = [] + self._replay_work_items: Optional[torch.Tensor] = None + self._replay_n_writes: Optional[torch.Tensor] = None self._spec_config = spec_config self._requires_uniform_kv_caches = requires_uniform_kv_caches self._reject_unmanaged_persistent_caches = reject_unmanaged_persistent_caches @@ -524,6 +535,16 @@ def _identify_managed_state_resources( for name, handler in self._resource_lookup.items() if isinstance(handler, ReplayPrevNumAcceptedHandler) ] + replay_work_items = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayWorkItemsHandler) + ] + replay_n_writes = [ + (name, handler) + for name, handler in self._resource_lookup.items() + if isinstance(handler, ReplayNWritesHandler) + ] # When speculative decoding is enabled, the backend must supply matching spec buffers. # When it is not enabled, spec buffers may still be registered by the backend (e.g. @@ -551,6 +572,14 @@ def _identify_managed_state_resources( f"Replay bundle mismatch: replay_prev_num_accepted has " f"{len(replay_prev_num_accepted)} entries, expected {n} (== len(ssm_managed))" ) + assert len(replay_work_items) == n, ( + f"Replay bundle mismatch: replay_work_items has {len(replay_work_items)} " + f"entries, expected {n} (== len(ssm_managed))" + ) + assert len(replay_n_writes) == n, ( + f"Replay bundle mismatch: replay_n_writes has {len(replay_n_writes)} " + f"entries, expected {n} (== len(ssm_managed))" + ) if self._spec_config is not None: if not use_replay: assert len(ssm_spec) == len(ssm_managed), ( @@ -574,6 +603,8 @@ def _identify_managed_state_resources( replay_old_dA_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, ) def _prepare_kv_cache_config( @@ -767,6 +798,8 @@ def _create_and_assign_state_views( replay_old_dA_cumsum: list = (), replay_cache_buf_idx: list = (), replay_prev_num_accepted: list = (), + replay_work_items: list = (), + replay_n_writes: list = (), ) -> Tuple[MambaHybridCacheManager, int]: """Create MambaHybridCacheManager and assign views for state resources. @@ -784,12 +817,15 @@ def _create_and_assign_state_views( conv_spec: List of speculative Conv resources. replay_old_x/B/dt/dA_cumsum: Per-layer replay cache resource lists. replay_cache_buf_idx/prev_num_accepted: Global replay resource lists. + replay_work_items/replay_n_writes: Per-forward replay metadata resources. Returns: Tuple of (manager, num_managed_mamba_layers). """ # Detect replay mode from presence of ReplayOldXHandler resources. use_replay = len(replay_old_x) > 0 + if use_replay and self._spec_config is None: + raise RuntimeError("Replay SSM state update requires speculative decoding config.") # Mamba state params can be derived from reference handlers and number of managed (non-speculative) resources. mamba_params = self._get_mamba_state_params( @@ -867,8 +903,73 @@ def _create_and_assign_state_views( for buf_name, _ in buf_list: self._caches[buf_name] = global_tensor + if replay_work_items: + self._replay_work_items = torch.empty( + self.info.max_num_state_slots, + REPLAY_WORK_ITEM_WIDTH, + device=self.info.device, + dtype=torch.int32, + ) + for buf_name, _ in replay_work_items: + self._caches[buf_name] = self._replay_work_items + + if replay_n_writes: + self._replay_n_writes = torch.zeros(1, device=self.info.device, dtype=torch.int32) + for buf_name, _ in replay_n_writes: + self._caches[buf_name] = self._replay_n_writes + return manager, num_managed_mamba_layers + def prepare_replay_metadata(self) -> None: + """Populate replay work items for the current AD batch.""" + if self._replay_work_items is None or self._replay_n_writes is None: + return + if not self.info.batch_info.is_use_replay(): + return + if not hasattr(self._kv_cache_manager, "get_replay_state_update_metadata"): + return + + replay_metadata = self._kv_cache_manager.get_replay_state_update_metadata() + self._replay_n_writes.zero_() + if replay_metadata is None: + return + + num_prefill, num_extend, _ = self.info.batch_info.get_num_sequences() + if num_extend == 0: + return + + slot_idx = self.info.get_arg("slot_idx", truncate=True) + cache_slot = slot_idx[num_prefill : num_prefill + num_extend].to(torch.int32) + cache_slot_idx = cache_slot.to(torch.long) + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens + cache_buf_idx = replay_metadata.cache_buf_idx + + position_in_decode_batch = torch.arange( + num_extend, dtype=torch.int32, device=slot_idx.device + ) + pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) + active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) + + # Keep field order and write-first partitioning in sync with the + # PyTorch replay metadata path in mamba2_metadata.py. + writes = pnat + replay_metadata.replay_step_width > replay_metadata.replay_history_size + writes_i32 = writes.to(torch.int32) + write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 + n_writes = torch.sum(writes_i32, dim=0, keepdim=True).to(torch.int32) + no_write_offsets = position_in_decode_batch - write_offsets + output_offsets = torch.where(writes, write_offsets, n_writes + no_write_offsets).to( + torch.long + ) + + work_items = self._replay_work_items[:num_extend] + work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH].scatter_( + 0, output_offsets, position_in_decode_batch + ) + work_items[:, REPLAY_WORK_CACHE_SLOT].scatter_(0, output_offsets, cache_slot) + work_items[:, REPLAY_WORK_PNAT].scatter_(0, output_offsets, pnat) + work_items[:, REPLAY_WORK_CACHE_BUF_IDX].scatter_(0, output_offsets, active_cache_buf_idx) + self._replay_n_writes.copy_(n_writes) + def _assign_kv_cache_views(self, kv_managed: Dict[str, KVPagedResourceHandler]) -> int: """Retrieve and assign buffer views for managed KV paged resources. @@ -923,6 +1024,8 @@ def _validate_no_unmanaged_persistent_caches( replay_old_dA_cumsum: list, replay_cache_buf_idx: list, replay_prev_num_accepted: list, + replay_work_items: list, + replay_n_writes: list, ) -> None: """Validate persistent cache resources are cache-manager backed. @@ -948,6 +1051,8 @@ def _validate_no_unmanaged_persistent_caches( replay_old_dA_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, ): managed_names.update(name for name, _ in replay_resources) @@ -1101,6 +1206,8 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: replay_old_dA_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, ) = self._identify_managed_state_resources() # Propagate replay mode into BatchInfo so SSM backends can branch on it @@ -1142,6 +1249,8 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: replay_old_dA_cumsum=replay_old_dA_cumsum, replay_cache_buf_idx=replay_cache_buf_idx, replay_prev_num_accepted=replay_prev_num_accepted, + replay_work_items=replay_work_items, + replay_n_writes=replay_n_writes, ) else: self._kv_cache_manager = KVCacheManager(**kv_cache_kwargs) @@ -1199,6 +1308,8 @@ def _create_kv_cache_manager(self, max_tokens: Optional[int] = None) -> Dict: replay_old_dA_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, ) # 8. Allocate remaining unmanaged resources @@ -1445,6 +1556,8 @@ def _clear_caches(self) -> None: for k in self._caches: self._caches[k] = None self._unmanaged_resources.clear() + self._replay_work_items = None + self._replay_n_writes = None def shutdown(self) -> None: """Shutdown and release all resources.""" diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 4ddd67707bb2..5cab7283f033 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -25,6 +25,12 @@ CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._utils import prefer_pinned +REPLAY_WORK_POSITION_IN_DECODE_BATCH = 0 +REPLAY_WORK_CACHE_SLOT = 1 +REPLAY_WORK_PNAT = 2 +REPLAY_WORK_CACHE_BUF_IDX = 3 +REPLAY_WORK_ITEM_WIDTH = 4 + @triton.jit def _cu_seqlens_triton_kernel( @@ -137,19 +143,26 @@ def cu_seqlens_to_chunk_indices_offsets( chunk_size: int) -> Tuple[torch.Tensor, torch.Tensor]: """ Args: - cu_seqlens (torch.Tensor): 1D tensor of cumulative sequence lengths, shape (num_seqs + 1,). The first element should be 0. Each entry represents the starting index of a sequence in the flattened token array. + cu_seqlens (torch.Tensor): 1D tensor of cumulative sequence lengths, + shape (num_seqs + 1,). The first element should be 0. Each entry + represents the starting index of a sequence in the flattened token + array. chunk_size (int): The size of each physical mamba chunk (number of tokens per chunk). Returns: Tuple[torch.Tensor, torch.Tensor]: A tuple containing: - chunk_indices (torch.Tensor): 1D tensor of indices indicating the physical chunk for each logical chunk. - - chunk_offsets (torch.Tensor): 1D tensor of offsets indicating the starting index of each logical chunk within its physical chunk. + - chunk_offsets (torch.Tensor): 1D tensor of offsets indicating + the starting index of each logical chunk within its physical + chunk. This function computes the chunk indices and offsets for the given cu_seqlens and chunk_size. Both are tensors of integers with length N, where N is the number of logical (pseudo) chunks. - A logical chunk is a sequence of tokens that are all part of the same sequence and are all in the same physical mamba chunk. + A logical chunk is a sequence of tokens that are all part of the same sequence + and are all in the same physical mamba chunk. In other words, a logical chunk changes every time we cross a sequence boundary or a physical mamba chunk boundary. - Logical chunks are needed to handle batched requests with initial states (see _state_passing_fwd and _chunk_scan_fwd). + Logical chunks are needed to handle batched requests with initial states + (see _state_passing_fwd and _chunk_scan_fwd). The chunk_indices tensor contains the index of the physical chunk for each logical chunk. The chunk_offsets tensor contains the offset (AKA starting index) of the logical chunk in the physical chunk. @@ -161,9 +174,12 @@ def cu_seqlens_to_chunk_indices_offsets( In this example, we have 2 sequences, each with 5 tokens. The physical chunk size is 8 tokens. We have three logical chunks: - - the first logical chunk starts at token 0 in the first physical chunk and contains all 5 tokens from the first sequence - - the second logical chunk starts at token 5 in the first physical chunk and contains first 3 tokens from the second sequence - - the third logical chunk starts at token 0 in the second physical chunk and contains the remaining 2 tokens from the second sequence + - the first logical chunk starts at token 0 in the first physical chunk and + contains all 5 tokens from the first sequence + - the second logical chunk starts at token 5 in the first physical chunk + and contains first 3 tokens from the second sequence + - the third logical chunk starts at token 0 in the second physical chunk + and contains the remaining 2 tokens from the second sequence """ total_seqlens = cu_seqlens[-1] @@ -232,6 +248,13 @@ def __init__(self, max_batch_size: int, chunk_size: int): # CUDA graph replays. self._state_indices_aliased_ptr = None + self.replay_work_items = torch.zeros(max_batch_size, + REPLAY_WORK_ITEM_WIDTH, + dtype=torch.int32, + device="cuda") + self.replay_n_writes = torch.zeros(1, dtype=torch.int32, device="cuda") + self.replay_num_decodes = 0 + # Pre-allocated buffers. self._arange_buffer = torch.arange(max_batch_size + 1, dtype=torch.int, @@ -241,6 +264,61 @@ def __init__(self, max_batch_size: int, chunk_size: int): dtype=torch.long, device="cuda") + def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, + num_contexts: int): + self.replay_num_decodes = 0 + if not getattr(kv_cache_manager, 'use_replay_state_update', False): + return + num_decodes = batch_size - num_contexts + self.replay_num_decodes = num_decodes + self.replay_n_writes.zero_() + if num_decodes == 0: + return + if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): + raise RuntimeError( + "Replay state update is enabled, but the KV cache manager " + "does not expose replay state update metadata.") + + replay_metadata = kv_cache_manager.get_replay_state_update_metadata() + if replay_metadata is None: + raise RuntimeError( + "Replay state update is enabled for a decode batch, but the " + "KV cache manager returned no replay state update metadata.") + + prev_num_accepted_tokens = replay_metadata.prev_num_accepted_tokens + cache_buf_idx = replay_metadata.cache_buf_idx + replay_step_width = replay_metadata.replay_step_width + replay_history_size = replay_metadata.replay_history_size + + position_in_decode_batch = torch.arange( + num_decodes, dtype=torch.int32, device=self.state_indices.device) + cache_slot = self.state_indices[num_contexts:batch_size] + cache_slot_idx = cache_slot.to(torch.long) + pnat = prev_num_accepted_tokens[cache_slot_idx].to(torch.int32) + active_cache_buf_idx = cache_buf_idx[cache_slot_idx].to(torch.int32) + + # Keep field order and write-first partitioning in sync with the + # AutoDeploy replay metadata path in shim/interface.py. + writes = (pnat + replay_step_width > replay_history_size) + writes_i32 = writes.to(torch.int32) + write_offsets = torch.cumsum(writes_i32, dim=0) - writes_i32 + n_writes = torch.sum(writes_i32, dim=0, keepdim=True).to(torch.int32) + no_write_offsets = position_in_decode_batch - write_offsets + output_offsets = torch.where(writes, write_offsets, + n_writes + no_write_offsets) + output_offsets = output_offsets.to(torch.long) + + work_items = self.replay_work_items[:num_decodes] + work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH].scatter_( + 0, output_offsets, position_in_decode_batch) + work_items[:, REPLAY_WORK_CACHE_SLOT].scatter_(0, output_offsets, + cache_slot) + work_items[:, REPLAY_WORK_PNAT].scatter_(0, output_offsets, pnat) + work_items[:, + REPLAY_WORK_CACHE_BUF_IDX].scatter_(0, output_offsets, + active_cache_buf_idx) + self.replay_n_writes.copy_(n_writes) + def prepare(self, attn_metadata: AttentionMetadata): batch_size = attn_metadata.seq_lens.shape[0] num_contexts = attn_metadata.num_contexts @@ -254,9 +332,11 @@ def prepare(self, attn_metadata: AttentionMetadata): and hasattr(kv_cache_manager, 'get_state_indices') and request_ids is not None): batch_request_ids = request_ids[:batch_size] + max_draft_len = getattr(kv_cache_manager, + "speculative_num_draft_tokens", 0) or 0 is_padding = [ - req_id == CUDA_GRAPH_DUMMY_REQUEST_ID - for req_id in batch_request_ids + CUDA_GRAPH_DUMMY_REQUEST_ID - max_draft_len <= req_id <= + CUDA_GRAPH_DUMMY_REQUEST_ID for req_id in batch_request_ids ] indices = kv_cache_manager.get_state_indices( batch_request_ids, is_padding) @@ -268,12 +348,10 @@ def prepare(self, attn_metadata: AttentionMetadata): # cudaStreamSynchronize per element. # # Safe under CUDA graphs only when the source buffer has a - # stable data pointer across all calls (currently true for - # CppMambaHybridCacheManager.cuda_state_indices, allocated - # once in __init__). If a future cache manager reallocates - # this buffer between iterations, captured kernels would - # still read from the address seen at capture time, so we - # assert stability here. + # stable data pointer across all calls. If a cache manager + # reallocates this buffer between iterations, captured kernels + # would still read from the address seen at capture time, so + # we assert stability here. if self._state_indices_aliased_ptr is None: self._state_indices_aliased_ptr = indices.data_ptr() else: @@ -297,6 +375,9 @@ def prepare(self, attn_metadata: AttentionMetadata): self.state_indices[:batch_size].copy_( self.state_indices_cpu[:batch_size], non_blocking=True) + self._prepare_replay_work_items(kv_cache_manager, batch_size, + num_contexts) + if num_contexts > 0: torch.cumsum(context_lens, dim=0, @@ -360,7 +441,7 @@ def prepare(self, attn_metadata: AttentionMetadata): # Complete any deferred recurrent-state block onboards scheduled by # CppMambaHybridCacheManager.prepare_resources(). prepare_resources # only enqueues the async cudaMemcpyAsync calls and sets a pending - # flag; we sync the onboard stream here, so the prior CPU-side prep + # flag; we sync the onboard stream here, so CPU-side prep # work in _prepare_tp_inputs overlaps with the in-flight transfers. # Cheap no-op on cache managers without this method or when no # transfers were scheduled this iteration. diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py index b7d24e72f460..03f9d8068679 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py @@ -413,11 +413,22 @@ def forward( # Speculative decoding only supported with Python path assert layer_cache is not None, \ "Speculative decoding requires Python MambaCacheManager" - # TODO: support dynamic speculation, will add current_draft_len later [TRTLLM-10319] - draft_token_num = spec_metadata.max_draft_len + 1 intermediate_conv_states = layer_cache.intermediate_conv_window use_replay = getattr(attn_metadata.kv_cache_manager, 'use_replay_state_update', False) + draft_token_num = spec_metadata.runtime_draft_len + 1 + if use_replay: + replay_metadata = (attn_metadata.kv_cache_manager. + get_replay_state_update_metadata()) + assert replay_metadata is not None, ( + "Mamba replay state update is enabled but replay " + "metadata was not allocated.") + replay_step_width = replay_metadata.replay_step_width + assert draft_token_num == replay_step_width, ( + "Mamba replay state update does not support dynamic " + "draft length yet. Runtime token width " + f"{draft_token_num} must match fixed replay step " + f"width {replay_step_width}.") intermediate_state_indices = _cached_arange( attn_metadata.kv_cache_manager.get_max_resource_count(), @@ -512,11 +523,9 @@ def convert_dt(): philox_kwargs = {} if use_stochastic_rounding: - # Both replay and flashinfer read from the cache manager's - # persistent per-slot Philox seed buffer; replay indexes by - # cache_batch_idx, flashinfer reads slot 0 from a (1,) - # view. In-place add_(1) keeps CUDA-graph replay fresh - # without allocating any new CUDA tensors per forward. + # Both replay and flashinfer use a single Philox seed. The + # cache manager owns the persistent buffer; passing a (1,) + # view avoids allocating CUDA tensors per forward. rand_seed = layer_cache.mamba_ssm_rand_seed assert rand_seed is not None, ( "Mamba SSM stochastic rounding is enabled but the " @@ -524,13 +533,13 @@ def convert_dt(): "_util.py passes mamba_ssm_stochastic_rounding=True " "to the cache manager.") rand_seed.add_(1) - if use_replay: - philox_kwargs['rand_seed'] = rand_seed - else: - philox_kwargs['rand_seed'] = rand_seed[:1] + philox_kwargs['rand_seed'] = rand_seed[:1] philox_kwargs['philox_rounds'] = self._philox_rounds if use_replay: + # replay_work_items is write-first for persistent_main and + # carries decode-batch position, cache slot, PNAT, and + # active cache buffer index for replay kernels. replay_selective_state_update( ssm_states, layer_cache.old_x, @@ -549,6 +558,9 @@ def convert_dt(): dt_softplus=self.delta_softplus, state_batch_indices=state_batch_indices, out=out_4d, + n_writes=mamba_metadata.replay_n_writes, + replay_work_items=( + mamba_metadata.replay_work_items[:num_decodes]), launch_with_pdl=True, **philox_kwargs, ) diff --git a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py index 547cc3153c5f..1e3cd8932e5e 100644 --- a/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py +++ b/tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py @@ -1,5 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +# ruff: noqa: E501 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,11 +24,72 @@ import triton import triton.language as tl -from tensorrt_llm._torch.modules.mamba import PAD_SLOT_ID from tensorrt_llm._utils import get_sm_version +from .mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from .softplus import softplus +_REPLAY_WORK_POSITION_IN_DECODE_BATCH = tl.constexpr(REPLAY_WORK_POSITION_IN_DECODE_BATCH) +_REPLAY_WORK_CACHE_SLOT = tl.constexpr(REPLAY_WORK_CACHE_SLOT) +_REPLAY_WORK_PNAT = tl.constexpr(REPLAY_WORK_PNAT) +_REPLAY_WORK_CACHE_BUF_IDX = tl.constexpr(REPLAY_WORK_CACHE_BUF_IDX) +_REPLAY_WORK_ITEM_WIDTH = tl.constexpr(REPLAY_WORK_ITEM_WIDTH) +MIN_REPLAY_TILE_SIZE = 16 + + +# Naming convention: +# - step: one decode invocation of replay_selective_state_update, with token +# positions [0, T). +# - window: the replay cache capacity [0, max_window). History from previous +# steps occupies [0, PNAT). In non-rectangle replay, this history is the only +# valid window-indexed range; current-step values stay in T-space. +# - rectangle nowrite: combines both ranges in window space, with history at +# [0, PNAT), current-step tokens at [PNAT, PNAT + T), and padding after +# PNAT + T masked invalid. +# - history_*: cached prior-token values, always indexed in window space. +# - step_*: current-step values in T-space. +# - step_*_in_window: current-step values remapped to rectangle window positions. +# - *_window: combined history + step values over the full window dimension. + + +@triton.jit +def _gdc_wait_with_memory_clobber(): + tl.inline_asm_elementwise( + "griddepcontrol.wait; // dummy $0", + "=r,~{memory}", + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +# Lazy global allocator for Triton TMA tensor descriptors. Required by any +# host- or device-built tensor_descriptor; without it Triton raises at first +# launch. +_TMA_ALLOCATOR_SET = False + + +def _ensure_tma_allocator() -> None: + global _TMA_ALLOCATOR_SET + if _TMA_ALLOCATOR_SET: + return + + def _alloc_fn(size, alignment, stream): + # Triton expects an int8 buffer of `size` bytes; alignment is enforced + # by the allocator returning a buffer satisfying it (PyTorch's + # cudaMalloc-backed tensors are 256B-aligned, so we're fine). + return torch.empty(size, device="cuda", dtype=torch.int8) + + triton.set_allocator(_alloc_fn) + _TMA_ALLOCATOR_SET = True + @triton.jit def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: @@ -51,16 +113,75 @@ def _stochastic_round_fp16x2(x: tl.tensor, rand: tl.tensor) -> tl.tensor: ) -# Precompute kernel: CB_scaled, decay_vec. Writes new cache (old_B, -# old_dt, old_dA_cumsum) to the WRITE buffer slot for next step's replay. -# Grid: (batch, nheads // HEADS_PER_BLOCK). +@triton.jit +def _stochastic_round_fp8x4_e4m3(x: tl.tensor, rand: tl.tensor) -> tl.tensor: + """Stochastic rounding: fp32 quad → fp8 e4m3 using Philox random bits. + + Uses PTX cvt.rs.satfinite.e4m3x4.f32 which combines stochastic rounding + and saturating cast in a single op (output is final fp8, no separate + clamp needed). The reversed source-register order {$4,$3,$2,$1} is + load-bearing — PTX packs leftmost source into the high byte but Triton's + pack=4 is little-endian, so the natural {$1,$2,$3,$4} order would + silently shuffle every group of 4 contiguous outputs. + + Requires SM_100a+ (Blackwell B200). Caller must gate at the wrapper + level — this kernel does not check. + + Adapted from vLLM PR #40012 (Apache-2.0). + """ + return tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + + +@triton.jit +def _bitrev32(x: tl.tensor) -> tl.tensor: + return tl.inline_asm_elementwise( + asm="brev.b32 $0, $1;", + constraints="=r,r", + args=(x,), + dtype=tl.uint32, + is_pure=True, + pack=1, + ) + + +@triton.jit +def _stochastic_round_int8_packed(x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor) -> tl.tensor: + """Stochastic rounding for int8 using one random uint32 per 4 values.""" + low = rand & 0x0000FFFF + high = (rand >> 16) & 0x0000FFFF + low_rev = _bitrev32(low) >> 16 + high_rev = _bitrev32(high) >> 16 + rand_pos = offs_n & 3 + rand16 = tl.where( + rand_pos == 0, + low, + tl.where(rand_pos == 1, low_rev, tl.where(rand_pos == 2, high, high_rev)), + ) + rand01 = rand16.to(tl.float32) * (1.0 / float(1 << 16)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +@triton.jit +def _stochastic_round_int16_packed(x: tl.tensor, rand: tl.tensor, offs_n: tl.tensor) -> tl.tensor: + """Stochastic rounding for int16 using one random uint32 per 2 values.""" + rand_bits = tl.where((offs_n & 1) == 0, rand, _bitrev32(rand)) + rand01 = (rand_bits & 0x00FFFFFF).to(tl.float32) * (1.0 / float(1 << 24)) + return tl.extra.cuda.libdevice.floor(x + rand01) + + +# Replay-style precompute body. Computes CB_scaled/decay_vec in T-space and +# writes this step's B/dt/dA_cumsum to the selected cache buffer. -@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) @triton.jit() -def _replay_precompute_kernel( +def _replay_precompute_impl( # Input pointers dt_ptr, dt_bias_ptr, @@ -70,14 +191,21 @@ def _replay_precompute_kernel( # Output pointers cb_scaled_ptr, decay_vec_ptr, - # Cache WRITE pointers (write-buffer for next step) + # Cache pointers (both buffers reachable via stride_*_dbuf). This + # kernel writes to either the active (= cache_buf_idx) or inactive + # (= 1 - cache_buf_idx) buffer depending on WRITE_CHECKPOINT — see + # comment block at top of kernel body. old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - # Double-buffer index (per cache slot) + # Double-buffer index (per cache slot) — selects this step's "active" + # buffer (= where the historical inputs for this step live). cache_buf_idx_ptr, + # Per-request accepted-tokens count (already-cached old tokens at + # [0, PNAT) of the active buffer; new tokens this step go after them + # on no-replay-write steps). + prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - pad_slot_id, # Dimensions T: tl.constexpr, dstate: tl.constexpr, @@ -126,37 +254,46 @@ def _replay_precompute_kernel( # Meta-parameters DT_SOFTPLUS: tl.constexpr, HAS_DT_BIAS: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, - LAUNCH_DEPENDENT_KERNELS: tl.constexpr, HEADS_PER_BLOCK: tl.constexpr, + # Checkpoint write flag — selects target buffer + offset for new-token + # cache writes. See "Cache write semantics" block below. + # Runtime (not constexpr): the only checkpoint-dependent code in + # this body is the write_buf/write_offset selection, which is plain + # arithmetic — no constexpr-shaped tile or whole-block gate. Letting + # it be runtime lets the dynamic dispatch kernel call us once with the + # per-slot needs_write flag instead of inlining two specializations. + needs_checkpoint_write, ): pid_b = tl.program_id(axis=0) pid_hg = tl.program_id(axis=1) # head-group index first_head = pid_hg * HEADS_PER_BLOCK - # Resolve cache index for writes - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) - - # Signal main kernel to start (internal PDL). Main's replay phase - # reads only from the READ buffer (written by the PREVIOUS step) — - # safe even if conv1d and this kernel are still running. Main's - # gdc_wait() gates the output phase, which reads conv1d outputs - # (x, C) and this kernel's outputs (cb_scaled, decay_vec). - if LAUNCH_DEPENDENT_KERNELS: - tl.extra.cuda.gdc_launch_dependents() + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - # Read buffer index: replay reads from buf_read. We WRITE to 1 - buf_read - # for next step's replay. - buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) - buf_write = 1 - buf_read + # --- Cache write semantics --- + # cache_buf_idx names this step's "active" buffer — the one with the + # historical inputs at [0, PNAT). The other buffer is "staging". + # + # Where do we write new tokens this step? + # WRITE_CHECKPOINT=False (no overflow): append to ACTIVE buffer at + # offset [PNAT : PNAT+T). Caller does NOT flip cache_buf_idx + # afterward; PNAT_next = PNAT + accepted. [0, PNAT) preserved. + # WRITE_CHECKPOINT=True (would overflow): write to STAGING buffer at + # [0, T). Caller flips cache_buf_idx afterward; next step's + # active = the one we just wrote. PNAT_next = accepted. The + # previous active-buffer history is folded into state by replay + # before the caller flips to the staging buffer. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + if needs_checkpoint_write: + write_buf = 1 - buf_active + write_offset = 0 + else: + write_buf = buf_active + write_offset = prev_num_accepted_tokens offs_t = tl.arange(0, BLOCK_SIZE_T) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) @@ -167,210 +304,183 @@ def _replay_precompute_kernel( causal_mask = offs_t[:, None] >= offs_t[None, :] valid_mask = causal_mask & t_mask[:, None] & t_mask[None, :] - # --- Loop 1: compute per-head dt/dA_cumsum/decay BEFORE gdc_wait --- - # These only depend on dt (from in_proj, not conv1d) and parameters (A, dt_bias). - # Store to cache; will reload after the wait for CB scaling. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - dt_base = dt_ptr + pid_b * stride_dt_batch + head_idx * stride_dt_head - dt = tl.load(dt_base + offs_t * stride_dt_T, mask=t_mask, other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt_bias = tl.load(dt_bias_ptr + head_idx * stride_dt_bias_head).to(tl.float32) - dt = dt + dt_bias - if DT_SOFTPLUS: - dt = softplus(dt) - - A = tl.load(A_ptr + head_idx * stride_A_head).to(tl.float32) - dA_cumsum = tl.cumsum(A * dt, axis=0) - decay_vec = tl.exp(dA_cumsum) - - # Store dt, dA_cumsum, decay_vec to cache - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_write * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - tl.store(old_dt_base + offs_t * stride_old_dt_T, dt, mask=t_mask) + # --- Pre-wait phase across HEADS_PER_BLOCK heads --- + # Compute dt, dA_cumsum, decay_vec, and scale_combo as head-block tiles. + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h # (H,) + + # Load dt (H, T) + dt_addrs = ( + dt_ptr + + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T + ) + dt = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt_bias = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt = dt + dt_bias[:, None] + if DT_SOFTPLUS: + dt = softplus(dt) - old_dA_cumsum_base = ( + A = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) # (H,) + dA_cumsum = tl.cumsum(A[:, None] * dt, axis=1) # (H, T) + decay_vec = tl.exp(dA_cumsum) # (H, T) + + # Cross-step continuity for old_dA_cumsum: when appending to active_buf at + # offset PNAT > 0, the previous step left a running cumsum at [0, PNAT) + # whose tail value lives at active_buf[head, PNAT-1]. Add that tail to + # this step's per-step-restarted cumsum before storing so the buffer + # holds one continuous cumsum across N back-to-back nowrites. Write path + # (write_buf = 1 - buf_active, write_offset = 0) starts fresh, no prefix. + # Both branches are on scalar runtime values (checkpoint predicate and + # PNAT), uniform across the block — use scalar if to short-circuit the load. + if needs_checkpoint_write or prev_num_accepted_tokens == 0: + prev_total = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + last_cumsum_ptrs = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_write * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T ) - tl.store(old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, dA_cumsum, mask=t_mask) + prev_total = tl.load(last_cumsum_ptrs).to(tl.float32) + + # Store dt, dA_cumsum to cache at [write_offset : write_offset+T) of write_buf. + old_dt_addrs = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T + ) + tl.store(old_dt_addrs, dt, mask=t_mask[None, :]) + + old_dA_cumsum_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store(old_dA_cumsum_addrs, dA_cumsum + prev_total[:, None], mask=t_mask[None, :]) + + # decay_vec scratch — always at offs_t. + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) + tl.store(decay_vec_addrs, decay_vec, mask=t_mask[None, :]) - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + head_idx * stride_dv_head - tl.store(decay_vec_base + offs_t * stride_dv_t, decay_vec, mask=t_mask) + # scale_combo (H, T, T) = exp(dA_cumsum[h, t1] - dA_cumsum[h, t2]) * dt[h, t2] + decay_matrix = tl.exp(dA_cumsum[:, :, None] - dA_cumsum[:, None, :]) # (H, T, T) + scale_combo = decay_matrix * dt[:, None, :] # (H, T, T) - # --- Wait for upstream kernel (external PDL) before loading B and C --- - # All dt processing above is independent of conv1d outputs. + # Wait for conv1d before loading this step's B/C. if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() + _gdc_wait_with_memory_clobber() # --- Load C and B once for the group (shared across HEADS_PER_BLOCK heads) --- group_idx = first_head // nheads_ngroups_ratio C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group - C_all = tl.load( + C_tile = tl.load( C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) - B_all = tl.load( + B_tile = tl.load( B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, mask=t_mask[:, None] & n_mask[None, :], other=0.0, ) - # Compute raw CB once, shared across all heads in this block. - raw_CB = tl.dot(C_all.to(tl.bfloat16), tl.trans(B_all).to(tl.bfloat16)) + # Compute raw CB once — shared across all heads in this block + raw_CB = tl.dot(C_tile.to(tl.bfloat16), tl.trans(B_tile).to(tl.bfloat16)) - # Store B to cache (once per group, only if this block covers the first heads) + # Store B to cache at [write_offset : write_offset+T) of write_buf. if first_head % nheads_ngroups_ratio == 0: old_B_base = ( old_B_ptr + cache_batch_idx * stride_old_B_cache - + buf_write * stride_old_B_dbuf + + write_buf * stride_old_B_dbuf + group_idx * stride_old_B_group ) tl.store( - old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - B_all, + old_B_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + B_tile, mask=t_mask[:, None] & n_mask[None, :], ) - # --- Loop 2: reload per-head dA_cumsum/dt from cache, scale CB --- - # The cache was just written above, so these loads should hit L2. - for h_local in range(HEADS_PER_BLOCK): - head_idx = first_head + h_local - - # Reload dt and dA_cumsum from cache (just written in loop 1) - old_dt_base = ( - old_dt_ptr - + cache_batch_idx * stride_old_dt_cache - + buf_write * stride_old_dt_dbuf - + head_idx * stride_old_dt_head - ) - dt = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to(tl.float32) - - old_dA_cumsum_base = ( - old_dA_cumsum_ptr - + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_write * stride_old_dA_cumsum_dbuf - + head_idx * stride_old_dA_cumsum_head - ) - dA_cumsum = tl.load( - old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 - ).to(tl.float32) - - # Scale raw_CB with per-head decay and dt - decay_matrix = tl.exp(dA_cumsum[:, None] - dA_cumsum[None, :]) - CB_scaled = tl.where(valid_mask, raw_CB * decay_matrix * dt[None, :], 0.0) - - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + head_idx * stride_cb_head - tl.store( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - CB_scaled, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - ) - - -# Main kernel: tl.dot replay + precomputed CB output. -# Grid: (cdiv(dim, M), batch, nheads). + # --- Post-wait CB precompute --- + # Combine raw_CB with scale_combo, apply the causal mask, and store one + # (H, T, T) tile. + CB_scaled_block = tl.where( + valid_mask[None, :, :], + raw_CB[None, :, :] * scale_combo, + 0.0, + ) # (H, T, T) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_t[None, None, :] * stride_cb_j + ) # (H, T, T) + cb_store_mask = (offs_t[None, :, None] < BLOCK_SIZE_T) & (offs_t[None, None, :] < BLOCK_SIZE_T) + tl.store(cb_scaled_addrs, CB_scaled_block, mask=cb_store_mask) -@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) -@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) -@triton.heuristics( - {"HAS_CACHE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None} -) -@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) -@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) -@triton.heuristics({"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), 16)}) +# Rectangle nowrite precompute body. Builds a window-space CB tile with +# history at [0, PNAT) and this step's values at [PNAT, PNAT + T). @triton.jit() -def _replay_state_update_kernel( - # Pointers - state_ptr, - # Cache READ pointers (read-buffer from previous step) - old_x_ptr, +def _rectangle_precompute_impl( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, # rectangle window: (batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K) + decay_vec_ptr, # total_decay * exp(cumAdt_new[t]): (batch, nheads, BLOCK_SIZE_T) + # Cache pointers (both buffers reachable via stride_*_dbuf). Nowrite + # path: read from buf_active at [0, PNAT), write new tokens at + # [PNAT, PNAT+T) of buf_active (same buffer). old_B_ptr, old_dt_ptr, old_dA_cumsum_ptr, - # Cache WRITE pointer (write-buffer for old_x only; B/dt/dA_cumsum written by precompute) - prev_num_accepted_tokens_ptr, cache_buf_idx_ptr, - # New input pointers - x_ptr, - C_ptr, - D_ptr, - z_ptr, - out_ptr, - # Precomputed pointers - cb_scaled_ptr, - decay_vec_ptr, + prev_num_accepted_tokens_ptr, state_batch_indices_ptr, - # Stochastic rounding - rand_seed_ptr, - pad_slot_id, # Dimensions T: tl.constexpr, - dim: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, # rectangle window bound dstate: tl.constexpr, nheads_ngroups_ratio: tl.constexpr, - # state strides - stride_state_batch, - stride_state_head, - stride_state_dim, - stride_state_dstate, - # old_x strides: (cache, T, nheads, dim) — single-buffered - stride_old_x_cache, - stride_old_x_T, - stride_old_x_head, - stride_old_x_dim, - # old_B strides: (cache, 2, T, ngroups, dstate) - stride_old_B_cache, - stride_old_B_dbuf, - stride_old_B_T, - stride_old_B_group, - stride_old_B_dstate, - # old_dt strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dt_cache, - stride_old_dt_dbuf, - stride_old_dt_head, - stride_old_dt_T, - # old_dA_cumsum strides: (cache, 2, nheads, T) — T contiguous for coalesced access - stride_old_dA_cumsum_cache, - stride_old_dA_cumsum_dbuf, - stride_old_dA_cumsum_head, - stride_old_dA_cumsum_T, - # x strides - stride_x_batch, - stride_x_T, - stride_x_head, - stride_x_dim, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, # C strides stride_C_batch, stride_C_T, stride_C_group, stride_C_dstate, - # D strides - stride_D_head, - stride_D_dim, - # z strides - stride_z_batch, - stride_z_T, - stride_z_head, - stride_z_dim, - # out strides - stride_out_batch, - stride_out_T, - stride_out_head, - stride_out_dim, - # cb_scaled strides + # cb_scaled strides (rectangle: batch, nheads, T, window) stride_cb_batch, stride_cb_head, stride_cb_t, @@ -379,225 +489,2923 @@ def _replay_state_update_kernel( stride_dv_batch, stride_dv_head, stride_dv_t, - # Meta - BLOCK_SIZE_M: tl.constexpr, - HAS_D: tl.constexpr, - HAS_Z: tl.constexpr, - HAS_CACHE_BATCH_INDICES: tl.constexpr, + # old_B strides: (cache, 2, T_max, ngroups, dstate) + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides: (cache, 2, nheads, T_max) + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides: (cache, 2, nheads, T_max) + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, LAUNCH_WITH_PDL: tl.constexpr, - USE_RS_ROUNDING: tl.constexpr, - PHILOX_ROUNDS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + USE_GATHER_FOR_NEW_TOKENS: tl.constexpr, ): - pid_m = tl.program_id(axis=0) - pid_b = tl.program_id(axis=1) - pid_h = tl.program_id(axis=2) - - if HAS_CACHE_BATCH_INDICES: - cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - if cache_batch_idx == pad_slot_id: - return - else: - cache_batch_idx = pid_b.to(tl.int64) + pid_b = tl.program_id(axis=0) + pid_hg = tl.program_id(axis=1) + first_head = pid_hg * HEADS_PER_BLOCK - # Double-buffer index: buf_read points to the buffer written by LAST step's - # precompute. THIS step's precompute writes to 1-buf_read, which will be - # read by NEXT step's main kernel. Anything not carried between steps is - # single-buffered. - buf_read = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) - offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + # Nowrite-only: write_buf = active, write_offset = PNAT. No flip after. + buf_active = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + write_buf = buf_active + write_offset = prev_num_accepted_tokens + + # Rectangle window layout: history at [0, PNAT), step tokens at [PNAT, PNAT+T). + # PNAT + T <= MAX is guaranteed on the nowrite path → no overlap. + + offs_t = tl.arange(0, BLOCK_SIZE_T) # current-step output rows + offs_window = tl.arange(0, BLOCK_SIZE_K) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - offs_t = tl.arange(0, BLOCK_SIZE_T) - m_mask = offs_m < dim - n_mask = offs_n < dstate t_mask = offs_t < T + n_mask = offs_n < dstate - # Load state - state_ptr += cache_batch_idx * stride_state_batch + pid_h * stride_state_head - state_ptrs = ( - state_ptr + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + # Window masks. Cache and matmul share this logical dimension. + is_history_position = offs_window < prev_num_accepted_tokens + safe_history_idx = tl.where(is_history_position, offs_window, 0) + step_idx_from_window = offs_window - prev_num_accepted_tokens + is_step_position = (step_idx_from_window >= 0) & (step_idx_from_window < T) + safe_step_idx = tl.where(is_step_position, step_idx_from_window, 0) + + offs_h = tl.arange(0, HEADS_PER_BLOCK) + heads_block = first_head + offs_h + + # Precompute this step's (H, T) dt and continuous dA_cumsum tiles. + dt_addrs = ( + dt_ptr + + pid_b * stride_dt_batch + + heads_block[:, None] * stride_dt_head + + offs_t[None, :] * stride_dt_T ) - state_mask = m_mask[:, None] & n_mask[None, :] - state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) - prev_num_accepted_tokens = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + dt_new = tl.load(dt_addrs, mask=t_mask[None, :], other=0.0).to(tl.float32) - # Phase 1: Replay via tl.dot fast-forward (reads from READ buffer) - group_idx = pid_h // nheads_ngroups_ratio + if HAS_DT_BIAS: + dt_bias_heads = tl.load(dt_bias_ptr + heads_block * stride_dt_bias_head).to(tl.float32) + dt_new = dt_new + dt_bias_heads[:, None] + if DT_SOFTPLUS: + dt_new = softplus(dt_new) - # Load precomputed dt and dA_cumsum from READ buffer - old_dt_base = ( + A_heads = tl.load(A_ptr + heads_block * stride_A_head).to(tl.float32) + dA_cumsum_step = tl.cumsum(A_heads[:, None] * dt_new, axis=1) + + if prev_num_accepted_tokens == 0: + dA_cumsum_prefix = tl.zeros((HEADS_PER_BLOCK,), dtype=tl.float32) + else: + dA_cumsum_prefix_ptrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head + + (prev_num_accepted_tokens - 1) * stride_old_dA_cumsum_T + ) + dA_cumsum_prefix = tl.load(dA_cumsum_prefix_ptrs).to(tl.float32) + + old_dt_write_addrs = ( old_dt_ptr + cache_batch_idx * stride_old_dt_cache - + buf_read * stride_old_dt_dbuf - + pid_h * stride_old_dt_head + + write_buf * stride_old_dt_dbuf + + heads_block[:, None] * stride_old_dt_head + + (write_offset + offs_t)[None, :] * stride_old_dt_T ) - old_dt_all = tl.load(old_dt_base + offs_t * stride_old_dt_T, mask=t_mask, other=0.0).to( - tl.float32 + tl.store(old_dt_write_addrs, dt_new, mask=t_mask[None, :]) + + old_dA_cumsum_write_addrs = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + write_buf * stride_old_dA_cumsum_dbuf + + heads_block[:, None] * stride_old_dA_cumsum_head + + (write_offset + offs_t)[None, :] * stride_old_dA_cumsum_T + ) + tl.store( + old_dA_cumsum_write_addrs, + dA_cumsum_step + dA_cumsum_prefix[:, None], + mask=t_mask[None, :], ) - old_dA_cumsum_base = ( + # Build the history side of combo_block from cached data. + group_idx = first_head // nheads_ngroups_ratio + + # Group-level: history B from active buffer at [0, PNAT) of the window. + old_B_read_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + buf_active * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + history_B = tl.load( + old_B_read_base + + safe_history_idx[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + mask=is_history_position[:, None] & n_mask[None, :], + other=0.0, + ) + + # Per-head read bases (H,) - broadcast with offs_window for 2D loads. + old_dt_read_h = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + buf_active * stride_old_dt_dbuf + + heads_block * stride_old_dt_head + ) + old_dA_cumsum_read_h = ( old_dA_cumsum_ptr + cache_batch_idx * stride_old_dA_cumsum_cache - + buf_read * stride_old_dA_cumsum_dbuf - + pid_h * stride_old_dA_cumsum_head + + buf_active * stride_old_dA_cumsum_dbuf + + heads_block * stride_old_dA_cumsum_head ) - old_dA_cumsum_all = tl.load( - old_dA_cumsum_base + offs_t * stride_old_dA_cumsum_T, mask=t_mask, other=0.0 + + # (H, window) loads at [0, PNAT) from previous steps. + history_mask_h = is_history_position[None, :] + history_dt = tl.load( + old_dt_read_h[:, None] + safe_history_idx[None, :] * stride_old_dt_T, + mask=history_mask_h, + other=0.0, ).to(tl.float32) + history_dA_cumsum = tl.load( + old_dA_cumsum_read_h[:, None] + safe_history_idx[None, :] * stride_old_dA_cumsum_T, + mask=history_mask_h, + other=0.0, + ).to(tl.float32) + # Combine cached history with this step's values at [PNAT, PNAT+T). + ht_mask = t_mask[None, :] # (1, T) + dA_cumsum_new = dA_cumsum_step + dA_cumsum_prefix[:, None] # (H, T) - # Load dA_cumsum at prev_k-1 directly via pointer math (avoids masked reduction). - # Clamp to [0, T-1] defensively — out-of-contract PNAT > T would read OOB. - prev_k_idx = tl.minimum(tl.maximum(prev_num_accepted_tokens - 1, 0), T - 1) - total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( - tl.float32 + # Production uses matching padded T and window sizes, where tl.gather is the cheap + # path. Some tests use a larger padded K than T; Triton rejects that gather + # axis mismatch, so use a slower one-hot sum fallback for those cases. The + # wrapper passes this as an explicit constexpr so fast-path compilations do + # not carry the fallback branch. + if USE_GATHER_FOR_NEW_TOKENS: + step_gather_idx = tl.broadcast_to(safe_step_idx[None, :], (HEADS_PER_BLOCK, BLOCK_SIZE_K)) + step_dt_in_window = tl.where( + is_step_position[None, :], + tl.gather(dt_new, step_gather_idx, axis=1), + 0.0, + ) # (H, window) + step_dA_cumsum_in_window = tl.where( + is_step_position[None, :], + tl.gather(dA_cumsum_new, step_gather_idx, axis=1), + 0.0, + ) # (H, window) + else: + step_in_window_selector = ( + offs_t[:, None] == (offs_window[None, :] - prev_num_accepted_tokens) + ) & is_step_position[None, :] + step_dt_in_window = tl.sum( + tl.where(step_in_window_selector[None, :, :], dt_new[:, :, None], 0.0), + axis=1, + ) # (H, window) + step_dA_cumsum_in_window = tl.sum( + tl.where(step_in_window_selector[None, :, :], dA_cumsum_new[:, :, None], 0.0), + axis=1, + ) # (H, window) + + # The write buffer stores continuous cumsum, so decay_vec_full[t] is + # exp(continuous_cumsum[PNAT+t]) directly. + decay_vec_full_block = tl.exp(dA_cumsum_new) # (H, T) + decay_vec_addrs = ( + decay_vec_ptr + + pid_b * stride_dv_batch + + heads_block[:, None] * stride_dv_head + + offs_t[None, :] * stride_dv_t + ) # (H, T) + tl.store(decay_vec_addrs, decay_vec_full_block, mask=ht_mask) + + # combo_block[t, j] = dt[j] * exp(cumsum[t] - cumsum[j]). + # History and step positions share the same formula because both halves use + # continuous cumsum values. + dt_factor_window = tl.where( + is_history_position[None, :], history_dt, step_dt_in_window + ) # (H, window) + neg_dA_cumsum_window = tl.where( + is_history_position[None, :], + -history_dA_cumsum, + -step_dA_cumsum_in_window, + ) # (H, window) + # exp_diff (H, T, window) = exp(-cumsum[j] (H, 1, window) + cumsum[t] (H, T, 1)). + exp_diff = tl.exp(neg_dA_cumsum_window[:, None, :] + dA_cumsum_new[:, :, None]) + combo_block = dt_factor_window[:, None, :] * exp_diff # (H, T, window) + + # Wait for conv1d before loading this step's B/C. + if LAUNCH_WITH_PDL: + _gdc_wait_with_memory_clobber() + + C_base = C_ptr + pid_b * stride_C_batch + group_idx * stride_C_group + step_B_base = B_ptr + pid_b * stride_B_batch + group_idx * stride_B_group + + C_tile = tl.load( + C_base + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + step_B = tl.load( + step_B_base + offs_t[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + step_B_in_window = tl.load( + step_B_base + safe_step_idx[:, None] * stride_B_T + offs_n[None, :] * stride_B_dstate, + mask=is_step_position[:, None] & n_mask[None, :], + other=0.0, ) + # Disjoint masks: history at [0, PNAT), step tokens at [PNAT, PNAT+T). + B_window = history_B + step_B_in_window + raw_rect_CB = tl.dot(C_tile.to(tl.bfloat16), tl.trans(B_window).to(tl.bfloat16)) - # Step 0 invariant: PNAT=0 means `state` is already last step's state (not - # two back). coeff is all-zero (offs_t < 0), total_decay is 1.0, so the - # replay leaves `state` unchanged — cache contents don't matter on step 0. - coeff = tl.exp(total_dA_cumsum - old_dA_cumsum_all) * old_dt_all - accepted_mask = t_mask & (offs_t < prev_num_accepted_tokens) - coeff = tl.where(accepted_mask, coeff, 0.0) + # Append new B to cache at [PNAT, PNAT+T) of write_buf (once per group). + if first_head % nheads_ngroups_ratio == 0: + old_B_write_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + write_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + tl.store( + old_B_write_base + + (write_offset + offs_t)[:, None] * stride_old_B_T + + offs_n[None, :] * stride_old_B_dstate, + step_B, + mask=t_mask[:, None] & n_mask[None, :], + ) + + # Causal mask (BLOCK_SIZE_T × window, shared across heads). + # Step tokens occupy runtime positions [PNAT, PNAT+T). + t_idx_2d = offs_t[:, None] + window_idx_2d = offs_window[None, :] + is_history_position_2d = window_idx_2d < prev_num_accepted_tokens + step_idx_from_window_2d = window_idx_2d - prev_num_accepted_tokens + is_step_causal_2d = ( + (step_idx_from_window_2d >= 0) + & (step_idx_from_window_2d < T) + & (step_idx_from_window_2d <= t_idx_2d) + ) + causal_combined = (is_history_position_2d | is_step_causal_2d) & t_mask[:, None] + + # rect_CB_scaled = where(causal, raw_rect_CB * combo_block, 0); store one + # (H, T, window) tile. + rect_CB_scaled_block = tl.where( + causal_combined[None, :, :], + raw_rect_CB[None, :, :] * combo_block, + 0.0, + ) # (H, T, window) + cb_scaled_addrs = ( + cb_scaled_ptr + + pid_b * stride_cb_batch + + heads_block[:, None, None] * stride_cb_head + + offs_t[None, :, None] * stride_cb_t + + offs_window[None, None, :] * stride_cb_j + ) # (H, T, window) + cb_store_mask_3d = (offs_t[None, :, None] < BLOCK_SIZE_T) & ( + offs_window[None, None, :] < BLOCK_SIZE_K + ) # (1, T, window) → broadcasts to (H, T, window) + tl.store(cb_scaled_addrs, rect_CB_scaled_block, mask=cb_store_mask_3d) + + +# Dynamic precompute kernel. Carries constexpr heuristics and dispatches each +# slot to the replay-style or rectangle precompute body. +# Grid: (batch, nheads // HEADS_PER_BLOCK). +@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics( + {"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), MIN_REPLAY_TILE_SIZE)} +) +@triton.heuristics( + { + "BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), MIN_REPLAY_TILE_SIZE + ) + } +) +@triton.jit() +def _dynamic_precompute_kernel( + # Input pointers + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + # Output pointers + cb_scaled_ptr, + decay_vec_ptr, + # Cache pointers + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # dt strides + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + # B strides + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # cb_scaled strides — wrapper allocates (T, window), so stride_cb_t is the padded window size. + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # Meta-parameters + DT_SOFTPLUS: tl.constexpr, + HAS_DT_BIAS: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + HEADS_PER_BLOCK: tl.constexpr, + # Compile-time pick: rectangle (with replay-write fallback) vs replay-only. + RECTANGLE: tl.constexpr, + RECTANGLE_USE_GATHER: tl.constexpr, +): + # Hoisted PDL signal: fire as the first thing every program does. + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() + + pid_b = tl.program_id(axis=0) + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + + pnat_local = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) + needs_write_runtime = pnat_local + T > MAX_REPLAY_BUFFER_LENGTH + # Replay precompute uses a per-slot write predicate; use rectangle only + # for no-write slots when enabled. + if needs_write_runtime or not RECTANGLE: + _replay_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + T, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + needs_write_runtime, + ) + else: + _rectangle_precompute_impl( + dt_ptr, + dt_bias_ptr, + A_ptr, + B_ptr, + C_ptr, + cb_scaled_ptr, + decay_vec_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + cache_buf_idx_ptr, + prev_num_accepted_tokens_ptr, + state_batch_indices_ptr, + T, + MAX_REPLAY_BUFFER_LENGTH, + dstate, + nheads_ngroups_ratio, + stride_dt_batch, + stride_dt_T, + stride_dt_head, + stride_dt_bias_head, + stride_A_head, + stride_B_batch, + stride_B_T, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + DT_SOFTPLUS, + HAS_DT_BIAS, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL, + HEADS_PER_BLOCK, + RECTANGLE_USE_GATHER, + ) + + +# Replay-style main body for one persistent work item. Used by write and +# replay-nowrite paths. + + +@triton.jit() +def _persistent_main_impl( + # Per-work-unit indices (computed by the persistent wrapper). + # `pid_b` indexes the row in decode-batch order. Cache metadata is already + # resolved by the persistent wrapper. + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + prev_num_accepted_tokens, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view, or + # the same `state_ptr` tensor when neither USE_TMA_LOAD_WRITE/NOWRITE nor + # USE_TMA_STORE is enabled (kernel ignores it via constexpr). + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + rand_seed_ptr, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides (double-buffered: cache, dbuf, T, head, dim) + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + WAIT_FOR_PDL_PREDECESSOR: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + # IS_DYNAMIC: when True (persistent_dynamic), is_write is per-slot from + # PNAT. When False (persistent_main), is_write is constexpr from + # WRITE_CHECKPOINT. See also WRITE_CHECKPOINT_IS_CONSTEXPR below. + IS_DYNAMIC: tl.constexpr, + # WRITE_CHECKPOINT_IS_CONSTEXPR: when True, force is_write to the + # WRITE_CHECKPOINT constexpr even in persistent_dynamic mode. The rectangle + # path uses this for the write arm after the outer kernel has already + # narrowed the runtime branch to write slots. When False, + # persistent_dynamic keeps one body with a runtime PNAT check. + WRITE_CHECKPOINT_IS_CONSTEXPR: tl.constexpr = False, + # TMA flags — picked inside body based on is_write. When is_write is + # constexpr (either IS_DYNAMIC=False or WRITE_CHECKPOINT_IS_CONSTEXPR=True), the + # use_tma_load = USE_TMA_LOAD_WRITE if is_write else USE_TMA_LOAD_NOWRITE + # ternary constexpr-folds and only one TMA load form survives. + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # IS_DYNAMIC: kernel-mode label, used by the outer _persistent_main_kernel + # to decide slot-range derivation and outer is_w dispatch strategy + # (constexpr WC for persistent_main; runtime is_w split -> 2 specialized + # impl calls for persistent_dynamic). Inside this impl, IS_DYNAMIC is + # NOT consulted at runtime -- WRITE_CHECKPOINT is the only constexpr that + # gates the write/nowrite codegen, in BOTH modes. + + # Compile-time invariant: QUANT_MAX > 0 must coincide with a quantized + # state dtype (int8 / int16 / float8e4nv) and only those. + tl.static_assert( + (QUANT_MAX > 0.0) + == ( + (state_ptr.dtype.element_ty == tl.int8) + or (state_ptr.dtype.element_ty == tl.int16) + or (state_ptr.dtype.element_ty == tl.float8e4nv) + ), + "QUANT_MAX > 0.0 must coincide with int8 / int16 / float8e4nv state dtype.", + ) + + # Resolve is_write: see WRITE_CHECKPOINT_IS_CONSTEXPR / IS_DYNAMIC docs in the param + # list above. Three cases: + # - WRITE_CHECKPOINT_IS_CONSTEXPR=True (RECT=1 is_w=True arm callers): use WRITE_CHECKPOINT + # constexpr. Caller knows the slot needs write; inner DCEs nowrite + # paths while still constexpr-DCEing the nowrite half. + # - IS_DYNAMIC=True (RECT=0 caller, persistent_dynamic): runtime + # branch on PNAT. Both write and nowrite codegen live in one body. + # - IS_DYNAMIC=False (persistent_main): WRITE_CHECKPOINT constexpr from caller. + if WRITE_CHECKPOINT_IS_CONSTEXPR: + is_write: tl.constexpr = WRITE_CHECKPOINT + elif IS_DYNAMIC: + is_write = (prev_num_accepted_tokens + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_write = WRITE_CHECKPOINT + if is_write: + write_offset = 0 + write_buf = 1 - active_buf + else: + write_offset = prev_num_accepted_tokens + write_buf = active_buf + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_WINDOW) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Load state. state_tma_descriptor is a host-built tensor_descriptor + # over a flat (cache*nheads*dim, dstate) view of state when any TMA + # path is enabled; raw `state_ptr` is the underlying tensor and is + # always passed. state_ptrs / state_ptr_raw are the raw-pointer view + # used for !TMA load and store paths. offs_y is the flat row index + # for TMA load/store; computed unconditionally (cheap int math; DCE'd + # when no TMA path is reachable). + state_mask = m_mask[:, None] & n_mask[None, :] + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state_ptr_raw = state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + state_ptrs = ( + state_ptr_raw + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) + # Load state. Branch on is_write (constexpr = WRITE_CHECKPOINT in BOTH + # modes after the outer-dispatch refactor), then constexpr-pick TMA-vs- + # tl.load per side. Outer `if` DCE's, only the matching side's + # constexpr-gated load survives -- same compile-time picking for both + # persistent_main and persistent_dynamic (the latter dispatches at the + # outer kernel level so each impl instance sees a constexpr WC). + if is_write: + if USE_TMA_LOAD_WRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + else: + if USE_TMA_LOAD_NOWRITE: + state = state_tma_descriptor.load([offs_y, 0]).to(tl.float32) + else: + state = tl.load(state_ptrs, mask=state_mask, other=0.0).to(tl.float32) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + state = state * decode_scale[:, None] + + # Phase 1: Replay via tl.dot fast-forward (reads from active_buf) + group_idx = pid_h // nheads_ngroups_ratio + + history_mask = offs_window < prev_num_accepted_tokens + + old_dt_base = ( + old_dt_ptr + + cache_batch_idx * stride_old_dt_cache + + active_buf * stride_old_dt_dbuf + + pid_h * stride_old_dt_head + ) + history_dt = tl.load( + old_dt_base + offs_window * stride_old_dt_T, + mask=history_mask, + other=0.0, + ).to(tl.float32) + + old_dA_cumsum_base = ( + old_dA_cumsum_ptr + + cache_batch_idx * stride_old_dA_cumsum_cache + + active_buf * stride_old_dA_cumsum_dbuf + + pid_h * stride_old_dA_cumsum_head + ) + history_dA_cumsum = tl.load( + old_dA_cumsum_base + offs_window * stride_old_dA_cumsum_T, + mask=history_mask, + other=0.0, + ).to(tl.float32) + + prev_k_idx = tl.minimum( + tl.maximum(prev_num_accepted_tokens - 1, 0), MAX_REPLAY_BUFFER_LENGTH - 1 + ) + total_dA_cumsum = tl.load(old_dA_cumsum_base + prev_k_idx * stride_old_dA_cumsum_T).to( + tl.float32 + ) + + coeff = tl.exp(total_dA_cumsum - history_dA_cumsum) * history_dt + + # Double buffering keeps old_x replay reads and checkpoint writes in + # disjoint buffers on write steps. + old_x_read_base = ( + old_x_ptr + + cache_batch_idx * stride_old_x_cache + + active_buf * stride_old_x_dbuf + + pid_h * stride_old_x_head + ) + old_x_write_base = ( + old_x_ptr + + cache_batch_idx * stride_old_x_cache + + write_buf * stride_old_x_dbuf + + pid_h * stride_old_x_head + ) + history_x = tl.load( + old_x_read_base + + offs_window[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=history_mask[:, None] & m_mask[None, :], + other=0.0, + ) + + old_B_base = ( + old_B_ptr + + cache_batch_idx * stride_old_B_cache + + active_buf * stride_old_B_dbuf + + group_idx * stride_old_B_group + ) + history_B = tl.load( + old_B_base + offs_window[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, + mask=history_mask[:, None] & n_mask[None, :], + other=0.0, + ).to(tl.float32) + + dB_scaled = coeff[:, None] * history_B + + total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) + state *= total_decay + + state += tl.dot(tl.trans(history_x).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) + + if is_write: + if USE_RS_ROUNDING: + # Generate random tensor for stochastic rounding. The amount of + # randomness needed depends on the SR codegen path: + # fp16 SR (cvt.rs.f16x2): 1 b32 per 2 outputs (pack=2) + # fp8 SR (cvt.rs.satfinite.e4m3x4): 1 b32 per 4 outputs (pack=4) + # int8 SR (16b chunks + bitrev16): 1 b32 per 4 outputs + # int16 SR (24b + bitrev32): 1 b32 per 2 outputs + # The PTX cvt.rs.* instructions consume a single 32-bit random + # and split the bits internally for 2 or 4 conversions. Generate + # only what's actually consumed and broadcast to fill the unused + # slots — saves Philox rounds proportionally. + if QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.float8e4nv: + RAND_DIVISOR: tl.constexpr = 4 # fp8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int8: + RAND_DIVISOR: tl.constexpr = 4 # int8 SR + elif QUANT_MAX > 0.0 and state_ptrs.dtype.element_ty == tl.int16: + RAND_DIVISOR: tl.constexpr = 2 # int16 SR + elif QUANT_MAX == 0.0: + RAND_DIVISOR: tl.constexpr = 2 # fp16 SR (only fp16 supported here) + else: + RAND_DIVISOR: tl.constexpr = 1 # unreachable; keeps constexpr initialized + + rand_seed = tl.load(rand_seed_ptr) + base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head + # Number of unique randoms per row = dstate / RAND_DIVISOR. + # randint4x emits 4 randoms per offset, so use that / 4 offsets. + offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // (4 * RAND_DIVISOR)) + rand_offsets_q = ( + base_rand + + offs_m[:, None] * stride_state_dim + + offs_n_q[None, :] * (stride_state_dstate * 4 * RAND_DIVISOR) + ) # (M, dstate / (4*RAND_DIVISOR)) + if PHILOX_ROUNDS > 0: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + else: + r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) + r01 = tl.join(r0, r1) + r23 = tl.join(r2, r3) + r0123 = tl.join(r01, r23) + rand_compact = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR)) + # Broadcast each unique rand to RAND_DIVISOR adjacent positions. + # Pack-group (pack=2 fp16 / pack=4 fp8) consumes adjacent positions; + # the unique rand lands at the asm's read slot; duplicates feed + # the dead slots. Triton's broadcast_to is stride-0 in IR. + if RAND_DIVISOR > 1: + rand_3d = rand_compact[:, :, None] + rand_3d = tl.broadcast_to( + rand_3d, + (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE // RAND_DIVISOR, RAND_DIVISOR), + ) + rand = tl.reshape(rand_3d, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) + else: + rand = rand_compact + + if QUANT_MAX > 0.0: + amax = tl.max(tl.abs(state), axis=1) + encode_scale = tl.where(amax == 0.0, 1.0, QUANT_MAX / amax) + decode_scale = 1.0 / encode_scale + state_scales_ptrs = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + + offs_m * stride_state_scales_dim + ) + tl.store(state_scales_ptrs, decode_scale, mask=m_mask) + state_q = state * encode_scale[:, None] + if USE_RS_ROUNDING and (state_ptrs.dtype.element_ty == tl.float8e4nv): + _state_q_fp8sr = _stochastic_round_fp8x4_e4m3(state_q, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_fp8sr) + else: + tl.store(state_ptrs, _state_q_fp8sr, mask=state_mask) + else: + if USE_RS_ROUNDING: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized SR fall-through expects int8 or int16; " + "fp8 SR is handled by the prior branch.", + ) + if state_ptrs.dtype.element_ty == tl.int8: + state_q = _stochastic_round_int8_packed(state_q, rand, offs_n[None, :]) + else: + state_q = _stochastic_round_int16_packed(state_q, rand, offs_n[None, :]) + elif state_ptrs.dtype.element_ty != tl.float8e4nv: + tl.static_assert( + (state_ptrs.dtype.element_ty == tl.int8) + or (state_ptrs.dtype.element_ty == tl.int16), + "Quantized RN with explicit round() expects int8 or int16.", + ) + state_q = tl.extra.cuda.libdevice.round(state_q) + state_q = tl.minimum(tl.maximum(state_q, -QUANT_MAX), QUANT_MAX) + _state_q_cast = state_q.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_q_cast) + else: + tl.store(state_ptrs, _state_q_cast, mask=state_mask) + elif USE_RS_ROUNDING: + tl.static_assert( + state_ptrs.dtype.element_ty == tl.float16, + "Non-quantized SR only supports fp16 state.", + ) + _state_sr = _stochastic_round_fp16x2(state, rand) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_sr) + else: + tl.store(state_ptrs, _state_sr, mask=state_mask) + else: + _state_cast = state.to(state_ptrs.dtype.element_ty) + if USE_TMA_STORE: + state_tma_descriptor.store([offs_y, 0], _state_cast) + else: + tl.store(state_ptrs, _state_cast, mask=state_mask) + + # Phase 2: Output using precomputed CB_scaled and decay_vec + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + if WAIT_FOR_PDL_PREDECESSOR: + _gdc_wait_with_memory_clobber() + + C_tile = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + + step_x = tl.load( + x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_write_base + + (write_offset + offs_t)[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + step_x, + mask=t_mask[:, None] & m_mask[None, :], + ) + step_x = step_x.to(tl.float32) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + init_out = tl.dot(C_tile.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] + cb_out = tl.dot(CB_scaled.to(tl.bfloat16), step_x.to(tl.bfloat16)) + output_tile = init_out + cb_out + + if HAS_D: + output_tile = output_tile + step_x * D[None, :] + + if HAS_Z: + z_tile = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + gated_output_tile = output_tile * z_tile * tl.sigmoid(z_tile) + output_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(output_ptrs, gated_output_tile, mask=t_mask[:, None] & m_mask[None, :]) + else: + output_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(output_ptrs, output_tile, mask=t_mask[:, None] & m_mask[None, :]) + + +# Rectangle nowrite main body for one persistent work item. Used only for +# nowrite slots when the persistent kernel runs with RECTANGLE=True. +@triton.jit() +def _persistent_rectangle_impl( + # Per-work-unit indices (computed by the persistent wrapper). + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + prev_num_accepted_tokens, + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor (same flat 2D view as + # replay path). Used when USE_TMA_LOAD; ignored otherwise. + state_tma_descriptor, + state_scales_ptr, # only consulted when QUANT_MAX > 0 + old_x_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides (double-buffered: cache, dbuf, T, head, dim) + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides (rectangle: batch, nheads, T, window) + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + WAIT_FOR_PDL_PREDECESSOR: tl.constexpr, + QUANT_MAX: tl.constexpr, + USE_TMA_LOAD: tl.constexpr = False, +): + # Nowrite-only: step tokens append at [PNAT, PNAT+T). + + # Rectangle window layout: history at [0, PNAT), step tokens at [PNAT, PNAT+T). + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + offs_t = tl.arange(0, BLOCK_SIZE_T) + offs_window = tl.arange(0, BLOCK_SIZE_K) + m_mask = offs_m < dim + n_mask = offs_n < dstate + t_mask = offs_t < T + + # Window masks use the same PNAT-runtime offset as rectangle precompute. + is_history_position = offs_window < prev_num_accepted_tokens + safe_history_idx = tl.where(is_history_position, offs_window, 0) + step_idx_from_window = offs_window - prev_num_accepted_tokens + is_step_position = (step_idx_from_window >= 0) & (step_idx_from_window < T) + safe_step_idx = tl.where(is_step_position, step_idx_from_window, 0) + + # Load state. Quant scale hoist: defer `* decode_scale` post-matmul. + if USE_TMA_LOAD: + offs_y = ( + cache_batch_idx.to(tl.int32) * (stride_state_batch // stride_state_dim).to(tl.int32) + + pid_h * dim + + pid_m * BLOCK_SIZE_M + ) + state = state_tma_descriptor.load([offs_y, 0]) + else: + state_ptr_local = ( + state_ptr + cache_batch_idx * stride_state_batch + pid_h * stride_state_head + ) + state_ptrs = ( + state_ptr_local + + offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate + ) + state_mask = m_mask[:, None] & n_mask[None, :] + state = tl.load(state_ptrs, mask=state_mask, other=0.0) + if QUANT_MAX > 0.0: + state_scales_base = ( + state_scales_ptr + + cache_batch_idx * stride_state_scales_cache + + pid_h * stride_state_scales_head + ) + decode_scale = tl.load( + state_scales_base + offs_m * stride_state_scales_dim, + mask=m_mask, + other=1.0, + ).to(tl.float32) + else: + state = state.to(tl.float32) + + # Group / pointer offset setup + group_idx = pid_h // nheads_ngroups_ratio + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + # Rectangle path: nowrite-only. write_buf == active_buf (no flip), so the + # read and write paths target the same dbuf slot. No race: write_offset= + # PNAT puts new tokens at [PNAT, PNAT+T), disjoint from the read range + # [0, PNAT). + old_x_read_base = ( + old_x_ptr + + cache_batch_idx * stride_old_x_cache + + active_buf * stride_old_x_dbuf + + pid_h * stride_old_x_head + ) + old_x_write_base = old_x_read_base + + if HAS_D: + D = tl.load( + D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 + ).to(tl.float32) + + # History x does not depend on conv1d/precompute. + history_x = tl.load( + old_x_read_base + + safe_history_idx[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + mask=is_history_position[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + + if WAIT_FOR_PDL_PREDECESSOR: + _gdc_wait_with_memory_clobber() + + C_tile = tl.load( + C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, + mask=t_mask[:, None] & n_mask[None, :], + other=0.0, + ) + step_x_in_window = tl.load( + x_ptr + safe_step_idx[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, + mask=is_step_position[:, None] & m_mask[None, :], + other=0.0, + ) + tl.store( + old_x_write_base + + offs_window[:, None] * stride_old_x_T + + offs_m[None, :] * stride_old_x_dim, + step_x_in_window, + mask=is_step_position[:, None] & m_mask[None, :], + ) + + step_x_in_window_for_dot = step_x_in_window.to(tl.bfloat16) + x_window_for_dot = history_x + step_x_in_window.to(tl.float32) + + if HAS_D: + step_in_window_selector = offs_t[:, None] == ( + offs_window[None, :] - prev_num_accepted_tokens + ) + step_x = tl.dot(step_in_window_selector.to(tl.bfloat16), step_x_in_window_for_dot) + + cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head + CB_scaled = tl.load( + cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_window[None, :] * stride_cb_j, + mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_window[None, :] < BLOCK_SIZE_K), + other=0.0, + ).to(tl.float32) + + decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head + decay_vec_full = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( + tl.float32 + ) + + state_out = ( + tl.dot(C_tile.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec_full[:, None] + ) + if QUANT_MAX > 0.0: + state_out = state_out * decode_scale[None, :] + + token_out = tl.dot(CB_scaled.to(tl.bfloat16), x_window_for_dot.to(tl.bfloat16)) + + output_tile = state_out + token_out + + if HAS_D: + output_tile = output_tile + step_x * D[None, :] + + if HAS_Z: + z_tile = tl.load( + z_ptr + offs_t[:, None] * stride_z_T + offs_m[None, :] * stride_z_dim, + mask=t_mask[:, None] & m_mask[None, :], + other=0.0, + ).to(tl.float32) + gated_output_tile = output_tile * z_tile * tl.sigmoid(z_tile) + output_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(output_ptrs, gated_output_tile, mask=t_mask[:, None] & m_mask[None, :]) + else: + output_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim + tl.store(output_ptrs, output_tile, mask=t_mask[:, None] & m_mask[None, :]) + + +# Persistent replay main kernel: 1D grid, persistent CTA loop. Heuristics cover +# both the replay-style and rectangle main bodies. +@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None}) +@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None}) +@triton.heuristics({"USE_RS_ROUNDING": lambda args: args["rand_seed_ptr"] is not None}) +@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])}) +@triton.heuristics( + {"BLOCK_SIZE_T": lambda args: max(triton.next_power_of_2(args["T"]), MIN_REPLAY_TILE_SIZE)} +) +@triton.heuristics( + { + "BLOCK_SIZE_WINDOW": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), MIN_REPLAY_TILE_SIZE + ) + } +) +@triton.heuristics( + { + "BLOCK_SIZE_K": lambda args: max( + triton.next_power_of_2(args["MAX_REPLAY_BUFFER_LENGTH"]), MIN_REPLAY_TILE_SIZE + ) + } +) +@triton.heuristics( + {"NUM_PID_M_BLOCKS": lambda args: triton.cdiv(args["dim"], args["BLOCK_SIZE_M"])} +) +@triton.jit() +def _persistent_main_kernel( + # Pointers + state_ptr, + # state_tma_descriptor: TMA tensor_descriptor over state's flat 2D view. + # Shared across BOTH the replay path (consumed by _persistent_main_impl + # when USE_TMA_LOAD_*/STORE) AND the rectangle path (consumed by + # _persistent_rectangle_impl when USE_TMA_LOAD) — same descriptor, same + # block_shape, just gated by separate constexprs per impl. Wrapper sets + # this to a TensorDescriptor when ANY of the three TMA flags is on, else + # to `state_ptr` (raw); each impl ignores it via its own constexpr when + # not consuming it. + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + prev_num_accepted_tokens_ptr, + cache_buf_idx_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + state_batch_indices_ptr, + replay_work_items_ptr, + rand_seed_ptr, + # Persistent-loop work-distribution scalars. Caller pre-sorts the batch + # write-first; the kernel uses (n_writes, batch_total, WRITE_CHECKPOINT) + # to derive its own slot range. Write half processes [0, n_writes), + # nowrite half processes [n_writes, batch_total). + # + # n_writes_ptr is a device pointer to a (1,) int32 tensor. Reading from + # device memory keeps the pointer stable across CUDA graph replay while + # allowing the value to change between iterations. + # When IS_DYNAMIC=True the value is unused (Triton DCEs the load). + n_writes_ptr, # device-side count of write-mode slots + batch_total, # total slot count + nheads, # total head count + # Dimensions + T: tl.constexpr, + MAX_REPLAY_BUFFER_LENGTH: tl.constexpr, + dim: tl.constexpr, + dstate: tl.constexpr, + nheads_ngroups_ratio: tl.constexpr, + # state strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + # state_scales strides + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + # old_x strides (double-buffered: cache, dbuf, T, head, dim) + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + # old_B strides + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + # old_dt strides + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + # old_dA_cumsum strides + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + # x strides + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + # C strides + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + # D strides + stride_D_head, + stride_D_dim, + # z strides + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + # out strides + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + # cb_scaled strides + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + # decay_vec strides + stride_dv_batch, + stride_dv_head, + stride_dv_t, + # Meta + BLOCK_SIZE_M: tl.constexpr, + HAS_D: tl.constexpr, + HAS_Z: tl.constexpr, + BLOCK_SIZE_DSTATE: tl.constexpr, + BLOCK_SIZE_T: tl.constexpr, + BLOCK_SIZE_WINDOW: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, + USE_RS_ROUNDING: tl.constexpr, + PHILOX_ROUNDS: tl.constexpr, + QUANT_MAX: tl.constexpr, + WRITE_CHECKPOINT: tl.constexpr, + LAUNCH_DEPENDENT_KERNELS: tl.constexpr, + # NUM_PERSISTENT is a runtime loop stride, so CTA-per-SM tuning can vary + # without changing the compiled kernel signature. Work decomposition uses + # constexpr NUM_PID_M_BLOCKS and runtime n_slots_local. + NUM_PERSISTENT, + NUM_LOOP_STAGES: tl.constexpr, + NUM_PID_M_BLOCKS: tl.constexpr, + FLATTEN: tl.constexpr, + WARP_SPECIALIZE: tl.constexpr, + IS_DYNAMIC: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr = MIN_REPLAY_TILE_SIZE, # rectangle window dimension + RECTANGLE: tl.constexpr = False, # dispatch nowrite slots to _persistent_rectangle_impl when true + # 3 TMA toggles per the 3 live paths per-compilation: + # USE_TMA_LOAD_WRITE — SSM state load when is_write + # USE_TMA_LOAD_NOWRITE — nowrite-path state load (rect when RECTANGLE, + # else replay-nowrite) + # USE_TMA_STORE — SSM state store (only fires on write + # path; no-op when not is_write) + # Wrapper picks USE_TMA_LOAD_NOWRITE = _use_tma_rect_load (if rectangle) + # or _use_tma_replay_nowrite_load (if not). + USE_TMA_LOAD_WRITE: tl.constexpr = False, + USE_TMA_LOAD_NOWRITE: tl.constexpr = False, + USE_TMA_STORE: tl.constexpr = False, +): + # PDL signal: fire once at kernel entry (not per work unit). + if LAUNCH_DEPENDENT_KERNELS: + tl.extra.cuda.gdc_launch_dependents() - # Zero stale rows beyond PNAT to prevent Inf/NaN from reaching tl.dot. - old_x_base = old_x_ptr + cache_batch_idx * stride_old_x_cache + pid_h * stride_old_x_head - old_x_all = tl.load( - old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - mask=accepted_mask[:, None] & m_mask[None, :], - other=0.0, - ) + # Load runtime n_writes from device memory. Read once at kernel entry; + # used only by the !IS_DYNAMIC slot-range derivation below. Triton + # DCEs the load when IS_DYNAMIC=True (n_writes is dead there). + n_writes = tl.load(n_writes_ptr) - # Apply the same accepted-row mask to old_B. - old_B_base = ( - old_B_ptr - + cache_batch_idx * stride_old_B_cache - + buf_read * stride_old_B_dbuf - + group_idx * stride_old_B_group - ) - old_B_all = tl.load( - old_B_base + offs_t[:, None] * stride_old_B_T + offs_n[None, :] * stride_old_B_dstate, - mask=accepted_mask[:, None] & n_mask[None, :], - other=0.0, - ).to(tl.float32) + # Derive this kernel's slot range. Two modes: + # IS_DYNAMIC=False (persistent_main): caller pre-sorts and splits halves; + # slot range is [0, n_writes) when WRITE_CHECKPOINT else [n_writes, batch_total) + # IS_DYNAMIC=True (persistent_dynamic): single launch covers full batch; + # each work-item dispatches via runtime PNAT check inside the impl. + if IS_DYNAMIC: + slot_lo = 0 + slot_hi = batch_total + else: + if WRITE_CHECKPOINT: + slot_lo = 0 + slot_hi = n_writes + else: + slot_lo = n_writes + slot_hi = batch_total + n_slots_local = slot_hi - slot_lo - # Scale B by coefficients - dB_scaled = coeff[:, None] * old_B_all + pid = tl.program_id(axis=0) + total_work = n_slots_local * NUM_PID_M_BLOCKS * nheads - # Apply total decay to initial state FIRST, then add contributions - total_decay = tl.where(prev_num_accepted_tokens > 0, tl.exp(total_dA_cumsum), 1.0) - state *= total_decay + if LAUNCH_WITH_PDL and ( + (total_work == 0 and pid == 0) or (NUM_LOOP_STAGES > 1 and pid < total_work) + ): + # This pre-loop path reads only replay partition metadata prepared + # before the PDL chain, not conv1d/precompute outputs. + # + # Empty PDL launches still need one waiting CTA, otherwise this kernel + # can retire before its upstream dependency and a later dependent launch + # can observe producer data too early. + # + # For loop-pipelined kernels, CTAs with work wait before the loop to + # work around a Triton PDL scheduling bug that can otherwise move + # producer-dependent loads ahead of the wait. + _gdc_wait_with_memory_clobber() - # tl.dot fast-forward: old_x^T @ dB_scaled -> (M, dstate) - state += tl.dot(tl.trans(old_x_all).to(tl.bfloat16), dB_scaled.to(tl.bfloat16)) - - # Write post-replay state - if USE_RS_ROUNDING: - # Stochastic rounding for fp16 state using Philox-4x32 PRNG. - # Each Philox call produces 4 random ints. We call randint4x on - # quarter-sized dstate offsets and join+reshape to get the full - # (M, dstate) random tensor — 4x fewer PRNG rounds. - rand_seed = tl.load(rand_seed_ptr + cache_batch_idx) - base_rand = cache_batch_idx * stride_state_batch + pid_h * stride_state_head - offs_n_q = tl.arange(0, BLOCK_SIZE_DSTATE // 4) - rand_offsets_q = ( - base_rand - + offs_m[:, None] * stride_state_dim - + offs_n_q[None, :] * (stride_state_dstate * 4) - ) # (M, dstate//4) - if PHILOX_ROUNDS > 0: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q, PHILOX_ROUNDS) + # Persistent loop. Decompose tile_id into (pid_h, pid_b_local, pid_m) + # with pid_m varying fastest (M-tile cache locality on state load), then + # slot, then head — mirrors the existing 3D grid's axis ordering + # (axis=0 fastest = pid_m). + for tile_id in tl.range( + pid, + total_work, + NUM_PERSISTENT, + flatten=FLATTEN, + num_stages=NUM_LOOP_STAGES, + warp_specialize=WARP_SPECIALIZE, + ): + pid_m = tile_id % NUM_PID_M_BLOCKS + pid_b_local = (tile_id // NUM_PID_M_BLOCKS) % n_slots_local + pid_h = tile_id // (NUM_PID_M_BLOCKS * n_slots_local) + work_item_idx = pid_b_local + slot_lo + if IS_DYNAMIC: + pid_b = work_item_idx + cache_batch_idx = tl.load(state_batch_indices_ptr + pid_b).to(tl.int64) + active_buf = tl.load(cache_buf_idx_ptr + cache_batch_idx).to(tl.int32) + pnat = tl.load(prev_num_accepted_tokens_ptr + cache_batch_idx) else: - r0, r1, r2, r3 = tl.randint4x(rand_seed, rand_offsets_q) - # Interleave 4 quarter-sized tensors → full (M, dstate) random tensor - r01 = tl.join(r0, r1) # (M, dstate//4, 2) - r23 = tl.join(r2, r3) # (M, dstate//4, 2) - r0123 = tl.join(r01, r23) # (M, dstate//4, 2, 2) - rand = tl.reshape(r0123, (BLOCK_SIZE_M, BLOCK_SIZE_DSTATE)) - tl.store(state_ptrs, _stochastic_round_fp16x2(state, rand), mask=state_mask) - else: - tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=state_mask) - - # Phase 2: Output using precomputed CB_scaled and decay_vec - x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head - C_ptr += pid_b * stride_C_batch + group_idx * stride_C_group - if HAS_Z: - z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head - out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head - - if HAS_D: - D = tl.load( - D_ptr + pid_h * stride_D_head + offs_m * stride_D_dim, mask=m_mask, other=0.0 - ).to(tl.float32) + work_item_base = replay_work_items_ptr + work_item_idx * _REPLAY_WORK_ITEM_WIDTH + pid_b = tl.load(work_item_base + _REPLAY_WORK_POSITION_IN_DECODE_BATCH) + cache_batch_idx = tl.load(work_item_base + _REPLAY_WORK_CACHE_SLOT).to(tl.int64) + pnat = tl.load(work_item_base + _REPLAY_WORK_PNAT) + active_buf = tl.load(work_item_base + _REPLAY_WORK_CACHE_BUF_IDX).to(tl.int32) + # Dispatch: when RECTANGLE is set, send nowrite slots to the rectangle + # impl. `replay_work_items` carries the cache slot, PNAT and active + # buffer for persistent_main; persistent_dynamic resolves those once + # here from the existing tensors. + if RECTANGLE: + if IS_DYNAMIC: + is_w = (pnat + T) > MAX_REPLAY_BUFFER_LENGTH + else: + is_w = WRITE_CHECKPOINT + if is_w: + # Pass WRITE_CHECKPOINT=True constexpr to specialize this + # impl call for the write path. Under IS_DYNAMIC=True, the + # kernel-level WRITE_CHECKPOINT is False (launcher default), + # but the OUTER is_w branch we are inside narrows the + # runtime path to writes-only, so we override to True here + # so the impl's constexpr-gated `if is_write:` blocks DCE + # to the write-only codegen. Under IS_DYNAMIC=False + # (persistent_main), the kernel-level WRITE_CHECKPOINT is + # itself True for this half (write half launches with + # WRITE_CHECKPOINT=True), and the outer is_w = WRITE_CHECKPOINT = True + # constexpr-folds; passing literal True here is consistent + # and constexpr-equivalent. + _persistent_main_impl( + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + pnat, + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + rand_seed_ptr, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + True, # WRITE_CHECKPOINT=True (write arm) + IS_DYNAMIC, + True, # WRITE_CHECKPOINT_IS_CONSTEXPR + USE_TMA_LOAD_WRITE, + USE_TMA_LOAD_NOWRITE, + USE_TMA_STORE, + ) + else: + _persistent_rectangle_impl( + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + pnat, + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_K, + LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR + QUANT_MAX, + USE_TMA_LOAD_NOWRITE, # USE_TMA_LOAD + ) + else: + _persistent_main_impl( + pid_m, + pid_b, + pid_h, + cache_batch_idx, + active_buf, + pnat, + state_ptr, + state_tma_descriptor, + state_scales_ptr, + old_x_ptr, + old_B_ptr, + old_dt_ptr, + old_dA_cumsum_ptr, + x_ptr, + C_ptr, + D_ptr, + z_ptr, + out_ptr, + cb_scaled_ptr, + decay_vec_ptr, + rand_seed_ptr, + T, + MAX_REPLAY_BUFFER_LENGTH, + dim, + dstate, + nheads_ngroups_ratio, + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_state_scales_cache, + stride_state_scales_head, + stride_state_scales_dim, + stride_old_x_cache, + stride_old_x_dbuf, + stride_old_x_T, + stride_old_x_head, + stride_old_x_dim, + stride_old_B_cache, + stride_old_B_dbuf, + stride_old_B_T, + stride_old_B_group, + stride_old_B_dstate, + stride_old_dt_cache, + stride_old_dt_dbuf, + stride_old_dt_head, + stride_old_dt_T, + stride_old_dA_cumsum_cache, + stride_old_dA_cumsum_dbuf, + stride_old_dA_cumsum_head, + stride_old_dA_cumsum_T, + stride_x_batch, + stride_x_T, + stride_x_head, + stride_x_dim, + stride_C_batch, + stride_C_T, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_T, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_T, + stride_out_head, + stride_out_dim, + stride_cb_batch, + stride_cb_head, + stride_cb_t, + stride_cb_j, + stride_dv_batch, + stride_dv_head, + stride_dv_t, + BLOCK_SIZE_M, + HAS_D, + HAS_Z, + BLOCK_SIZE_DSTATE, + BLOCK_SIZE_T, + BLOCK_SIZE_WINDOW, + LAUNCH_WITH_PDL and (NUM_LOOP_STAGES == 1), # WAIT_FOR_PDL_PREDECESSOR + USE_RS_ROUNDING, + PHILOX_ROUNDS, + QUANT_MAX, + WRITE_CHECKPOINT, + IS_DYNAMIC, + False, # WRITE_CHECKPOINT_IS_CONSTEXPR + USE_TMA_LOAD_WRITE, + USE_TMA_LOAD_NOWRITE, + USE_TMA_STORE, + ) - # Wait for precompute kernel (PDL) before reading its outputs. - # With chained PDL (conv1d → precompute → main), gdc_wait() ensures - # precompute has completed — which transitively ensures conv1d has - # completed (precompute waited on conv1d via its own gdc_wait). - # All loads below (x, C from conv1d; CB_scaled, decay_vec from precompute) - # are safe after this point. - if LAUNCH_WITH_PDL: - tl.extra.cuda.gdc_wait() - # Load conv1d outputs: C_all and x_all - C_all = tl.load( - C_ptr + offs_t[:, None] * stride_C_T + offs_n[None, :] * stride_C_dstate, - mask=t_mask[:, None] & n_mask[None, :], - other=0.0, - ) +# ============================================================================ +# Python wrapper +# ============================================================================ - x_all = tl.load( - x_ptr + offs_t[:, None] * stride_x_T + offs_m[None, :] * stride_x_dim, - mask=t_mask[:, None] & m_mask[None, :], - other=0.0, - ) - # Store new x to cache (single-buffered; replay already read the old data) - tl.store( - old_x_base + offs_t[:, None] * stride_old_x_T + offs_m[None, :] * stride_old_x_dim, - x_all, - mask=t_mask[:, None] & m_mask[None, :], - ) - x_all = x_all.to(tl.float32) - # Load precomputed CB_scaled and decay_vec - cb_scaled_base = cb_scaled_ptr + pid_b * stride_cb_batch + pid_h * stride_cb_head - CB_scaled = tl.load( - cb_scaled_base + offs_t[:, None] * stride_cb_t + offs_t[None, :] * stride_cb_j, - mask=(offs_t[:, None] < BLOCK_SIZE_T) & (offs_t[None, :] < BLOCK_SIZE_T), - other=0.0, - ).to(tl.float32) +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} - decay_vec_base = decay_vec_ptr + pid_b * stride_dv_batch + pid_h * stride_dv_head - decay_vec = tl.load(decay_vec_base + offs_t * stride_dv_t, mask=t_mask, other=0.0).to( - tl.float32 - ) - # init_out = C_all @ state^T * decay_vec - init_out = tl.dot(C_all.to(tl.bfloat16), tl.trans(state).to(tl.bfloat16)) * decay_vec[:, None] +# --------------------------------------------------------------------------- +# Default tunings — looked up by (effective_batch, dtype, sr) when the caller +# leaves mode/knobs as None. +# +# Effective batch = raw_batch × nheads_per_rank. Our sweep was at TP=8 with +# the standard Mamba2 nheads; at call time we compute it from the input +# tensor shape so callers at other TP / nheads pick up the right cell. +# +# Schema: dict[(dtype_str, sr_str)] → list[(eff_batch_threshold, mode, knobs)] +# sorted by threshold ascending. Lookup finds the first threshold ≥ eff_b +# (so missing intermediate batches fall up to the next tuned cell). If +# eff_b exceeds the largest threshold, use the largest entry. +# +# Each `knobs` dict only contains keys for the chosen mode; the wrapper +# unpacks them with the same name as the matching kwargs. Caller-provided +# kwargs always win over table values. +# +# This table is intentionally NOT parameterized by T or max_window. Our +# sweep was T=6, max_window=16. Callers outside that regime silently get +# the same numbers — they may be suboptimal but they're correct. +# +# Source: emit_tuning_from_noise.py. Auto-generated from noise-cleaned per-cell search +# winners (best of pd / pm by bucket_expected_renorm). Effective batch = raw_batch × 16. +# The search predates a Triton PDL scheduling bug. +# Most knobs are unchanged; large PDL-hoist/precompute regressions got spot retunes. +# Missing dtype/SR combos fall back via the _resolve_tuning chain: +# RN→SR for same dtype, then bf16/int16→fp16/SR and fp8→int8/SR. +_DEFAULT_TUNING: dict[tuple[str, str], list[tuple[int, str, dict]]] = { + ("fp16", "SR"): [ + ( + 16, + "persistent_dynamic", + { + "_block_size_m": 4, + "_cta_per_sm": 4, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=1, score=6.40us (B200 precompute-retune noise-cleaned 5x500) + ( + 32, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 4, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=2, score=6.99us (B200 precompute-retune noise-cleaned 5x500) + ( + 64, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 10, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=4, score=7.05us (B200 precompute-retune noise-cleaned 5x500) + ( + 128, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=8, score=7.76us (B200 PDL-hoist default 5x200) + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=16, score=9.27us (B200 PDL-hoist default 5x200) + ( + 512, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 2, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=32, score=13.08us (B200 precompute-retune noise-cleaned 5x500) + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 9, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=64, score=17.51us (B200 PDL-hoist default 5x200) + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=26.39us (B200 PDL-hoist default 5x200) + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=43.57us (B200 PDL-hoist default 5x200) + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=73.8us (B200 PDL-hoist default 5x200) + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=134.03us (B200 PDL-hoist default 5x200) + ], + ("int8", "SR"): [ + ( + 16, + "persistent_dynamic", + { + "_block_size_m": 4, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=1, score=6.54us (B200 precompute-retune noise-cleaned 5x500) + ( + 32, + "persistent_dynamic", + { + "_block_size_m": 4, + "_cta_per_sm": 10, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=2, score=7.44us (B200 PDL-hoist default 5x200) + ( + 64, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 5, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 2, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=4, score=7.97us (B200 PDL-hoist default 5x200) + ( + 128, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages": 1, + "_num_stages": 1, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=8, score=8.15us (B200 PDL-hoist default 5x200) + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 10, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages": 1, + "_num_stages": 1, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=16, score=9.88us (B200 PDL-hoist default 5x200) + ( + 512, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 16, + "_cta_per_sm_nowrite": 3, + "_cta_per_sm_write": 9, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=32, score=12.99us (B200 PDL-hoist default 5x200) + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 3, + "_cta_per_sm_write": 10, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 5, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=64, score=18.51us (B200 PDL-retune noise-cleaned 5x500) + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 5, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 2, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=28.83us (B200 precompute-retune noise-cleaned 5x500) + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 3, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 4, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=46.60us (B200 precompute-retune noise-cleaned 5x500) + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 6, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 5, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 2, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, + "_num_warps_write": 4, + "_precompute_num_warps": 4, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=81.07us (B200 precompute-retune noise-cleaned 5x500) + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 64, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 3, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 2, + "_num_stages_nowrite": 4, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 4, + "_precompute_num_warps": 1, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=149.25us (B200 PDL-hoist default 5x200) + ], + ("fp8", "SR"): [ + # --- ALL UNCHANGED (not tuned this round) --- + ( + 16, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 3, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 5, + "_num_warps": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=1, score=6.46us (B200 PDL-retune noise-cleaned 5x500) + ( + 32, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 10, + "_flatten": False, + "_heads_per_block": 1, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=2, score=6.72us (B200 precompute retune 5x100) + ( + 64, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 9, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=4, score=7.24us (B200 PDL-hoist default 5x200) + ( + 128, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 5, + "_flatten": False, + "_heads_per_block": 4, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=8, score=7.79us (B200 PDL-hoist default 5x200) + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 16, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=16, score=8.77us (B200 PDL-hoist default 5x200) + ( + 512, + "persistent_dynamic", + { + "_block_size_m": 32, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages": 1, + "_num_stages": 3, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=32, score=10.71us (B200 PDL-hoist default 5x200) + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 7, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 2, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=64, score=16.0us (B200 PDL-hoist default 5x200) + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 5, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=24.5us (B200 PDL-retune noise-cleaned 5x500) + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 10, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 5, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=40.04us (B200 PDL-retune noise-cleaned 5x500) + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 5, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 1, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": False, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=67.94us (B200 precompute retune 5x100) + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 6, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 5, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 2, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=126.74us (B200 PDL-retune noise-cleaned 5x500) + ], + ("fp32", "RN"): [ + ( + 16, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 6, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 1, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=1, score=6.42us (B200 precompute-retune noise-cleaned 5x500) + ( + 32, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 2, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 5, + "_num_warps": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=2, score=6.83us (B200 PDL-hoist default 5x200) + ( + 64, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 1, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=4, score=7.10us (B200 precompute-retune noise-cleaned 5x500) + ( + 128, + "persistent_dynamic", + { + "_block_size_m": 8, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 2, + "_num_loop_stages": 1, + "_num_stages": 4, + "_num_warps": 1, + "_precompute_num_warps": 2, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=8, score=8.14us (B200 PDL-hoist default 5x200) + ( + 256, + "persistent_dynamic", + { + "_block_size_m": 32, + "_cta_per_sm": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": False, + "_use_tma_replay_write_store": False, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=16, score=9.66us (B200 PDL-hoist default 5x200) + ( + 512, + "persistent_dynamic", + { + "_block_size_m": 32, + "_cta_per_sm": 7, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages": 1, + "_num_stages": 2, + "_num_warps": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=32, score=13.59us (B200 PDL-hoist default 5x200) + ( + 1024, + "persistent_main", + { + "_block_size_m_nowrite": 64, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 6, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 1, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 1, + "_num_stages_write": 1, + "_num_warps_nowrite": 2, + "_num_warps_write": 2, + "_precompute_num_warps": 8, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": True, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": False, + }, + ), # raw_batch=64, score=20.05us (B200 PDL-retune noise-cleaned 5x500) + ( + 2048, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 3, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=128, score=31.39us (B200 PDL-retune noise-cleaned 5x500) + ( + 4096, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 4, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 2, + "_num_stages_write": 3, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 8, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=256, score=51.27us (B200 PDL-retune noise-cleaned 5x500) + ( + 8192, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 4, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 16, + "_num_loop_stages_nowrite": 5, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 4, + "_num_stages_write": 4, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 4, + "_use_tma_rect_load": True, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=512, score=88.57us (B200 precompute retune 5x100) + ( + 16384, + "persistent_main", + { + "_block_size_m_nowrite": 32, + "_block_size_m_write": 32, + "_cta_per_sm_nowrite": 8, + "_cta_per_sm_write": 8, + "_flatten": False, + "_heads_per_block": 8, + "_num_loop_stages_nowrite": 2, + "_num_loop_stages_write": 1, + "_num_stages_nowrite": 3, + "_num_stages_write": 5, + "_num_warps_nowrite": 1, + "_num_warps_write": 1, + "_precompute_num_warps": 1, + "_use_tma_rect_load": False, + "_use_tma_replay_nowrite_load": False, + "_use_tma_replay_write_load": True, + "_use_tma_replay_write_store": True, + "_warp_specialize": False, + "nowrite_first": True, + "rectangle_for_nowrite": True, + }, + ), # raw_batch=1024, score=170.51us (B200 PDL-retune noise-cleaned 5x500) + ], +} +_PD_TO_PM_SPLIT_MAP = { # pd unsplit knob -> (pm_write_knob, pm_nowrite_knob) + "_block_size_m": ("_block_size_m_write", "_block_size_m_nowrite"), + "_num_warps": ("_num_warps_write", "_num_warps_nowrite"), + "_num_stages": ("_num_stages_write", "_num_stages_nowrite"), + # CPS / LS are persistent-loop knobs; pd uses _cta_per_sm + _num_loop_stages + # as unsplit, pm uses _cta_per_sm_write/_nowrite + _num_loop_stages_write/_nowrite. + "_cta_per_sm": ("_cta_per_sm_write", "_cta_per_sm_nowrite"), + "_num_loop_stages": ("_num_loop_stages_write", "_num_loop_stages_nowrite"), +} - # cb_out = CB_scaled @ x_all - cb_out = tl.dot(CB_scaled.to(tl.bfloat16), x_all.to(tl.bfloat16)) - out_all = init_out + cb_out +def _bridge_tuning_knobs(knobs: dict, from_mode: str, to_mode: str) -> dict: + """Convert a tuning dict between pd ↔ pm knob namespaces. - if HAS_D: - out_all = out_all + x_all * D[None, :] + pd → pm: copy each unsplit value to both write/nowrite split knobs; drop + the unsplit form (pm doesn't read it). + pm → pd: take the nowrite split value as the unsplit knob; drop the + write/nowrite split forms (pd doesn't read them). + Shape knobs that exist in both modes (_heads_per_block, _flatten, + _warp_specialize, TMA flags, rectangle_for_nowrite) carry over unchanged. + """ + out = dict(knobs) + if from_mode == "persistent_dynamic" and to_mode == "persistent_main": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if unsplit in out: + out.setdefault(pm_w, out[unsplit]) + out.setdefault(pm_nw, out[unsplit]) + del out[unsplit] + elif from_mode == "persistent_main" and to_mode == "persistent_dynamic": + for unsplit, (pm_w, pm_nw) in _PD_TO_PM_SPLIT_MAP.items(): + if pm_nw in out: + out.setdefault(unsplit, out[pm_nw]) + out.pop(pm_w, None) + out.pop(pm_nw, None) + return out - if HAS_Z: - for t in range(T): - z_t = tl.load( - z_ptr + t * stride_z_T + offs_m * stride_z_dim, mask=m_mask, other=0.0 - ).to(tl.float32) - out_t = tl.sum(tl.where((offs_t == t)[:, None], out_all, 0.0), axis=0) - out_t = out_t * z_t * tl.sigmoid(z_t) - tl.store(out_ptr + t * stride_out_T + offs_m * stride_out_dim, out_t, mask=m_mask) - else: - out_all_ptrs = out_ptr + offs_t[:, None] * stride_out_T + offs_m[None, :] * stride_out_dim - tl.store(out_all_ptrs, out_all, mask=t_mask[:, None] & m_mask[None, :]) +def _resolve_tuning( + batch: int, + nheads_per_rank: int, + dt_str: str, + sr_str: str, +) -> tuple[str, dict] | None: + """Look up the default mode + knobs for this (eff_batch, dt, sr) cell. -# Python wrapper + Returns (mode, knobs_dict) or None if the table has no entry covering this + dtype/sr after fallbacks. + """ + eff_b = batch * max(1, nheads_per_rank) + # Lookup chain. Order: + # 1. Exact (dt, sr). + # 2. (dt, SR) if RN missing for that dtype. + # 3. Cross-dtype fallback for dtypes we haven't tuned: + # bf16 / int16 → fp16/SR + # fp8 → int8/SR + # Unknown dtype → raise. + valid_dtypes = {"fp32", "fp16", "bf16", "int8", "int16", "fp8"} + if dt_str not in valid_dtypes: + raise ValueError( + f"replay_selective_state_update: unsupported state dtype {dt_str!r}; " + f"expected one of {sorted(valid_dtypes)}" + ) + keys_to_try = [(dt_str, sr_str)] + if sr_str == "RN": + keys_to_try.append((dt_str, "SR")) + if dt_str in ("bf16", "int16"): + keys_to_try.append(("fp16", "SR")) + elif dt_str == "fp8": + keys_to_try.append(("int8", "SR")) + entries = None + for k in keys_to_try: + if k in _DEFAULT_TUNING: + entries = _DEFAULT_TUNING[k] + break + if entries is None: + return None + # Find first threshold >= eff_b; if none, use largest entry. + for thresh, mode, knobs in entries: + if eff_b <= thresh: + return mode, dict(knobs) + thresh, mode, knobs = entries[-1] + return mode, dict(knobs) def replay_selective_state_update( @@ -614,22 +3422,77 @@ def replay_selective_state_update( B: torch.Tensor, C: torch.Tensor, out: torch.Tensor, + # Required persistent-mode plumbing (REQUIRED for both pd and pm; pd + # ignores both internally but the wrapper still demands them): + # n_writes : (1,) int32 device tensor with the count of write-mode + # slots in the batch. pm uses it to size the two halves; + # pd ignores it (per-slot runtime PNAT check). + n_writes: torch.Tensor, + replay_work_items: torch.Tensor, + state_batch_indices: torch.Tensor, D: torch.Tensor | None = None, z: torch.Tensor | None = None, dt_bias: torch.Tensor | None = None, dt_softplus: bool = False, - state_batch_indices: torch.Tensor | None = None, - pad_slot_id: int = PAD_SLOT_ID, rand_seed: torch.Tensor | None = None, philox_rounds: int = 10, + state_scales: torch.Tensor | None = None, launch_with_pdl=False, use_internal_pdl=True, + rectangle_for_nowrite: bool | None = None, + nowrite_first: bool | None = None, + mode: str | None = None, _block_size_m: int | None = None, _num_warps: int | None = None, _num_stages: int | None = None, _precompute_num_warps: int | None = None, - _precompute_num_stages: int | None = None, _heads_per_block: int | None = None, + # Per-main knobs override shared values for one half of the persistent_main + # launches. Default None ties the half-specific value to the shared knob. + # The two main kernels (write vs nowrite) have + # different per-slot work — write does a state shift + store, nowrite + # just appends — so the optimum (M, W, S, H) can differ. Precompute + # knobs are intentionally NOT split: shared precompute wins (cheaper + # launch, hotter precompute outputs in L2). Persistent CPS / LS knobs + # are also split per-main since the two persistent_main launches have + # different grid sizes. + _block_size_m_write: int | None = None, + _block_size_m_nowrite: int | None = None, + _num_warps_write: int | None = None, + _num_warps_nowrite: int | None = None, + _num_stages_write: int | None = None, + _num_stages_nowrite: int | None = None, + # Note: heads_per_block / precompute_num_warps are NOT split — they only + # affect the precompute kernel, which is shared across write/nowrite. + # TMA state-tensor toggles — 4 independent paths (see replay design notes + # item #17 for measured perf profiles). Each is False=raw load/store, True= + # use a host-built TMA tensor_descriptor for that path. + _use_tma_rect_load: bool | None = None, # rectangle path state load (nowrite-only) + _use_tma_replay_write_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=True + _use_tma_replay_write_store: bool | None = None, # SSM state store when WRITE_CHECKPOINT=True + _use_tma_replay_nowrite_load: bool | None = None, # SSM state load when WRITE_CHECKPOINT=False + _require_tma_state_layout: bool = False, + # Persistent-mode tuning kwargs (consulted for both pd and pm; pd uses + # _cta_per_sm / _num_loop_stages, pm uses the _write/_nowrite splits): + # _cta_per_sm : int — CTAs per SM in the 1D persistent grid. Internally + # expanded to `num_persistent = _cta_per_sm × NUM_SMS`. + # _num_loop_stages : int — `num_stages` arg on the inner `tl.range(...)` + # persistent loop. Note: this is loop-level, NOT the kernel-arg + # `num_stages` (which only pipelines dot-feeding loads). + # _flatten : bool — `flatten` arg on `tl.range(...)`. + # _warp_specialize : bool — `warp_specialize` arg on `tl.range(...)`. + _cta_per_sm: int | None = None, + _num_loop_stages: int | None = None, + _flatten: bool | None = None, + _warp_specialize: bool | None = None, + # Per-main persistent-specific knobs. Same rationale as the BLOCK_SIZE_M + # split above: the two persistent_main launches (write half vs nowrite + # half) have different grid sizes and per-work-item costs, so they may + # want different cta_per_sm / num_loop_stages. + _cta_per_sm_write: int | None = None, + _cta_per_sm_nowrite: int | None = None, + _num_loop_stages_write: int | None = None, + _num_loop_stages_nowrite: int | None = None, ): """ Replay SSM state update with precomputed CB and tl.dot fast-forward. @@ -649,17 +3512,19 @@ def replay_selective_state_update( main blocks until precompute completes before loading conv1d outputs (x, C) and precompute outputs (CB_scaled, decay_vec). - Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which - buffer (0 or 1) to READ from for replay. The WRITE buffer is 1 - read. - Caller must flip cache_buf_idx[slot] after each call. + Uses double-buffered cache tensors. cache_buf_idx[slot] indicates which + buffer (0 or 1) to read from for replay. Checkpoint-write steps write the + new history to the inactive buffer; no-write steps append to the active + buffer. The caller must update cache_buf_idx and PNAT with the same + checkpoint predicate used by the kernel. Arguments: state: (cache, nheads, dim, dstate) in-place. After the call, contains the state after replaying prev_num_accepted_tokens old tokens. - old_x: (cache, T, nheads, dim) bf16 — old x cache (single-buffered). - old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered old B cache. - old_dt: (cache, 2, nheads, T) fp32 — double-buffered processed dt. - old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered cumulative A*dt. + old_x: (cache, 2, max_window, nheads, dim) bf16 replay history cache. + old_B: (cache, 2, max_window, ngroups, dstate) bf16 replay history cache. + old_dt: (cache, 2, nheads, max_window) fp32 processed dt history. + old_dA_cumsum: (cache, 2, nheads, max_window) fp32 cumulative A*dt history. cache_buf_idx: (cache,) int32 — which buffer to read (0 or 1). prev_num_accepted_tokens: (cache,) int32. x: (batch, T, nheads, dim) new token inputs. @@ -668,34 +3533,102 @@ def replay_selective_state_update( B: (batch, T, ngroups, dstate). C: (batch, T, ngroups, dstate). out: (batch, T, nheads, dim) preallocated output. + n_writes: (1,) int32 device tensor with the write-mode count. + replay_work_items: (batch, 4) int32 device tensor, sorted write-first. + Persistent-main consumes all fields; persistent-dynamic only + requires the argument for a stable wrapper signature. + state_batch_indices: (batch,) int32 cache slot mapping. D: (nheads, dim) optional feed-through parameter. z: (batch, T, nheads, dim) optional silu gate. dt_bias: (nheads, dim) optional, with stride(-1)==0 (tie_hdim). - state_batch_indices: (batch,) optional cache slot mapping. - rand_seed: optional (cache_size,) int64 CUDA tensor of per-cache-slot - Philox PRNG seeds. The caller bumps this tensor in-place for each - replay invocation so CUDA graph replay still gets fresh draws; the - kernel indexes it by cache_batch_idx. When provided, state is - stochastically rounded to fp16 on store. + rand_seed: optional single-element int64 CUDA tensor of Philox PRNG + seed. The caller bumps this tensor in-place for each replay + invocation so CUDA graph replay still gets fresh draws. When + provided, state is stochastically rounded on store. Supported for + state.dtype in (fp16, int8, int16, fp8_e4m3fn). fp16+SR and + fp8+SR both require sm_100a (Blackwell B200+) — wrapper asserts + this loudly. When None, standard deterministic rounding is used. philox_rounds: number of Philox PRNG rounds (default 10). + state_scales: required when state.dtype in (int8, int16, fp8_e4m3fn). + Shape (cache_size, nheads, dim), fp32. Per-(head, dim) channel + decode scale (= 1 / encode_scale). The kernel writes scales on + replay-write steps and reads them on load (broadcast over dstate). + Ignored for non-quantized state dtypes. launch_with_pdl: enable external PDL (conv1d → precompute chain). Defaults False; caller opts in when the upstream chain is PDL-safe. Ignored on hardware that doesn't support PDL (sm < 90). use_internal_pdl: enable internal PDL (precompute → main overlap). Defaults True; override for testing only. Ignored on hardware that doesn't support PDL (sm < 90). + nowrite_first: benchmark/tuning knob for mode="persistent_main". + When None, use the tuning-table value if present. When true, + launch the nowrite half before the write half. - _-prefixed kwargs (_block_size_m, _num_warps, _num_stages, - _precompute_num_warps, _precompute_num_stages, _heads_per_block) are - benchmark-only overrides; production callers should leave them None - to use the heuristic-tuned defaults. + _-prefixed kwargs are tuning overrides; production callers should + leave them None to use the tuning-table defaults. """ + sm_version = get_sm_version() + # PDL needs sm >= 90. - if get_sm_version() < 90: + if sm_version < 90: launch_with_pdl = False use_internal_pdl = False + # Mode selection: + # mode=None (default): look up the table-tuned mode + knobs for this + # (effective_batch, dtype, sr) cell. See `_resolve_tuning` above. + # mode="persistent_dynamic": single persistent-CTA kernel covering the + # full batch. Each work-item dispatches via runtime PNAT check + # (is_write = (pnat + T) > MAX). No write/nowrite split. + # The wrapper requires replay_work_items for a uniform signature, but + # the dynamic kernel ignores its contents. + # mode="persistent_main": persistent-CTA kernel with two launches + # (write half + nowrite half). Caller MUST pre-sort replay_work_items + # write-first; the n_writes tensor partitions the persistent loop + # into the two halves with the right WRITE_CHECKPOINT constexpr + # each time. RECTANGLE constexpr (= rectangle_for_nowrite) picks + # rect vs replay for the nowrite half. nowrite_first controls + # launch order only. + # Note: mode-and-knob resolution from the default-tuning table happens + # below, after we have `batch` and `nheads`. + caller_forced_tma_state = ( + _use_tma_rect_load is True + or _use_tma_replay_write_load is True + or _use_tma_replay_write_store is True + or _use_tma_replay_nowrite_load is True + ) + + # --- Hardware support gates --- + # fp8 e4m3fn (any rounding mode) needs SM 89+ for the fp32↔e4m3 cvt PTX + # instructions (Ada Lovelace introduced them; Hopper/Blackwell carry them). + if state.dtype == torch.float8_e4m3fn: + assert sm_version >= 89, ( + "fp8_e4m3fn state requires SM 89+ (Ada Lovelace / Hopper / Blackwell) " + f"for fp32↔fp8 cvt PTX instructions; current SM is {sm_version}." + ) + + # PTX cvt.rs.* (stochastic rounding) family lands on Blackwell only. + # Wrapper fails loud; framework decides fall-back (e.g. drop SR, use RN). + # int8 / int16 SR uses pure-Triton libdevice.floor + uniform noise — no + # PTX SR instruction needed, runs anywhere. + if rand_seed is not None: + if state.dtype == torch.float16: + assert sm_version >= 100, ( + "fp16 stochastic rounding (PTX cvt.rs.f16x2.f32) requires " + f"sm_100a (Blackwell B200+); current SM is {sm_version}." + ) + elif state.dtype == torch.float8_e4m3fn: + assert sm_version >= 100, ( + "fp8 stochastic rounding (PTX cvt.rs.satfinite.e4m3x4.f32) " + f"requires sm_100a (Blackwell B200+); current SM is {sm_version}." + ) + else: + assert state.dtype in (torch.int8, torch.int16), ( + "stochastic rounding is supported only for state.dtype in " + f"(fp16, int8, int16, fp8_e4m3fn), got {state.dtype}." + ) + # --- Unsqueeze inputs to canonical shapes --- if state.dim() == 3: state = state.unsqueeze(1) @@ -733,20 +3666,211 @@ def replay_selective_state_update( cache_size, nheads, dim, dstate = state.shape batch, T, _, _ = x.shape + device = x.device ngroups = B.shape[2] assert nheads % ngroups == 0 + # --- Quantization plumbing --- + # QUANT_MAX > 0 ⇔ state is int8 / int16 / fp8_e4m3fn. Kernel-entry + # static_assert on the Triton side mirrors this invariant. + quant_max = _QUANT_MAX_BY_DTYPE.get(state.dtype, 0.0) + is_quantized = quant_max > 0.0 + + # --- Default-tuning lookup --- + # Resolve (mode, knobs) from the table when caller leaves them None. + # Caller-provided kwargs always win. If the caller forces a different + # mode, translate table knobs into that mode's namespace: persistent_dynamic + # knobs fan out to both persistent_main halves, while persistent_main + # nowrite knobs seed persistent_dynamic. + _dt_str = { + torch.float32: "fp32", + torch.float16: "fp16", + torch.bfloat16: "bf16", + torch.int8: "int8", + torch.int16: "int16", + torch.float8_e4m3fn: "fp8", + }.get(state.dtype, str(state.dtype)) + _sr_str = "SR" if rand_seed is not None else "RN" + _table_entry = _resolve_tuning(batch, nheads, _dt_str, _sr_str) + if _table_entry is None: + raise ValueError( + "replay_selective_state_update has no default tuning for " + f"state dtype {_dt_str!r} with rounding mode {_sr_str!r}." + ) + if _table_entry is not None: + _table_mode, _table_knobs = _table_entry + if mode is None: + mode = _table_mode + if mode != _table_mode: + # Bridge across modes — see header comment above. + _table_knobs = _bridge_tuning_knobs(_table_knobs, _table_mode, mode) + # Fill None-valued kwargs from table. We can't reliably mutate + # locals() for re-read, so re-bind each kwarg explicitly. + if rectangle_for_nowrite is None and "rectangle_for_nowrite" in _table_knobs: + rectangle_for_nowrite = bool(_table_knobs["rectangle_for_nowrite"]) + if nowrite_first is None and "nowrite_first" in _table_knobs: + nowrite_first = bool(_table_knobs["nowrite_first"]) + _block_size_m = ( + _block_size_m if _block_size_m is not None else _table_knobs.get("_block_size_m") + ) + _num_warps = _num_warps if _num_warps is not None else _table_knobs.get("_num_warps") + _num_stages = _num_stages if _num_stages is not None else _table_knobs.get("_num_stages") + _heads_per_block = ( + _heads_per_block + if _heads_per_block is not None + else _table_knobs.get("_heads_per_block") + ) + _precompute_num_warps = ( + _precompute_num_warps + if _precompute_num_warps is not None + else _table_knobs.get("_precompute_num_warps") + ) + _block_size_m_write = ( + _block_size_m_write + if _block_size_m_write is not None + else _table_knobs.get("_block_size_m_write") + ) + _block_size_m_nowrite = ( + _block_size_m_nowrite + if _block_size_m_nowrite is not None + else _table_knobs.get("_block_size_m_nowrite") + ) + _num_warps_write = ( + _num_warps_write + if _num_warps_write is not None + else _table_knobs.get("_num_warps_write") + ) + _num_warps_nowrite = ( + _num_warps_nowrite + if _num_warps_nowrite is not None + else _table_knobs.get("_num_warps_nowrite") + ) + _num_stages_write = ( + _num_stages_write + if _num_stages_write is not None + else _table_knobs.get("_num_stages_write") + ) + _num_stages_nowrite = ( + _num_stages_nowrite + if _num_stages_nowrite is not None + else _table_knobs.get("_num_stages_nowrite") + ) + _cta_per_sm = _cta_per_sm if _cta_per_sm is not None else _table_knobs.get("_cta_per_sm") + _num_loop_stages = ( + _num_loop_stages + if _num_loop_stages is not None + else _table_knobs.get("_num_loop_stages") + ) + # persistent_main uses split write/nowrite tuning knobs. + _num_loop_stages_write = ( + _num_loop_stages_write + if _num_loop_stages_write is not None + else _table_knobs.get("_num_loop_stages_write") + ) + _num_loop_stages_nowrite = ( + _num_loop_stages_nowrite + if _num_loop_stages_nowrite is not None + else _table_knobs.get("_num_loop_stages_nowrite") + ) + _cta_per_sm_write = ( + _cta_per_sm_write + if _cta_per_sm_write is not None + else _table_knobs.get("_cta_per_sm_write") + ) + _cta_per_sm_nowrite = ( + _cta_per_sm_nowrite + if _cta_per_sm_nowrite is not None + else _table_knobs.get("_cta_per_sm_nowrite") + ) + _flatten = _flatten if _flatten is not None else _table_knobs.get("_flatten") + _warp_specialize = ( + _warp_specialize + if _warp_specialize is not None + else _table_knobs.get("_warp_specialize") + ) + if _use_tma_rect_load is None: + _use_tma_rect_load = bool(_table_knobs.get("_use_tma_rect_load", False)) + if _use_tma_replay_write_load is None: + _use_tma_replay_write_load = bool(_table_knobs.get("_use_tma_replay_write_load", False)) + if _use_tma_replay_write_store is None: + _use_tma_replay_write_store = bool( + _table_knobs.get("_use_tma_replay_write_store", False) + ) + if _use_tma_replay_nowrite_load is None: + _use_tma_replay_nowrite_load = bool( + _table_knobs.get("_use_tma_replay_nowrite_load", False) + ) + # Final defaults for optional mode flags if neither caller nor table set them. + if mode is None: + mode = "persistent_dynamic" + if rectangle_for_nowrite is None: + rectangle_for_nowrite = False + if nowrite_first is None: + nowrite_first = False + _use_tma_rect_load = bool(_use_tma_rect_load) + _use_tma_replay_write_load = bool(_use_tma_replay_write_load) + _use_tma_replay_write_store = bool(_use_tma_replay_write_store) + _use_tma_replay_nowrite_load = bool(_use_tma_replay_nowrite_load) + if sm_version < 90: + _use_tma_rect_load = False + _use_tma_replay_write_load = False + _use_tma_replay_write_store = False + _use_tma_replay_nowrite_load = False + assert mode in ("persistent_dynamic", "persistent_main"), ( + f"unknown mode {mode!r}; expected 'persistent_dynamic' or 'persistent_main'" + ) + if is_quantized: + assert state_scales is not None, ( + f"state.dtype={state.dtype} requires state_scales tensor " + "(shape (cache_size, nheads, dim), fp32)." + ) + assert state_scales.shape == (cache_size, nheads, dim), ( + f"state_scales shape mismatch: expected {(cache_size, nheads, dim)}, " + f"got {state_scales.shape}." + ) + assert state_scales.dtype == torch.float32, ( + f"state_scales must be fp32, got {state_scales.dtype}." + ) + assert state_scales.device == state.device + + # Cache window capacity comes from old_x.shape[2]; it may equal T or be + # larger when retaining replay history. Replay and rectangle window tile + # sizes are derived independently from MAX_REPLAY_BUFFER_LENGTH so + # max_window can exceed BLOCK_SIZE_T freely. + max_window = old_x.shape[2] + assert T <= max_window, f"T={T} exceeds cache max_window={max_window}" + assert x.shape == (batch, T, nheads, dim) assert dt.shape == x.shape assert A.shape == (nheads, dim, dstate) assert B.shape == (batch, T, ngroups, dstate) assert C.shape == B.shape - assert old_x.shape == (cache_size, T, nheads, dim) - assert old_B.shape == (cache_size, 2, T, ngroups, dstate) - assert old_dt.shape == (cache_size, 2, nheads, T) - assert old_dA_cumsum.shape == (cache_size, 2, nheads, T) + assert old_x.shape == (cache_size, 2, max_window, nheads, dim) + assert old_B.shape == (cache_size, 2, max_window, ngroups, dstate) + assert old_dt.shape == (cache_size, 2, nheads, max_window) + assert old_dA_cumsum.shape == (cache_size, 2, nheads, max_window) assert cache_buf_idx.shape == (cache_size,) + assert cache_buf_idx.dtype == torch.int32, ( + f"cache_buf_idx must be int32, got {cache_buf_idx.dtype}" + ) assert prev_num_accepted_tokens.shape == (cache_size,) + assert prev_num_accepted_tokens.dtype == torch.int32, ( + f"prev_num_accepted_tokens must be int32, got {prev_num_accepted_tokens.dtype}" + ) + assert isinstance(state_batch_indices, torch.Tensor), ( + f"state_batch_indices must be a torch.Tensor, got {type(state_batch_indices).__name__}" + ) + assert state_batch_indices.device == device, ( + f"state_batch_indices must be on device {device}, got {state_batch_indices.device}" + ) + assert state_batch_indices.dtype == torch.int32, ( + f"state_batch_indices must be int32, got {state_batch_indices.dtype}" + ) + assert state_batch_indices.shape == (batch,), ( + f"state_batch_indices must have shape (batch={batch},), " + f"got {tuple(state_batch_indices.shape)}" + ) + assert state_batch_indices.is_contiguous(), "state_batch_indices must be contiguous" if rand_seed is not None: assert rand_seed.dtype == torch.int64, ( f"rand_seed dtype must be int64, got {rand_seed.dtype}" @@ -754,11 +3878,8 @@ def replay_selective_state_update( assert rand_seed.dim() == 1, ( f"rand_seed must be a 1D tensor; got shape {tuple(rand_seed.shape)}" ) - if rand_seed.shape[0] == 1 and cache_size > 1: - rand_seed = rand_seed.expand(cache_size).contiguous() - assert rand_seed.shape[0] >= cache_size, ( - f"rand_seed must have length 1 or >= cache_size ({cache_size}); " - f"got shape {tuple(rand_seed.shape)}" + assert rand_seed.shape[0] == 1, ( + f"rand_seed must have length 1; got shape {tuple(rand_seed.shape)}" ) tie_hdim = ( @@ -769,12 +3890,22 @@ def replay_selective_state_update( ) assert tie_hdim - device = x.device - BLOCK_SIZE_T = max(triton.next_power_of_2(T), 16) + BLOCK_SIZE_T = max(triton.next_power_of_2(T), MIN_REPLAY_TILE_SIZE) + # Rectangle window bound = max_window. Computed unconditionally + # so the launch sites can refer to it; only used on the rectangle path. + # If this differs from BLOCK_SIZE_T, rectangle precompute uses a slower + # one-hot fallback because tl.gather requires matching padded dimension sizes. + # Production uses matching padded T and window sizes. + BLOCK_SIZE_K = max(triton.next_power_of_2(max_window), MIN_REPLAY_TILE_SIZE) + rectangle_use_gather = BLOCK_SIZE_T == BLOCK_SIZE_K - # Allocate precomputed intermediates (per-call, not cached). + # Allocate precomputed intermediates (per-call, not cached). Always + # allocate (T, window) — the largest layout that any path uses. Replay-style + # paths only touch the first T columns; rectangle/dynamic use the full window. + # The few extra unused columns per row are negligible (~6KB per layer at + # production sizes) and let the dispatch helpers share one buffer. cb_scaled = torch.empty( - batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_T, device=device, dtype=torch.float32 + batch, nheads, BLOCK_SIZE_T, BLOCK_SIZE_K, device=device, dtype=torch.float32 ) decay_vec = torch.empty(batch, nheads, BLOCK_SIZE_T, device=device, dtype=torch.float32) @@ -782,129 +3913,206 @@ def replay_selective_state_update( (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) if z is not None else (0, 0, 0, 0) ) - # Kernel tuning: BLOCK_SIZE_M, num_warps, HEADS_PER_BLOCK, precompute_num_warps. - # Dtype-aware heuristic from B200 sweeps (batch 1-512, T=6/32, TP=8, conv1d + - # chained PDL). Keyed on total_heads, BLOCK_SIZE_T, and state dtype; 16-bit - # states prefer different tiles from fp32 due to lower bandwidth. Philox - # gets its own branch — stochastic rounding shifts compute toward CUDA cores, - # so small-batch configs want more warps to hide the extra work. - total_heads = batch * nheads heads_per_group = nheads // ngroups - state_is_16bit = state.dtype in (torch.float16, torch.bfloat16) - use_philox = rand_seed is not None - if BLOCK_SIZE_T <= 16: - if use_philox and state_is_16bit: - # Philox: more warps at small batch to hide CUDA core work. - # At large batch, converges to non-Philox fp16 config. - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 4, 4, 4, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - elif state_is_16bit: - if total_heads <= 16: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # fp32 state (no Philox — fp32 doesn't need stochastic rounding) - if total_heads <= 32: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 4, 1 - elif total_heads <= 64: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 1, 2, 1 - elif total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 8, 2, 2, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 1, 2, 1 - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(2, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 32, 4, 2, 1 - else: # T > 16 - if state_is_16bit: - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 16, - 1, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 1, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 1, - 4, - min(2, heads_per_group), - ) - else: # fp32 state - if total_heads <= 128: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = 16, 2, 4, 1 - elif total_heads <= 256: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 32, - 2, - 4, - min(2, heads_per_group), - ) - elif total_heads <= 512: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 2, - min(4, heads_per_group), - ) - else: - BLOCK_SIZE_M, num_warps, precompute_num_warps, heads_per_block = ( - 64, - 2, - 4, - min(2, heads_per_group), - ) - if _block_size_m is not None: - BLOCK_SIZE_M = _block_size_m - if _num_warps is not None: - num_warps = _num_warps - if _heads_per_block is not None: - heads_per_block = _heads_per_block - if _precompute_num_warps is not None: - precompute_num_warps = _precompute_num_warps + if mode == "persistent_dynamic": + assert _block_size_m is not None, "persistent_dynamic requires _block_size_m tuning" + assert _num_warps is not None, "persistent_dynamic requires _num_warps tuning" + assert _num_stages is not None, "persistent_dynamic requires _num_stages tuning" + assert _cta_per_sm is not None, "persistent_dynamic requires _cta_per_sm tuning" + assert _num_loop_stages is not None, "persistent_dynamic requires _num_loop_stages tuning" + else: + assert _block_size_m_write is not None, ( + "persistent_main requires _block_size_m_write tuning" + ) + assert _block_size_m_nowrite is not None, ( + "persistent_main requires _block_size_m_nowrite tuning" + ) + assert _num_warps_write is not None, "persistent_main requires _num_warps_write tuning" + assert _num_warps_nowrite is not None, "persistent_main requires _num_warps_nowrite tuning" + assert _num_stages_write is not None, "persistent_main requires _num_stages_write tuning" + assert _num_stages_nowrite is not None, ( + "persistent_main requires _num_stages_nowrite tuning" + ) + assert _cta_per_sm_write is not None, "persistent_main requires _cta_per_sm_write tuning" + assert _cta_per_sm_nowrite is not None, ( + "persistent_main requires _cta_per_sm_nowrite tuning" + ) + assert _num_loop_stages_write is not None, ( + "persistent_main requires _num_loop_stages_write tuning" + ) + assert _num_loop_stages_nowrite is not None, ( + "persistent_main requires _num_loop_stages_nowrite tuning" + ) + assert _heads_per_block is not None, "replay default tuning requires _heads_per_block" + assert _precompute_num_warps is not None, "replay default tuning requires _precompute_num_warps" + BLOCK_SIZE_M = _block_size_m if _block_size_m is not None else _block_size_m_nowrite + num_warps = _num_warps if _num_warps is not None else _num_warps_nowrite + precompute_num_warps = _precompute_num_warps + heads_per_block = int(_heads_per_block) + assert heads_per_block > 0, "heads_per_block must be positive" + heads_per_block = min(heads_per_block, heads_per_group) + while heads_per_group % heads_per_block != 0 or heads_per_block & (heads_per_block - 1) != 0: + heads_per_block -= 1 - HAS_CACHE_BATCH_INDICES = state_batch_indices is not None + # Per-main knob resolution: each _*_{write,nowrite} arg, if not None, + # overrides the corresponding shared value for ONE main launch only. + # Default (None) = tied to shared value (current behavior). + BLOCK_SIZE_M_WRITE = _block_size_m_write if _block_size_m_write is not None else BLOCK_SIZE_M + BLOCK_SIZE_M_NOWRITE = ( + _block_size_m_nowrite if _block_size_m_nowrite is not None else BLOCK_SIZE_M + ) + NUM_WARPS_WRITE = _num_warps_write if _num_warps_write is not None else num_warps + NUM_WARPS_NOWRITE = _num_warps_nowrite if _num_warps_nowrite is not None else num_warps + NUM_STAGES_WRITE = _num_stages_write if _num_stages_write is not None else _num_stages + NUM_STAGES_NOWRITE = _num_stages_nowrite if _num_stages_nowrite is not None else _num_stages + # Persistent-only per-main: + CTA_PER_SM_WRITE = _cta_per_sm_write if _cta_per_sm_write is not None else _cta_per_sm + CTA_PER_SM_NOWRITE = _cta_per_sm_nowrite if _cta_per_sm_nowrite is not None else _cta_per_sm + NUM_LOOP_STAGES_WRITE = ( + _num_loop_stages_write if _num_loop_stages_write is not None else _num_loop_stages + ) + NUM_LOOP_STAGES_NOWRITE = ( + _num_loop_stages_nowrite if _num_loop_stages_nowrite is not None else _num_loop_stages + ) - with torch.cuda.device(device.index): - # --- Precompute kernel --- - assert nheads % heads_per_block == 0, ( - f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + assert nheads % heads_per_block == 0, ( + f"nheads ({nheads}) must be divisible by heads_per_block ({heads_per_block})" + ) + assert heads_per_block <= heads_per_group, ( + f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + ) + assert heads_per_group % heads_per_block == 0, ( + f"heads_per_block ({heads_per_block}) must divide heads_per_group ({heads_per_group})" + ) + + # state_scales pointer + strides: real tensor when quantized, otherwise + # zero-strided dummy (kernel never reads it because QUANT_MAX==0.0). + if is_quantized: + state_scales_arg = state_scales + state_scales_strides = ( + state_scales.stride(0), + state_scales.stride(1), + state_scales.stride(2), + ) + else: + state_scales_arg = state # any valid pointer; gated by QUANT_MAX == 0 + state_scales_strides = (0, 0, 0) + + # Per-path TMA descriptors for state — write-side and nowrite-side. + # Kernels see state as a 2D row-space: [cache/head/dim row, dstate]. + # Dense tensors use the original flat view. Block-reuse cache tensors may + # have a gap between slots (conv state packed after SSM state), but the + # per-slot SSM rows are still dense and can use the same 2D row-space with + # explicit strides. Unsupported layouts fall back to raw loads/stores. + use_tma_state = ( + _use_tma_rect_load + or _use_tma_replay_write_load + or _use_tma_replay_write_store + or _use_tma_replay_nowrite_load + ) + if use_tma_state: + state_row_stride = state.stride(2) + tma_state_supported = ( + state.stride(-1) == 1 + and state_row_stride > 0 + and state.stride(1) == dim * state_row_stride + and state.stride(0) % state_row_stride == 0 + and (state_row_stride * state.element_size()) % 16 == 0 ) - assert heads_per_block <= heads_per_group, ( - f"heads_per_block ({heads_per_block}) must not cross group boundary ({heads_per_group})" + if not tma_state_supported: + if _require_tma_state_layout or caller_forced_tma_state: + raise AssertionError( + "TMA state layout requires inner stride 1, dense dim rows " + "within each head, cache stride aligned to dim-row stride, " + "and 16-byte-aligned row stride; got " + f"shape={tuple(state.shape)} strides={state.stride()}" + ) + _use_tma_rect_load = False + _use_tma_replay_write_load = False + _use_tma_replay_write_store = False + _use_tma_replay_nowrite_load = False + use_tma_state = False + + # Each kernel launch consumes the descriptor whose block_shape[0] matches + # its BLOCK_SIZE_M constexpr. With M-split (Mw != Mnw) the two sides need + # distinct descriptors; otherwise the descriptor's block_shape[0] would + # mismatch the kernel's BLOCK_SIZE_M and downstream tl.dot/arithmetic on + # the loaded tile fails shape inference at compile time. When no TMA flag + # is on, both variables hold raw `state` as a dummy; kernels never + # reference it because their constexprs are all False. + if use_tma_state: + from triton.tools.tensor_descriptor import TensorDescriptor + + _ensure_tma_allocator() + if state.is_contiguous(): + state_tma_base = state.view(-1, state.shape[-1]) + make_state_tma_descriptor = TensorDescriptor.from_tensor + else: + slot_stride_rows = state.stride(0) // state_row_stride + state_rows = (cache_size - 1) * slot_stride_rows + nheads * dim + state_tma_base = state + state_tma_shape = [state_rows, dstate] + state_tma_strides = [state_row_stride, state.stride(3)] + + def make_state_tma_descriptor(_state, block_shape): + return TensorDescriptor( + _state, + state_tma_shape, + state_tma_strides, + block_shape=block_shape, + ) + + _dstate_pow2 = triton.next_power_of_2(dstate) + state_tma_descriptor_write = make_state_tma_descriptor( + state_tma_base, + block_shape=[BLOCK_SIZE_M_WRITE, _dstate_pow2], ) - _replay_precompute_kernel[(batch, nheads // heads_per_block)]( + if BLOCK_SIZE_M_NOWRITE == BLOCK_SIZE_M_WRITE: + state_tma_descriptor_nowrite = state_tma_descriptor_write + else: + state_tma_descriptor_nowrite = make_state_tma_descriptor( + state_tma_base, + block_shape=[BLOCK_SIZE_M_NOWRITE, _dstate_pow2], + ) + else: + state_tma_descriptor_write = state # dummy; all consuming constexprs are false + state_tma_descriptor_nowrite = state # dummy; all consuming constexprs are false + + # Work items are sorted write-first for persistent_main. Each row carries + # decode-batch position, cache slot, PNAT, and active cache buffer index. + assert isinstance(replay_work_items, torch.Tensor), ( + f"replay_work_items must be a torch.Tensor, got {type(replay_work_items).__name__}" + ) + assert replay_work_items.device == device, ( + f"replay_work_items must be on device {device}, got {replay_work_items.device}" + ) + assert replay_work_items.dtype == torch.int32, ( + f"replay_work_items must be int32, got {replay_work_items.dtype}" + ) + assert replay_work_items.shape == (batch, REPLAY_WORK_ITEM_WIDTH), ( + "replay_work_items must have shape " + f"(batch={batch}, {REPLAY_WORK_ITEM_WIDTH}), got " + f"{tuple(replay_work_items.shape)}" + ) + assert replay_work_items.is_contiguous(), "replay_work_items must be contiguous" + assert isinstance(n_writes, torch.Tensor), ( + f"n_writes must be a torch.Tensor, got {type(n_writes).__name__}" + ) + assert n_writes.device == device, f"n_writes must be on device {device}, got {n_writes.device}" + assert n_writes.dtype == torch.int32, f"n_writes must be int32, got {n_writes.dtype}" + assert n_writes.shape == (1,), f"n_writes must have shape (1,), got {tuple(n_writes.shape)}" + replay_work_items_arg = replay_work_items + + precomp_grid = (batch, nheads // heads_per_block) + d_strides = (D.stride(0), D.stride(1)) if D is not None else (0, 0) + + # ---- Launch helpers (close over locals) ------------------------------- + # Each helper is a thin closure that calls one Triton kernel with the + # full positional + kwarg argument list. Mode-dependent constexprs + # (write_checkpoint_mode, rectangle) are passed in. + + def launch_dynamic_precompute(rectangle: bool): + _dynamic_precompute_kernel[precomp_grid]( dt, dt_bias, A, @@ -916,68 +4124,242 @@ def replay_selective_state_update( old_dt, old_dA_cumsum, cache_buf_idx, + prev_num_accepted_tokens, state_batch_indices, - pad_slot_id, T, + max_window, dstate, nheads // ngroups, - # dt strides dt.stride(0), dt.stride(1), dt.stride(2), dt_bias.stride(0) if dt_bias is not None else 0, A.stride(0), - # B strides B.stride(0), B.stride(1), B.stride(2), B.stride(3), - # C strides C.stride(0), C.stride(1), C.stride(2), C.stride(3), - # cb_scaled strides cb_scaled.stride(0), cb_scaled.stride(1), cb_scaled.stride(2), cb_scaled.stride(3), - # decay_vec strides decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), - # old_B strides old_B.stride(0), old_B.stride(1), old_B.stride(2), old_B.stride(3), old_B.stride(4), - # old_dt strides old_dt.stride(0), old_dt.stride(1), old_dt.stride(2), old_dt.stride(3), - # old_dA_cumsum strides old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), dt_softplus, - HAS_CACHE_BATCH_INDICES=HAS_CACHE_BATCH_INDICES, LAUNCH_WITH_PDL=launch_with_pdl, LAUNCH_DEPENDENT_KERNELS=use_internal_pdl, HEADS_PER_BLOCK=heads_per_block, + RECTANGLE=rectangle, + RECTANGLE_USE_GATHER=rectangle_use_gather, num_warps=precompute_num_warps, - **({"num_stages": _precompute_num_stages} if _precompute_num_stages else {}), launch_pdl=launch_with_pdl, ) - # --- Main kernel --- - def grid(META): - return (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + # ---- launch_persistent_main ------------------------------------------ + # Persistent main launches the write and nowrite halves separately. + # Replay work items are write-first, and the device n_writes tensor + # partitions [0, n_writes) from [n_writes, batch). + + # Resolve persistent-mode tuning knobs. None values fall back to stable + # defaults so direct callers do not need to mirror benchmark search args. + _num_sms = torch.cuda.get_device_properties(device).multi_processor_count + cta_per_sm_arg = _cta_per_sm if _cta_per_sm else 1 + num_persistent_arg = cta_per_sm_arg * _num_sms + num_loop_stages_arg = _num_loop_stages if _num_loop_stages else 2 + flatten_arg = True if _flatten is None else bool(_flatten) + warp_specialize_arg = False if _warp_specialize is None else bool(_warp_specialize) + # Per-launch work-item count. At small batch, total_work may be < the + # full persistent grid; capping `grid` at `min(NUM_PERSISTENT, total_work)` + # avoids launching empty CTAs that pay setup cost for no work. Correctness: + # the kernel's `tl.range(pid, total_work, NUM_PERSISTENT)` ensures each + # tile_id is covered exactly once across all live pids in [0, grid) when + # grid <= NUM_PERSISTENT (each CTA does 1 tile; loop step >= total_work + # exits immediately) AND when grid == NUM_PERSISTENT (each CTA loops over + # multiple tiles). NUM_PERSISTENT is now a runtime int (see kernel def + # docstring at _persistent_main_kernel) so changing cta_per_sm does NOT + # trigger a new Triton compile — same kernel binary, different loop step. + # Runtime value, despite the Triton-style uppercase name. + _num_pid_m = (dim + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M + + def launch_persistent_main( + write_checkpoint_mode: bool, + *, + launch_dependent_kernels: bool = False, + rectangle: bool = False, + ): + # `n_writes` is the (1,) int32 device tensor with the write count. + # Both halves always launch; an empty half has a zero-length slot + # range and the persistent loop does no work. + block_size_m = BLOCK_SIZE_M_WRITE if write_checkpoint_mode else BLOCK_SIZE_M_NOWRITE + launch_num_warps = NUM_WARPS_WRITE if write_checkpoint_mode else NUM_WARPS_NOWRITE + launch_num_stages = NUM_STAGES_WRITE if write_checkpoint_mode else NUM_STAGES_NOWRITE + ctas_per_sm = CTA_PER_SM_WRITE if write_checkpoint_mode else CTA_PER_SM_NOWRITE + ctas_per_sm = ctas_per_sm if ctas_per_sm else 1 + num_loop_stages = ( + NUM_LOOP_STAGES_WRITE if write_checkpoint_mode else NUM_LOOP_STAGES_NOWRITE + ) + num_loop_stages = num_loop_stages if num_loop_stages else 2 + num_persistent = ctas_per_sm * _num_sms + num_pid_m_local = (dim + block_size_m - 1) // block_size_m + # Grid sizing: cap at min(full persistent grid, upper-bound total work). + # We use `batch` as the upper bound on slots-per-half — overcounting + # by a few CTAs is fine since the kernel derives the exact slot range + # from n_writes at runtime. + total_work_launch = max(1, batch * num_pid_m_local * nheads) + grid = (min(num_persistent, total_work_launch),) + # Per-path TMA descriptor — block_shape[0] must match block_size_m. + selected_state_tma_descriptor = ( + state_tma_descriptor_write if write_checkpoint_mode else state_tma_descriptor_nowrite + ) + _persistent_main_kernel[grid]( + state, + selected_state_tma_descriptor, + state_scales_arg, + old_x, + old_B, + old_dt, + old_dA_cumsum, + prev_num_accepted_tokens, + cache_buf_idx, + x, + C, + D, + z, + out, + cb_scaled, + decay_vec, + state_batch_indices, + replay_work_items_arg, + rand_seed, + n_writes, + batch, + nheads, + T, + max_window, + dim, + dstate, + nheads // ngroups, + state.stride(0), + state.stride(1), + state.stride(2), + state.stride(3), + state_scales_strides[0], + state_scales_strides[1], + state_scales_strides[2], + old_x.stride(0), + old_x.stride(1), + old_x.stride(2), + old_x.stride(3), + old_x.stride(4), + old_B.stride(0), + old_B.stride(1), + old_B.stride(2), + old_B.stride(3), + old_B.stride(4), + old_dt.stride(0), + old_dt.stride(1), + old_dt.stride(2), + old_dt.stride(3), + old_dA_cumsum.stride(0), + old_dA_cumsum.stride(1), + old_dA_cumsum.stride(2), + old_dA_cumsum.stride(3), + x.stride(0), + x.stride(1), + x.stride(2), + x.stride(3), + C.stride(0), + C.stride(1), + C.stride(2), + C.stride(3), + d_strides[0], + d_strides[1], + z_strides[0], + z_strides[1], + z_strides[2], + z_strides[3], + out.stride(0), + out.stride(1), + out.stride(2), + out.stride(3), + cb_scaled.stride(0), + cb_scaled.stride(1), + cb_scaled.stride(2), + cb_scaled.stride(3), + decay_vec.stride(0), + decay_vec.stride(1), + decay_vec.stride(2), + block_size_m, + LAUNCH_WITH_PDL=use_internal_pdl, + PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=write_checkpoint_mode, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + NUM_PERSISTENT=num_persistent, + NUM_LOOP_STAGES=num_loop_stages, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=False, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=False: WC fixed per launch; impl + # constexpr-folds the LOAD pick. When WRITE_CHECKPOINT=True (write half), + # NOWRITE_LOAD is dummy False; when WRITE_CHECKPOINT=False, WRITE_LOAD/STORE + # dummy False. NOWRITE_LOAD picks rect-load (RECTANGLE) or + # replay-nowrite-load. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load and write_checkpoint_mode), + USE_TMA_LOAD_NOWRITE=bool( + (_use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load) + and not write_checkpoint_mode + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store and write_checkpoint_mode), + num_warps=launch_num_warps, + **({"num_stages": launch_num_stages} if launch_num_stages else {}), + launch_pdl=use_internal_pdl, + ) - _replay_state_update_kernel[grid]( + def launch_persistent_dynamic_main( + n_writes_tensor: torch.Tensor, + launch_dependent_kernels: bool = False, + rectangle: bool = False, + ): + # Single-launch persistent kernel covering the whole batch with + # runtime per-slot WRITE_CHECKPOINT branch. No half-split, no + # n_writes needed (the kernel ignores n_writes_tensor when + # IS_DYNAMIC=True; Triton DCEs the load). is_write is computed + # at runtime per work-item from the loaded PNAT. + # We still pass the same tensor as persistent_main so the kernel + # signature is uniform. + # Grid sizing: cap at total_work (= batch * num_pid_m * nheads) for + # the dynamic case (full-batch coverage); see launch_persistent_main + # comment for correctness rationale. + total_work_launch = max(1, batch * _num_pid_m * nheads) + grid = (min(num_persistent_arg, total_work_launch),) + # Persistent-dynamic kernel uses a single BLOCK_SIZE_M (same as the + # wrapper's BLOCK_SIZE_M == BLOCK_SIZE_M_WRITE tied convention), so + # the write-side descriptor matches. Both write and nowrite slots + # in this kernel share that BSM. + _persistent_main_kernel[grid]( state, + state_tma_descriptor_write, + state_scales_arg, old_x, old_B, old_dt, @@ -992,73 +4374,140 @@ def grid(META): cb_scaled, decay_vec, state_batch_indices, + replay_work_items_arg, rand_seed, - pad_slot_id, + n_writes_tensor, + batch, + nheads, T, + max_window, dim, dstate, nheads // ngroups, - # state strides state.stride(0), state.stride(1), state.stride(2), state.stride(3), - # old_x strides (single-buffered: cache, T, nheads, dim) + state_scales_strides[0], + state_scales_strides[1], + state_scales_strides[2], old_x.stride(0), old_x.stride(1), old_x.stride(2), old_x.stride(3), - # old_B strides + old_x.stride(4), old_B.stride(0), old_B.stride(1), old_B.stride(2), old_B.stride(3), old_B.stride(4), - # old_dt strides old_dt.stride(0), old_dt.stride(1), old_dt.stride(2), old_dt.stride(3), - # old_dA_cumsum strides old_dA_cumsum.stride(0), old_dA_cumsum.stride(1), old_dA_cumsum.stride(2), old_dA_cumsum.stride(3), - # x strides x.stride(0), x.stride(1), x.stride(2), x.stride(3), - # C strides C.stride(0), C.stride(1), C.stride(2), C.stride(3), - # D strides - *(D.stride(0), D.stride(1)) if D is not None else (0, 0), - # z strides + d_strides[0], + d_strides[1], z_strides[0], z_strides[1], z_strides[2], z_strides[3], - # out strides out.stride(0), out.stride(1), out.stride(2), out.stride(3), - # cb_scaled strides cb_scaled.stride(0), cb_scaled.stride(1), cb_scaled.stride(2), cb_scaled.stride(3), - # decay_vec strides decay_vec.stride(0), decay_vec.stride(1), decay_vec.stride(2), BLOCK_SIZE_M, LAUNCH_WITH_PDL=use_internal_pdl, PHILOX_ROUNDS=philox_rounds if rand_seed is not None else 0, + QUANT_MAX=quant_max, + WRITE_CHECKPOINT=False, + LAUNCH_DEPENDENT_KERNELS=launch_dependent_kernels and use_internal_pdl, + NUM_PERSISTENT=num_persistent_arg, + NUM_LOOP_STAGES=num_loop_stages_arg, + FLATTEN=flatten_arg, + WARP_SPECIALIZE=warp_specialize_arg, + IS_DYNAMIC=True, + RECTANGLE=rectangle, + # 3 TMA flags. IS_DYNAMIC=True: is_write is runtime per slot; + # impl's load TMA picks per-slot (constexpr ternary becomes a + # runtime branch — both load forms emitted, ~negligible cost). + # NOWRITE_LOAD picks rect-load when RECTANGLE, else + # replay-nowrite-load. STORE only fires on runtime is_write. + USE_TMA_LOAD_WRITE=bool(_use_tma_replay_write_load), + USE_TMA_LOAD_NOWRITE=bool( + _use_tma_rect_load if rectangle else _use_tma_replay_nowrite_load + ), + USE_TMA_STORE=bool(_use_tma_replay_write_store), num_warps=num_warps, **({"num_stages": _num_stages} if _num_stages else {}), launch_pdl=use_internal_pdl, ) + + # ---- Mode dispatch ---------------------------------------------------- + with torch.cuda.device(device.index): + if mode == "persistent_dynamic": + # Single-launch persistent kernel covering the full batch. Each + # work-item dispatches via runtime PNAT check. Kernel ignores + # n_writes (Triton DCEs the load) when IS_DYNAMIC=True; we still + # pass the wrapper-provided tensor as required by the signature. + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + launch_persistent_dynamic_main( + n_writes, + launch_dependent_kernels=False, + rectangle=rectangle_for_nowrite, + ) + elif mode == "persistent_main": + # Persistent-CTA main kernel. One shared dynamic_precompute + # (per-slot dispatch via PNAT) feeds two persistent_main + # launches (write half + nowrite half). The first main launch in + # program order signals the second one when internal PDL is on. + # + # Caller-provided contract: `n_writes` is a (1,) int32 device + # tensor (the kernel reads it at runtime, after the precompute); + # `replay_work_items` is a (batch, 4) int32 device tensor + # pre-sorted write-first. + def launch_nowrite(launch_dependent_kernels: bool): + launch_persistent_main( + write_checkpoint_mode=False, + launch_dependent_kernels=launch_dependent_kernels, + rectangle=rectangle_for_nowrite, + ) + + launch_dynamic_precompute(rectangle=rectangle_for_nowrite) + if nowrite_first: + launch_nowrite(launch_dependent_kernels=True) + launch_persistent_main( + write_checkpoint_mode=True, + launch_dependent_kernels=False, + rectangle=False, # write always uses the replay-style path + ) + else: + launch_persistent_main( + write_checkpoint_mode=True, + launch_dependent_kernels=True, + rectangle=False, # write always uses the replay-style path + ) + launch_nowrite(launch_dependent_kernels=False) + else: + raise ValueError( + f"mode={mode!r} is not supported. Supported modes: " + f"'persistent_dynamic', 'persistent_main'." + ) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 5ae9ce95fbd0..5a21819d7df7 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -98,9 +98,19 @@ def get_kv_cache_manager_cls( logger.info("Hybrid linear model has 0 mamba layers; using " "KVCacheManager without mamba caching") return _non_hybrid_kv_cache_manager_cls(config, kv_cache_config) + if use_py_mamba_cache_manager(): + if kv_cache_config.enable_block_reuse: + raise ValueError( + "TRTLLM_USE_PY_MAMBA=1 forces " + "MixedMambaHybridCacheManager, which does not support " + "block reuse. Disable block reuse or unset " + "TRTLLM_USE_PY_MAMBA to use CppMambaHybridCacheManager.") + logger.info( + "Using MixedMambaHybridCacheManager for hybrid mamba model") + return MixedMambaHybridCacheManager if kv_cache_config.enable_block_reuse: return CppMambaHybridCacheManager - if use_cpp_mamba_cache_manager() or use_py_mamba_cache_manager(): + if use_cpp_mamba_cache_manager(): logger.info( "Using MixedMambaHybridCacheManager for hybrid mamba model") return MixedMambaHybridCacheManager @@ -1654,7 +1664,11 @@ def _create_kv_cache_manager( quant_config, 'mamba_ssm_stochastic_rounding', False) if quant_config is not None else False - use_replay = sm >= 80 + use_replay = spec_config is not None and sm >= 80 + if spec_config is None: + logger.info( + "Replay kernel requires speculative decoding; using non-replay path" + ) # Block reuse (prefix caching): replay leaves SSM state at a # checkpoint after speculation. The next decode step replays forward diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 7f1af74d4c83..43535f838160 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -17,7 +17,7 @@ import os from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import TYPE_CHECKING, Dict, List, Optional, Union +from typing import TYPE_CHECKING, Dict, List, NamedTuple, Optional, Union import torch import triton @@ -26,8 +26,7 @@ import tensorrt_llm.bindings if TYPE_CHECKING: - from tensorrt_llm._torch.attention_backend.interface import \ - AttentionMetadata + from tensorrt_llm._torch.attention_backend.interface import AttentionMetadata from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig from tensorrt_llm._torch.pyexecutor.llm_request import ( @@ -49,6 +48,20 @@ GB = 1 << 30 +# Replay kernels pad the token/window dimension to at least 16 for tensor-core +# tiles, so history sizes below 16 are no faster when not writing and do +# expensive checkpoint writes more often. This minimum tensor-core tile size is +# present in all tensor-core generations, with newer GPUs adding larger tile +# operations, not smaller ones. In the other direction, history sizes above 16 +# currently pad to 32, which is substantially slower with the current kernel +# design. Keep the default floor at 16 while still allowing larger T, which +# degenerates to checkpointing every step instead of rejecting the request. If +# larger T becomes common, we could consider padding larger histories to powers +# of 2 or multiples of 16, but this is untested. For smaller T, we could explore +# combining larger histories with new kernel designs that stay efficient when +# the window is only partly full. +MIN_REPLAY_HISTORY_SIZE = 16 + def get_tensor_size_bytes(tensor): """Calculate tensor size in bytes.""" @@ -166,6 +179,14 @@ def use_py_mamba_cache_manager() -> bool: return py +class ReplayStateUpdateMetadata(NamedTuple): + """Shared tensors and fixed sizes for replay state updates.""" + prev_num_accepted_tokens: torch.Tensor + cache_buf_idx: torch.Tensor + replay_step_width: int + replay_history_size: int + + class BaseMambaCacheManager(ABC): """Abstract interface for accessing mamba/recurrent state caches.""" @@ -177,6 +198,11 @@ def get_state_indices(self, *args, **kwargs) -> torch.Tensor: """ ... + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + """Return replay metadata tensors and fixed replay sizes.""" + return None + @abstractmethod def get_conv_states(self, layer_idx: int) -> torch.Tensor: """Return conv states for specific layer. @@ -371,11 +397,11 @@ class SpeculativeState(State): # CUDA graph replay uses fresh SR draws without allocating RNG tensors. # (cache,) int64 - shared across layers mamba_ssm_rand_seed: torch.Tensor | None = None - old_x: torch.Tensor | None = None # (layers, cache, T, nheads, dim) - old_B: torch.Tensor | None = None # (layers, cache, 2, T, ngroups, dstate) + old_x: torch.Tensor | None = None # (layers, cache, 2, history, nheads, dim) + old_B: torch.Tensor | None = None # (layers, cache, 2, history, ngroups, dstate) # Processed dt: softplus(raw_dt + dt_bias), clamped to dt_limit. - old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, T) fp32 - old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, T) fp32 + old_dt: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 def __init__( self, @@ -401,6 +427,8 @@ def __init__( self.speculative_num_draft_tokens = speculative_num_draft_tokens self.spec_state_size = spec_state_size self._use_replay_state_update = use_replay_state_update + self.replay_history_size: Optional[int] = None + self.replay_step_width: Optional[int] = None # When True, allocate the per-slot Philox seed buffer even outside # the replay path so the non-replay flashinfer SR kernel reads a # persistent deterministic seed instead of a per-call torch.randint. @@ -459,13 +487,13 @@ def __init__( ssm_state_shape = (nheads, head_dim, d_state) # create mamba conv and ssm states - conv_states = torch.empty( + conv_states = torch.zeros( size=(num_local_layers, max_batch_size) + conv_state_shape, dtype=dtype, device=device, ) - ssm_states = torch.empty( + ssm_states = torch.zeros( size=(num_local_layers, max_batch_size) + ssm_state_shape, dtype=self.mamba_ssm_cache_dtype, device=device, @@ -485,6 +513,7 @@ def __init__( # create state container if speculative_num_draft_tokens is not None: T = speculative_num_draft_tokens + 1 + self.replay_step_width = T # Conv intermediate cache — same for both paths intermediate_conv_window_cache = torch.zeros( @@ -504,19 +533,18 @@ def __init__( assert n_groups % tp_size == 0, \ "replay state update requires n_groups divisible by tp_size" n_groups_per_rank = n_groups // tp_size + self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) # Compact replay cache. - # old_x is single-buffered (written by main kernel after replay). - # old_B, old_dt, old_dA_cumsum are double-buffered (written by - # precompute kernel concurrently with main kernel via PDL). spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( - max_batch_size, dtype=int, device=device) + max_batch_size, dtype=torch.int32, device=device) spec_kwargs['cache_buf_idx'] = torch.zeros(max_batch_size, dtype=torch.int32, device=device) spec_kwargs['old_x'] = torch.zeros(num_local_layers, max_batch_size, - T, + 2, + self.replay_history_size, nheads, head_dim, dtype=dtype, @@ -524,7 +552,7 @@ def __init__( spec_kwargs['old_B'] = torch.zeros(num_local_layers, max_batch_size, 2, - T, + self.replay_history_size, n_groups_per_rank, d_state, dtype=dtype, @@ -533,16 +561,17 @@ def __init__( max_batch_size, 2, nheads, - T, + self.replay_history_size, dtype=torch.float32, device=device) - spec_kwargs['old_dA_cumsum'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - nheads, - T, - dtype=torch.float32, - device=device) + spec_kwargs['old_dA_cumsum'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + dtype=torch.float32, + device=device) ssm_spec_cache = [ spec_kwargs['old_x'], spec_kwargs['old_B'], spec_kwargs['old_dt'], spec_kwargs['old_dA_cumsum'] @@ -571,7 +600,8 @@ def __init__( f"conv_state size: {get_tensor_size_bytes(conv_states) / GB:.2f}GB, " f"ssm_state size: {get_tensor_size_bytes(ssm_states) / GB:.2f}GB, " f"ssm_spec_cache size: {get_tensor_size_bytes(ssm_spec_cache) / GB:.2f}GB, " - f"intermediate_conv_window_cache size: {get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" + "intermediate_conv_window_cache size: " + f"{get_tensor_size_bytes(intermediate_conv_window_cache) / GB:.2f}GB" ) else: self.mamba_cache = self.State( @@ -591,6 +621,15 @@ def __init__( # mamba cache index, maps request_id -> state indices self.mamba_cache_index: Dict[int, int] = {} + self._dummy_request_ids: set[int] = set() + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = torch.zeros(max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros(max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) # Permanent slot shared by every CUDA-graph padding sentinel id # (CUDA_GRAPH_DUMMY_REQUEST_ID - runtime_draft_len, one per @@ -609,12 +648,15 @@ def __init__( dtype=torch.int32, device=device) - # Store max_batch_size for resource management + # Physical tensor rows include reserved dummy slots. Resource capacity + # is the number of real request slots remaining after those + # reservations. self._max_batch_size = max_batch_size + self._max_resource_count = len(self.mamba_cache_free_blocks) def get_max_resource_count(self) -> int: - """Return the maximum number of sequences that can be cached.""" - return self._max_batch_size + """Return the maximum number of real requests that can be cached.""" + return self._max_resource_count def filter_ctx_requests_by_capacity(self, context_requests: list) -> list: """Return the prefix of *context_requests* that fits in the @@ -648,6 +690,7 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): if (isinstance(self.mamba_cache, self.SpeculativeState) and self._use_replay_state_update): self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 if self._mamba_ssm_rand_seed is not None: # Deterministic per-slot rotation on fresh assignment. # `block` is pulled from mamba_cache_free_blocks, which @@ -660,6 +703,8 @@ def _prepare_mamba_cache_blocks(self, request_ids: List[int]): self._seed_rank_offset)) def prepare_resources(self, scheduled_batch: ScheduledRequests): + requests = (scheduled_batch.context_requests + + scheduled_batch.generation_requests) context_ids = [ i.py_request_id for i in scheduled_batch.context_requests ] @@ -668,6 +713,7 @@ def prepare_resources(self, scheduled_batch: ScheduledRequests): ] request_ids = context_ids + generation_ids self._prepare_mamba_cache_blocks(request_ids) + self._refresh_dummy_request_mask([req.is_dummy for req in requests]) def _is_padding_sentinel(self, request_id: int) -> bool: # cuda_graph_runner caches one dummy per runtime_draft_len value @@ -685,24 +731,35 @@ def add_dummy_requests(self, request_ids: List[int], **kwargs): # slot and are freed individually. if not request_ids: return + self._dummy_request_ids.update(request_ids) for r in request_ids: if r in self.mamba_cache_index: + block = self.mamba_cache_index[r] + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 continue if self._is_padding_sentinel(r): - self.mamba_cache_index[r] = self._padding_slot + block = self._padding_slot elif (r == ATTENTION_DP_DUMMY_REQUEST_ID and self._attention_dp_dummy_slot is not None): - self.mamba_cache_index[r] = self._attention_dp_dummy_slot + block = self._attention_dp_dummy_slot else: if len(self.mamba_cache_free_blocks) == 0: raise RuntimeError("run out of mamba cache blocks") block = self.mamba_cache_free_blocks.pop() - self.mamba_cache_index[r] = block + self.mamba_cache_index[r] = block + if (isinstance(self.mamba_cache, self.SpeculativeState) + and self._use_replay_state_update): + self.mamba_cache.prev_num_accepted_tokens[block] = 0 + self.mamba_cache.cache_buf_idx[block] = 0 def free_resources(self, request: LlmRequest): request_id = request.py_request_id if request_id not in self.mamba_cache_index: return + self._dummy_request_ids.discard(request_id) block = self.mamba_cache_index.pop(request_id) # Reserved slots must not re-enter the real-request free pool. if block != self._padding_slot and \ @@ -711,7 +768,25 @@ def free_resources(self, request: LlmRequest): def get_state_indices(self, request_ids: List[int], is_padding: List[bool]) -> List[int]: - return [self.mamba_cache_index[rid] for rid in request_ids] + assert len(request_ids) == len(is_padding) + indices = [self.mamba_cache_index[rid] for rid in request_ids] + is_dummy = [ + rid in self._dummy_request_ids or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices + + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.as_tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) def get_conv_states(self, layer_idx: int) -> torch.Tensor: layer_offset = self.mamba_layer_offsets[layer_idx] @@ -795,7 +870,23 @@ def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: @property def use_replay_state_update(self) -> bool: - return self._use_replay_state_update + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + if (not self._use_replay_state_update + or not isinstance(self.mamba_cache, self.SpeculativeState) + or self.mamba_cache.prev_num_accepted_tokens is None + or self.mamba_cache.cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=( + self.mamba_cache.prev_num_accepted_tokens), + cache_buf_idx=self.mamba_cache.cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) def shutdown(self): """Release tensor memory.""" @@ -838,14 +929,33 @@ def update_mamba_states(self, attn_metadata: "AttentionMetadata", src_state_indices = self.intermediate_state_indices[:num_gens] if self._use_replay_state_update: - # SSM state is handled incrementally by the kernel. Update the - # number of accepted tokens and flip the double-buffer index so the - # next step's replay reads from the buffer that was just written by - # the precompute kernel. + # SSM state is handled incrementally by the kernel. Mirror the + # kernel's per-slot checkpoint predicate from the previous PNAT and + # fixed replay step width: checkpoint steps write a fresh history + # buffer and flip, while no-checkpoint steps append to the active + # buffer and keep reading from it next step. + accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + + num_gens] + prev_num_accepted_tokens = \ + self.mamba_cache.prev_num_accepted_tokens[state_indices_d] + wrote_checkpoint = (prev_num_accepted_tokens + + self.replay_step_width + > self.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, accepted_tokens, + prev_num_accepted_tokens + accepted_tokens) + cache_buf_idx = self.mamba_cache.cache_buf_idx[state_indices_d] + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + next_num_accepted_tokens = torch.where(is_dummy_request, + prev_num_accepted_tokens, + next_num_accepted_tokens) self.mamba_cache.prev_num_accepted_tokens[state_indices_d] = \ - num_accepted_tokens[num_contexts:num_contexts + num_gens] + next_num_accepted_tokens self.mamba_cache.cache_buf_idx[state_indices_d] = \ - 1 - self.mamba_cache.cache_buf_idx[state_indices_d] + torch.where(is_dummy_request, cache_buf_idx, + torch.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx)) else: # Legacy: copy accepted SSM state from intermediate cache. ssm_states = self.mamba_cache.temporal @@ -984,7 +1094,15 @@ def get_mamba_ssm_rand_seed(self) -> Optional[torch.Tensor]: @property def use_replay_state_update(self) -> bool: - return getattr(self._impl, 'use_replay_state_update', False) + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + get_metadata = getattr(self._impl, 'get_replay_state_update_metadata', + None) + if get_metadata is None: + return None + return get_metadata() def get_intermediate_ssm_states(self, layer_idx: int) -> Optional[torch.Tensor]: @@ -1121,7 +1239,9 @@ def __init__( ) -> None: # mamba hybrid cache requires block reuse to be disabled in KV cache config - assert not kv_cache_config.enable_block_reuse, "mamba hybrid cache requires block reuse to be disabled in KV cache config" + assert not kv_cache_config.enable_block_reuse, ( + "mamba hybrid cache requires block reuse to be disabled in KV cache config" + ) pool_size = _get_mamba_hybrid_pool_size(max_batch_size, mapping) @@ -1383,6 +1503,12 @@ def __init__( # accessors (get_mamba_ssm_cache_dtype, use_replay_state_update) work # on ranks with no local mamba layers. self._use_replay_state_update = use_replay_state_update + self.replay_step_width: Optional[int] = ( + spec_config.tokens_per_gen_step + if spec_config is not None and use_replay_state_update else None) + self.replay_history_size: Optional[int] = ( + max(MIN_REPLAY_HISTORY_SIZE, self.replay_step_width) + if self.replay_step_width is not None else None) # Same allocation gate as PythonMambaCacheManager: the rand_seed # buffer must exist whenever SR can fire, not only on the replay path. self._mamba_ssm_stochastic_rounding = mamba_ssm_stochastic_rounding @@ -1470,7 +1596,9 @@ def __init__( self.linear_attention_metadata = LinearAttentionMetadata() self.linear_attention_metadata.cache_type = LinearCacheType.RECURRENT_STATES.value self.linear_attention_metadata.all_recurrent_states_bytes = self.ssm_bytes + self.conv_bytes - self.linear_attention_metadata.states_snapshot_interval = kv_cache_config.mamba_state_cache_interval if kv_cache_config.enable_block_reuse else 0 + self.linear_attention_metadata.states_snapshot_interval = ( + kv_cache_config.mamba_state_cache_interval + if kv_cache_config.enable_block_reuse else 0) # RNN model params for disagg TP-mismatch split/concat. conv_section_map = {"nemotron_hybrid": 1, "qwen3_next": 2} self.linear_attention_metadata.rnn_num_heads = self._rnn_num_heads @@ -1558,6 +1686,11 @@ def __init__( dtype=torch.long, device="cpu") self._request_id_to_state_index = {} + self._request_id_to_is_dummy = {} + # Batch-order mask aligned with state_indices; duplicate dummy request + # IDs mark every batch row even when they share one cache slot. + self._dummy_request_mask = None + self._dummy_request_mask_host = None self.kv_cache_config = kv_cache_config self.is_estimating_kv_cache = is_estimating_kv_cache @@ -1647,6 +1780,8 @@ def shutdown(self): self.prev_num_accepted_tokens = None self.cache_buf_idx = None self.mamba_ssm_rand_seed = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None self.old_x = None self.old_B = None self.old_dt = None @@ -1713,10 +1848,16 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): self._pending_state_transfers = self.impl.copy_linear_attention_block_batch( self.requests) self._setup_state_indices() + # Reset replay double-buffer state for fresh context blocks. A reused + # block (prefix-cache hit or block recycled across requests) may carry + # stale prev_num_accepted_tokens / cache_buf_idx values from a prior + # owner; the replay kernel reads these on the first decode step. num_contexts = len(scheduled_batch.context_requests) if num_contexts > 0: ctx_slots = self.cuda_state_indices[:num_contexts].long() - if self._use_replay_state_update and self.prev_num_accepted_tokens is not None: + if (self._use_replay_state_update + and self.prev_num_accepted_tokens is not None + and self.cache_buf_idx is not None): self.prev_num_accepted_tokens[ctx_slots] = 0 self.cache_buf_idx[ctx_slots] = 0 if self.old_x is not None: @@ -1798,15 +1939,32 @@ def update_mamba_states(self, # writes through the view's real strides (~85% of HBM peak, one launch # per state). if self._use_replay_state_update: - # SSM state is handled incrementally by the replay kernel. Update - # the per-slot accepted-token counter and flip the double-buffer - # index so the next step reads from the buffer that was just - # written by the precompute kernel. + # SSM state is handled incrementally by the kernel. Mirror the + # kernel's checkpoint predicate from the previous PNAT and fixed + # replay step width: checkpoint steps flip buffers, while no-write + # steps append to the active history. slots = state_indices_d.long() - accepted = num_accepted_tokens[num_contexts:num_contexts + num_gens] - self.prev_num_accepted_tokens[slots] = accepted.to( - self.prev_num_accepted_tokens.dtype) - self.cache_buf_idx[slots] = 1 - self.cache_buf_idx[slots] + accepted = num_accepted_tokens[num_contexts:num_contexts + + num_gens].to( + self.prev_num_accepted_tokens. + dtype) + prev_num_accepted_tokens = self.prev_num_accepted_tokens[slots] + wrote_checkpoint = (prev_num_accepted_tokens + + self.replay_step_width + > self.replay_history_size) + next_num_accepted_tokens = torch.where( + wrote_checkpoint, accepted, prev_num_accepted_tokens + accepted) + cache_buf_idx = self.cache_buf_idx[slots] + assert self._dummy_request_mask is not None + is_dummy_request = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + next_num_accepted_tokens = torch.where(is_dummy_request, + prev_num_accepted_tokens, + next_num_accepted_tokens) + self.prev_num_accepted_tokens[slots] = next_num_accepted_tokens + self.cache_buf_idx[slots] = torch.where( + is_dummy_request, cache_buf_idx, + torch.where(wrote_checkpoint, 1 - cache_buf_idx, cache_buf_idx)) else: # Legacy: copy the accepted SSM state from the intermediate buffer. _promote_mamba_state_triton(self.all_ssm_states, @@ -1822,6 +1980,20 @@ def update_mamba_states(self, src_state_indices, num_accepted_draft_tokens, state_indices_d) + @torch.inference_mode() + def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: + if self._dummy_request_mask is None: + return + + n = len(is_dummy) + assert n <= self._dummy_request_mask_host.shape[0] + self._dummy_request_mask_host.zero_() + if n > 0: + self._dummy_request_mask_host[:n].copy_( + torch.tensor(is_dummy, dtype=torch.bool)) + self._dummy_request_mask.copy_(self._dummy_request_mask_host, + non_blocking=True) + def get_num_available_tokens(self, token_num_upper_bound: int, max_num_draft_tokens: int = 0, @@ -1911,6 +2083,7 @@ def free_resources(self, request: LlmRequest, pin_on_release: bool = False): if request in self.requests: self.requests.remove(request) self._request_id_to_state_index.pop(request.py_request_id, None) + self._request_id_to_is_dummy.pop(request.py_request_id, None) super().free_resources(request, pin_on_release) def _setup_state_indices(self) -> None: @@ -1958,18 +2131,22 @@ def _setup_state_indices(self) -> None: f"prepopulated_token_num={req.prepopulated_prompt_len}, " f"context_chunk_size={req.context_chunk_size if not req.is_context_finished else 'N/A'}, " f"block_index for next step is {block_indices[bad_i]}, " - f"\nblock_ids={self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" + "\nblock_ids=" + f"{self.impl.get_cache_block_ids(req.py_request_id, LinearCacheType.RECURRENT_STATES.value)}" ) self._host_state_indices[:n] = values self.cuda_state_indices.copy_(self._host_state_indices, non_blocking=True) + self._refresh_dummy_request_mask( + [req.is_dummy for req in self.requests]) # Build request_id → pool block offset mapping so that # get_state_indices can return indices in arbitrary request order. for i, req in enumerate(self.requests): self._request_id_to_state_index[ req.py_request_id] = self._host_state_indices[i].item() + self._request_id_to_is_dummy[req.py_request_id] = req.is_dummy def get_state_indices(self, request_ids: Optional[List[int]] = None, @@ -1979,7 +2156,18 @@ def get_state_indices(self, # not the internal self.requests order. This is critical when # the batch is reordered after prepare_resources (e.g. disagg # serving sorts generation_requests by py_batch_idx). - return [self._request_id_to_state_index[rid] for rid in request_ids] + indices = [ + self._request_id_to_state_index[rid] for rid in request_ids + ] + if is_padding is None: + is_padding = [False] * len(request_ids) + assert len(request_ids) == len(is_padding) + is_dummy = [ + self._request_id_to_is_dummy.get(rid, False) or padding + for rid, padding in zip(request_ids, is_padding) + ] + self._refresh_dummy_request_mask(is_dummy) + return indices return self.cuda_state_indices def calc_next_context_chunk_size(self, request: LlmRequest) -> int: @@ -2001,7 +2189,9 @@ def calc_next_context_chunk_size(self, request: LlmRequest) -> int: if current >= prompt_len: return 0 if not self.kv_cache_config.enable_block_reuse: - assert current == 0, f"Expected context_current_position to be 0 when block reuse is disabled, but got {current}" + assert current == 0, ( + "Expected context_current_position to be 0 when block reuse is " + f"disabled, but got {current}") return prompt_len - current step = self.linear_attention_metadata.states_snapshot_interval stop_positions = calc_context_stop_positions(prompt_len, @@ -2029,6 +2219,8 @@ def _setup_states(self) -> None: self.local_num_mamba_layers, num_blocks_in_pool ] + self.conv_state_shape) + self.all_ssm_states.zero_() + self.all_conv_states.zero_() def _setup_mtp_intermediate_states(self, spec_config, max_batch_size) -> None: @@ -2104,9 +2296,17 @@ def _setup_replay_buffers(self, spec_config) -> None: # Without spec_config or replay we still keep the seed buffer # (above) so the non-MTP flashinfer SR path has a persistent # rand_seed source. + self.prev_num_accepted_tokens = None + self.cache_buf_idx = None + self.old_x = None + self.old_B = None + self.old_dt = None + self.old_dA_cumsum = None + self._dummy_request_mask = None + self._dummy_request_mask_host = None return - T = spec_config.max_draft_len + 1 + history_size = self.replay_history_size num_local_mamba_layers = self.local_num_mamba_layers nheads, head_dim, d_state = self.ssm_state_shape n_groups_per_rank = self._n_groups_per_rank @@ -2118,10 +2318,16 @@ def _setup_replay_buffers(self, spec_config) -> None: self.cache_buf_idx = torch.zeros(cache_size, dtype=torch.int32, device=device) - # x is not double-buffered + self._dummy_request_mask = torch.zeros(self.max_batch_size, + dtype=torch.bool, + device=device) + self._dummy_request_mask_host = torch.zeros(self.max_batch_size, + dtype=torch.bool, + pin_memory=prefer_pinned()) self.old_x = torch.zeros(num_local_mamba_layers, cache_size, - T, + 2, + history_size, nheads, head_dim, dtype=self.conv_state_dtype, @@ -2130,7 +2336,7 @@ def _setup_replay_buffers(self, spec_config) -> None: self.old_B = torch.zeros(num_local_mamba_layers, cache_size, 2, - T, + history_size, n_groups_per_rank, d_state, dtype=self.conv_state_dtype, @@ -2139,20 +2345,36 @@ def _setup_replay_buffers(self, spec_config) -> None: cache_size, 2, nheads, - T, + history_size, dtype=torch.float32, device=device) self.old_dA_cumsum = torch.zeros(num_local_mamba_layers, cache_size, 2, nheads, - T, + history_size, dtype=torch.float32, device=device) @property def use_replay_state_update(self) -> bool: - return self._use_replay_state_update + return self.get_replay_state_update_metadata() is not None + + def get_replay_state_update_metadata( + self) -> Optional[ReplayStateUpdateMetadata]: + prev_num_accepted_tokens = getattr(self, 'prev_num_accepted_tokens', + None) + cache_buf_idx = getattr(self, 'cache_buf_idx', None) + if (not self._use_replay_state_update + or prev_num_accepted_tokens is None or cache_buf_idx is None + or self.replay_step_width is None + or self.replay_history_size is None): + return None + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=self.replay_step_width, + replay_history_size=self.replay_history_size) def get_mamba_ssm_cache_dtype(self) -> torch.dtype: return self.ssm_state_dtype diff --git a/tensorrt_llm/_torch/pyexecutor/resource_manager.py b/tensorrt_llm/_torch/pyexecutor/resource_manager.py index 183d5ab0c715..e4ad48050606 100644 --- a/tensorrt_llm/_torch/pyexecutor/resource_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/resource_manager.py @@ -93,6 +93,7 @@ class BaseResourceManager(ABC): @abstractmethod def get_max_resource_count(self) -> int: + """Return the maximum number of real requests this manager can admit.""" raise NotImplementedError @abstractmethod diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 7a72ffebc88f..24e350a14297 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7425,7 +7425,7 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, enable_block_reuse=True, mamba_ssm_cache_dtype="float16", mamba_state_cache_interval=mamba_state_cache_interval, - free_gpu_memory_fraction=0.6, + free_gpu_memory_fraction=0.5, ), max_batch_size=max_batch_size, tensor_parallel_size=tp_size, diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 030f5d77f721..5f73a1f86881 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -13,6 +13,7 @@ from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import CUDA_GRAPH_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.llm_request import ATTENTION_DP_DUMMY_REQUEST_ID from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MIN_REPLAY_HISTORY_SIZE, CppMambaCacheManager, CppMambaHybridCacheManager, PythonMambaCacheManager, @@ -27,7 +28,9 @@ skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -def _make_mgr(max_batch_size=4, max_draft_len=2, enable_attention_dp=False): +def _make_mgr( + max_batch_size=4, max_draft_len=2, enable_attention_dp=False, use_replay_state_update=False +): mapping = Mapping(world_size=1, tp_size=1, pp_size=1, enable_attention_dp=enable_attention_dp) pool = _get_mamba_hybrid_pool_size(max_batch_size, mapping) return PythonMambaCacheManager( @@ -43,9 +46,36 @@ def _make_mgr(max_batch_size=4, max_draft_len=2, enable_attention_dp=False): dtype=torch.float16, ssm_cache_dtype=torch.float16, speculative_num_draft_tokens=max_draft_len, + use_replay_state_update=use_replay_state_update, ) +@skip_no_cuda +@pytest.mark.parametrize("enable_attention_dp", [False, True]) +def test_python_mamba_resource_count_excludes_reserved_dummy_slots(enable_attention_dp): + max_batch_size = 4 + mgr = _make_mgr( + max_batch_size=max_batch_size, + max_draft_len=2, + enable_attention_dp=enable_attention_dp, + ) + + assert mgr.get_max_resource_count() == max_batch_size + assert len(mgr.mamba_cache_free_blocks) == max_batch_size + + +@skip_no_cuda +def test_replay_inactive_without_spec_config(): + mgr = _make_mgr( + max_batch_size=2, + max_draft_len=None, + use_replay_state_update=True, + ) + + assert mgr.use_replay_state_update is False + assert mgr.get_replay_state_update_metadata() is None + + @skip_no_cuda def test_padding_slot_not_held_by_parked_real(): """Padding must not resolve to a slot owned by a parked real @@ -117,6 +147,81 @@ def _fake(rid): assert mgr._padding_slot == shared +@skip_no_cuda +def test_replay_update_mamba_states_uses_history_window(): + """Replay path accumulates PNAT until layer kernels write a checkpoint.""" + mgr = _make_mgr(max_batch_size=4, max_draft_len=5, use_replay_state_update=True) + assert mgr.replay_step_width == 6 + assert mgr.replay_history_size == MIN_REPLAY_HISTORY_SIZE + assert mgr.mamba_cache.prev_num_accepted_tokens.dtype == torch.int32 + assert mgr.mamba_cache.cache_buf_idx.dtype == torch.int32 + assert mgr.mamba_cache.old_x.shape[3] == MIN_REPLAY_HISTORY_SIZE + assert mgr.mamba_cache.old_B.shape[3] == MIN_REPLAY_HISTORY_SIZE + assert mgr.mamba_cache.old_dt.shape[4] == MIN_REPLAY_HISTORY_SIZE + assert mgr.mamba_cache.old_dA_cumsum.shape[4] == MIN_REPLAY_HISTORY_SIZE + + mgr._prepare_mamba_cache_blocks([100, 101]) + slot_appended = mgr.mamba_cache_index[100] + slot_checkpointed = mgr.mamba_cache_index[101] + + mgr.mamba_cache.prev_num_accepted_tokens[slot_appended] = 7 + mgr.mamba_cache.prev_num_accepted_tokens[slot_checkpointed] = 13 + mgr.mamba_cache.cache_buf_idx[slot_appended] = 0 + mgr.mamba_cache.cache_buf_idx[slot_checkpointed] = 1 + mgr.mamba_cache.conv.zero_() + mgr.mamba_cache.intermediate_conv_window.zero_() + mgr.mamba_cache.intermediate_conv_window[:, 0, 2] = 11.0 + mgr.mamba_cache.intermediate_conv_window[:, 1, 2] = 13.0 + + state_indices = torch.tensor( + [slot_appended, slot_checkpointed], dtype=torch.int32, device="cuda" + ) + attn = SimpleNamespace(num_seqs=2, num_contexts=0) + mgr.update_mamba_states( + attn, + torch.tensor([3, 3], dtype=torch.int32, device="cuda"), + state_indices=state_indices, + ) + + assert mgr.mamba_cache.prev_num_accepted_tokens[slot_appended].item() == 10 + assert mgr.mamba_cache.prev_num_accepted_tokens[slot_checkpointed].item() == 3 + assert mgr.mamba_cache.cache_buf_idx[slot_appended].item() == 0 + assert mgr.mamba_cache.cache_buf_idx[slot_checkpointed].item() == 0 + assert torch.all(mgr.mamba_cache.conv[:, slot_appended] == 11.0) + assert torch.all(mgr.mamba_cache.conv[:, slot_checkpointed] == 13.0) + + +@skip_no_cuda +def test_replay_update_mamba_states_skips_dummy_slots(): + mgr = _make_mgr(max_batch_size=2, max_draft_len=5, use_replay_state_update=True) + mgr._prepare_mamba_cache_blocks([100]) + mgr.add_dummy_requests([CUDA_GRAPH_DUMMY_REQUEST_ID]) + + real_slot = mgr.mamba_cache_index[100] + dummy_slot = mgr.mamba_cache_index[CUDA_GRAPH_DUMMY_REQUEST_ID] + mgr.mamba_cache.prev_num_accepted_tokens[real_slot] = 13 + mgr.mamba_cache.prev_num_accepted_tokens[dummy_slot] = 13 + mgr.mamba_cache.cache_buf_idx[real_slot] = 1 + mgr.mamba_cache.cache_buf_idx[dummy_slot] = 1 + + state_indices = torch.tensor( + mgr.get_state_indices([100, CUDA_GRAPH_DUMMY_REQUEST_ID], [False, True]), + dtype=torch.int32, + device="cuda", + ) + attn = SimpleNamespace(num_seqs=2, num_contexts=0) + mgr.update_mamba_states( + attn, + torch.tensor([3, 3], dtype=torch.int32, device="cuda"), + state_indices=state_indices, + ) + + assert mgr.mamba_cache.prev_num_accepted_tokens[real_slot].item() == 3 + assert mgr.mamba_cache.prev_num_accepted_tokens[dummy_slot].item() == 13 + assert mgr.mamba_cache.cache_buf_idx[real_slot].item() == 0 + assert mgr.mamba_cache.cache_buf_idx[dummy_slot].item() == 1 + + @skip_no_cuda def test_attention_dp_dummy_has_reserved_slot_with_batch_size_one(): mgr = _make_mgr(max_batch_size=1, max_draft_len=0, enable_attention_dp=True) diff --git a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py deleted file mode 100644 index a755579a8600..000000000000 --- a/tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py +++ /dev/null @@ -1,1044 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Standalone benchmark for replay_selective_state_update (Triton kernel). - -Suitable for nsight-compute (ncu) and nsight-systems (nsys) capture. - -Fixed model config: NVIDIA-Nemotron-3-Super-120B-A12B at TP=8 - nheads=16, head_dim=64, d_state=128, ngroups=1 - -mtp_len is the per-request sequence length processed by replay: in MTP it -equals num_draft_tokens + 1 target token, so --mtp-lengths 6 models 5 drafts -+ 1 target. - -Baseline kernel (--baseline [triton|flashinfer]): - Calls selective_state_update with T=mtp_len tokens and disable_state_update=True, - matching the MTP scoring pass in mamba2_mixer.py exactly. - -Example usage: - # Basic sweep - python benchmark_replay_selective_state_update.py \\ - --batch-sizes 1,2,4 --mtp-lengths 1,4,8 --warmup 5 --iters 20 - - # With CUDA graph (default) and Triton baseline: - python benchmark_replay_selective_state_update.py --baseline \\ - --batch-sizes 1,2,4 --mtp-lengths 5,10,20 - - # nsys capture (NVTX ranges visible in timeline) - nsys profile --capture-range=cudaProfilerApi \\ - python benchmark_replay_selective_state_update.py --profile - - # ncu capture - ncu --target-processes all \\ - python benchmark_replay_selective_state_update.py --profile \\ - --batch-sizes 1 --mtp-lengths 4 --warmup 5 --iters 5 -""" - -import argparse -import importlib -import itertools -import os -import statistics -import sys -from datetime import datetime -from pathlib import Path - -import torch -from einops import repeat - - -def _import_mamba_kernels_fast(): - """Load kernel modules directly (~40s faster than a full tensorrt_llm init). - Use --full-import as the fallback if module dependencies change. - - Strategy: stub the parent packages (tensorrt_llm, tensorrt_llm._torch, - tensorrt_llm._torch.modules) in sys.modules with __path__ set, but do - NOT execute their __init__.py. Then load the leaf kernel modules. - When a kernel body imports e.g. tensorrt_llm._utils.get_sm_version, - Python's machinery resolves it against our stub's __path__ and loads - only _utils.py — skipping the heavy tensorrt_llm package init. - """ - import types - - repo_root = Path(__file__).resolve().parents[5] - trtllm_dir = repo_root / "tensorrt_llm" - mamba_pkg = "tensorrt_llm._torch.modules.mamba" - mamba_dir = trtllm_dir / "_torch" / "modules" / "mamba" - - def _stub_pkg(fqn: str, pkg_dir: Path): - """Register a stub package in sys.modules without running its - __init__.py. Sets __path__ so Python can resolve submodule imports - against the real directory on disk.""" - if fqn in sys.modules: - return - stub = types.ModuleType(fqn) - stub.__path__ = [str(pkg_dir)] - sys.modules[fqn] = stub - - # Stub the parent chain so `from tensorrt_llm._utils import ...` (and - # similar) work without triggering tensorrt_llm/__init__.py. - _stub_pkg("tensorrt_llm", trtllm_dir) - _stub_pkg("tensorrt_llm._torch", trtllm_dir / "_torch") - _stub_pkg("tensorrt_llm._torch.modules", trtllm_dir / "_torch" / "modules") - - def _load(mod_name: str, file_name: str): - fqn = f"{mamba_pkg}.{mod_name}" if mod_name else mamba_pkg - if fqn in sys.modules: - return sys.modules[fqn] - spec = importlib.util.spec_from_file_location( - fqn, - mamba_dir / file_name, - submodule_search_locations=[str(mamba_dir)] if file_name == "__init__.py" else [], - ) - mod = importlib.util.module_from_spec(spec) - sys.modules[fqn] = mod - spec.loader.exec_module(mod) - return mod - - # 1. Package __init__ (defines PAD_SLOT_ID = -1) - _load("", "__init__.py") - # 2. softplus helper (used by both kernel modules) - _load("softplus", "softplus.py") - # 3. The actual kernels - replay_mod = _load("replay_selective_state_update", "replay_selective_state_update.py") - base_mod = _load("selective_state_update", "selective_state_update.py") - conv1d_mod = _load("causal_conv1d_triton", "causal_conv1d_triton.py") - - return ( - replay_mod.replay_selective_state_update, - base_mod.selective_state_update, - conv1d_mod.causal_conv1d_update, - ) - - -def _import_mamba_kernels_full(): - """Import via the standard tensorrt_llm package (slow but safe).""" - from tensorrt_llm._torch.modules.mamba.causal_conv1d_triton import causal_conv1d_update - from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( - replay_selective_state_update, - ) - from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update - - return replay_selective_state_update, selective_state_update, causal_conv1d_update - - -# Use fast import by default; --full-import parsed later but we need the -# functions at module level. Check sys.argv early. -if "--full-import" in sys.argv: - replay_selective_state_update, selective_state_update, causal_conv1d_update = ( - _import_mamba_kernels_full() - ) -else: - try: - replay_selective_state_update, selective_state_update, causal_conv1d_update = ( - _import_mamba_kernels_fast() - ) - except Exception as e: # noqa: BLE001 - exit loudly; don't hide a fast-import regression - print( - f"ERROR: fast import failed ({type(e).__name__}: {e})\n" - "Re-run with --full-import for the slow but stable path, " - "then file a bug or fix _import_mamba_kernels_fast.", - file=sys.stderr, - ) - sys.exit(1) - -# Model config defaults (Nemotron-3-Super-120B full model). -# --tp-size divides nheads and ngroups to get the per-GPU slice. -# TP=1: nheads=128, ngroups=8 -# TP=4: nheads=32, ngroups=2 -# TP=8: nheads=16, ngroups=1 (default) -NHEADS = 128 -HEAD_DIM = 64 -D_STATE = 128 -NGROUPS = 8 -TP_SIZE = 8 # default; overridden by --tp-size - -# L2 flush buffer: ~128 MB — larger than L2 on A100/H100/B200 -_L2_FLUSH_SIZE = 32 * 1024 * 1024 # float32 elements → 128 MB -_l2_flush: torch.Tensor | None = None - - -def _init_l2_flush() -> None: - global _l2_flush - _l2_flush = torch.empty(_L2_FLUSH_SIZE, dtype=torch.float32, device="cuda") - - -def _flush_l2() -> None: - """Evict L2 by writing to a large buffer then synchronising.""" - assert _l2_flush is not None - _l2_flush.fill_(0.0) - torch.cuda.synchronize() - - -# Tensor construction helpers - - -def _build_tensors( - batch: int, - mtp_len: int, - state_dtype: torch.dtype, - act_dtype: torch.dtype, - nheads: int, - head_dim: int, - d_state: int, - ngroups: int, -): - """ - Build all tensors for one benchmark configuration. - - nheads/ngroups are already TP-split (i.e. full_nheads // tp_size). - - Returns: - state0 : (batch, nheads, head_dim, d_state) – initial SSM state - x, dt, B, C : (batch, mtp_len, ...) – token inputs for both kernels - A, dt_bias, D : SSM parameters (float32, tie_hdim strides) - prev_tokens : (batch,) - out_incr : pre-allocated output for replay kernel (batch, mtp_len, nheads, head_dim) - out_base : pre-allocated output for baseline kernel (batch, mtp_len, nheads, head_dim) - intermediate_states_buffer: for baseline kernel (batch, mtp_len, nheads, head_dim, d_state) - """ - device = "cuda" - - torch.manual_seed(42) - - # --- SSM parameters (float32, tie_hdim strides) --- - A_base = -torch.rand(nheads, device=device) - 0.5 - A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) # stride(-1)=0, stride(-2)=0 - - dt_bias_base = torch.randn(nheads, device=device, dtype=torch.float32) - dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) # stride(-1)=0 - - D_base = torch.randn(nheads, device=device, dtype=torch.float32) - D = repeat(D_base, "h -> h p", p=head_dim) - - # --- SSM state --- - state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=state_dtype) - - # --- Cache tensors for replay kernel --- - # old_x: single-buffered (cache, T, nheads, dim) - old_x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - # old_B: double-buffered (cache, 2, T, ngroups, dstate) - old_B = torch.randn(batch, 2, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - # old_dt: double-buffered (cache, 2, nheads, T) fp32 — T contiguous - old_dt = torch.randn(batch, 2, nheads, mtp_len, device=device, dtype=torch.float32) - # old_dA_cumsum: double-buffered (cache, 2, nheads, T) fp32 — T contiguous - old_dA_cumsum = torch.randn(batch, 2, nheads, mtp_len, device=device, dtype=torch.float32) - # cache_buf_idx: which buffer to read (0 or 1) - cache_buf_idx = torch.zeros(batch, device=device, dtype=torch.int32) - - # --- Token inputs (used by both replay and baseline kernels) --- - x = torch.randn(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - # dt must match D's dtype (fp32) for flashinfer — force it for all paths. - dt_base = torch.randn(batch, mtp_len, nheads, device=device, dtype=torch.float32) - dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) # tie_hdim - B = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - C = torch.randn(batch, mtp_len, ngroups, d_state, device=device, dtype=act_dtype) - - # prev_tokens placeholder — overwritten per-run - prev_tokens = torch.zeros(batch, device=device, dtype=torch.int32) - - out_incr = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - out_base = torch.zeros(batch, mtp_len, nheads, head_dim, device=device, dtype=act_dtype) - - intermediate_states_buffer = torch.zeros( - batch, mtp_len, nheads, head_dim, d_state, device=device, dtype=state_dtype - ) - - # --- Conv1d tensors (for --with-conv1d mode) --- - d_inner = nheads * head_dim - conv_dim = d_inner + 2 * ngroups * d_state - d_conv = 4 # conv kernel width for Nemotron/Mamba2 - - # xbc_input: (batch, conv_dim, mtp_len) — "hot" input from in_proj. - # Match production layout: in_proj output is (batch*mtp_len, conv_dim) - # contiguous, then .view(batch, mtp_len, conv_dim).transpose(1, 2) - # gives strides (mtp_len*conv_dim, 1, conv_dim) — NOT the standard - # (conv_dim*mtp_len, mtp_len, 1) of a freshly allocated 3D tensor. - # Conv1d preserves input strides in its output, so downstream split - # + view inherits the correct layout without needing .contiguous(). - xbc_input_flat = torch.randn(batch * mtp_len, conv_dim, device=device, dtype=act_dtype) - xbc_input = xbc_input_flat.view(batch, mtp_len, conv_dim).transpose(1, 2) - # conv_state: (batch, conv_dim, d_conv) — "cold" cache - conv_state = torch.randn(batch, conv_dim, d_conv, device=device, dtype=act_dtype) - # conv_weight: (conv_dim, d_conv) — parameter - conv_weight = torch.randn(conv_dim, d_conv, device=device, dtype=act_dtype) - # conv_bias: (conv_dim,) — parameter - conv_bias = torch.randn(conv_dim, device=device, dtype=act_dtype) - - return ( - state0, - old_x, - old_B, - old_dt, - old_dA_cumsum, - cache_buf_idx, - x, - dt, - B, - C, - A, - dt_bias, - D, - prev_tokens, - out_incr, - out_base, - intermediate_states_buffer, - xbc_input, - conv_state, - conv_weight, - conv_bias, - d_inner, - conv_dim, - ) - - -# Timing helpers - - -def _compute_stats(latencies_us: list[float]) -> tuple[float, float, float]: - """Return (median_us, p95_us, p99_us) from a list of latencies.""" - median_us = statistics.median(latencies_us) - s = sorted(latencies_us) - p95_us = s[int(0.95 * len(s))] - p99_us = s[int(0.99 * len(s))] - return median_us, p95_us, p99_us - - -def _time_kernel_cuda_graph( - args, - run_fn, - reset_fn, - tag: str, -) -> tuple[float, float, float]: - """ - All-in-one CUDA graph timing. - - Captures a single graph containing warmup iterations followed by timed - iterations with per-iteration event pairs recorded inside the graph. - One replay, one sync, then all timings are read. - """ - warmup = args.warmup - iters = args.iters - - start_events = [torch.cuda.Event(enable_timing=True, external=True) for _ in range(iters)] - end_events = [torch.cuda.Event(enable_timing=True, external=True) for _ in range(iters)] - - # Eager warmup before graph capture (triggers Triton autotune if active) - reset_fn() - run_fn() - torch.cuda.synchronize() - - reset_fn() - torch.cuda.synchronize() - - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g): - # Warmup iterations (unrolled into the graph) - for _ in range(warmup): - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - run_fn() - - # Timed iterations with events inside the graph - for i in range(iters): - reset_fn() - if args.l2_flush: - _l2_flush.fill_(0.0) - start_events[i].record() - run_fn() - end_events[i].record() - - torch.cuda.synchronize() - - # Single replay - torch.cuda.nvtx.range_push(tag) - g.replay() - torch.cuda.synchronize() - torch.cuda.nvtx.range_pop() - - latencies_us = [start_events[i].elapsed_time(end_events[i]) * 1000.0 for i in range(iters)] - return _compute_stats(latencies_us) - - -def _time_kernel_eager( - args, - run_fn, - reset_fn, - tag: str, -) -> tuple[float, float, float]: - """Non-CUDA-graph timing path (for debugging, ncu, etc.).""" - # Warmup - for _ in range(args.warmup): - reset_fn() - run_fn() - torch.cuda.synchronize() - - start_event = torch.cuda.Event(enable_timing=True) - end_event = torch.cuda.Event(enable_timing=True) - - latencies_us: list[float] = [] - torch.cuda.nvtx.range_push(tag) - for _ in range(args.iters): - reset_fn() - if args.l2_flush: - _flush_l2() # includes synchronize - start_event.record() - run_fn() - end_event.record() - torch.cuda.synchronize() - latencies_us.append(start_event.elapsed_time(end_event) * 1000.0) - torch.cuda.nvtx.range_pop() - - return _compute_stats(latencies_us) - - -def _time_kernel(args, run_fn, reset_fn, tag: str) -> tuple[float, float, float]: - """Dispatch to CUDA-graph or eager timing path.""" - if args.cuda_graph: - return _time_kernel_cuda_graph(args, run_fn, reset_fn, tag) - return _time_kernel_eager(args, run_fn, reset_fn, tag) - - -# Per-config benchmark (consolidated baseline + replay) - - -def _bench_config( - args, - batch: int, - mtp_len: int, - prev_ks: list[int], - state_dtype: torch.dtype, - act_dtype: torch.dtype, - baseline_fn, -) -> None: - """ - Benchmark one (batch, mtp_len, dtype) configuration. - - Runs the baseline kernel (if baseline_fn is not None) followed by the - replay kernel for each prev_k value. Tensors are built once and - shared across all runs in this config. - """ - state_dtype_name = str(state_dtype).split(".")[-1] - act_dtype_name = str(act_dtype).split(".")[-1] - - ( - state0, - old_x0, - old_B0, - old_dt0, - old_dA_cumsum0, - cache_buf_idx0, - x, - dt, - B, - C, - A, - dt_bias, - D, - prev_tokens, - out_incr, - out_base, - intermediate_states_buffer, - xbc_input0, - conv_state0, - conv_weight, - conv_bias, - d_inner, - conv_dim, - ) = _build_tensors( - batch, - mtp_len, - state_dtype, - act_dtype, - args.tp_nheads, - args.head_dim, - args.d_state, - args.tp_ngroups, - ) - - nheads = args.tp_nheads - ngroups = args.tp_ngroups - head_dim = args.head_dim - d_state = args.d_state - with_conv1d = getattr(args, "with_conv1d", False) - use_philox = getattr(args, "philox_rounding", False) - - # Philox rounding: allocate rand_seed tensor - rand_seed = None - if use_philox: - if state_dtype != torch.float16: - raise ValueError(f"--philox-rounding requires --state-dtypes fp16, got {state_dtype}") - if args.baseline == "triton": - raise ValueError( - "--philox-rounding not supported with --baseline triton " - "(only flashinfer and replay support it)" - ) - rand_seed = torch.randint(0, 2**62, (1,), device="cuda", dtype=torch.int64) - - state_work = state0.clone() - old_x_work = old_x0.clone() - old_B_work = old_B0.clone() - old_dt_work = old_dt0.clone() - old_dA_cumsum_work = old_dA_cumsum0.clone() - cache_buf_idx_work = cache_buf_idx0.clone() - xbc_input_work = xbc_input0.clone() - conv_state_work = conv_state0.clone() - - def _reset(): - state_work.copy_(state0) - old_x_work.copy_(old_x0) - old_B_work.copy_(old_B0) - old_dt_work.copy_(old_dt0) - old_dA_cumsum_work.copy_(old_dA_cumsum0) - cache_buf_idx_work.copy_(cache_buf_idx0) - if with_conv1d: - conv_state_work.copy_(conv_state0) - - def _reset_conv1d_realistic(): - """Realistic reset: cold cache, L2 flush, then hot in_proj output.""" - # 1. Reset cold state (cache tensors, SSM state) - state_work.copy_(state0) - old_x_work.copy_(old_x0) - old_B_work.copy_(old_B0) - old_dt_work.copy_(old_dt0) - old_dA_cumsum_work.copy_(old_dA_cumsum0) - cache_buf_idx_work.copy_(cache_buf_idx0) - conv_state_work.copy_(conv_state0) - # 2. L2 flush (evicts cold state from cache) - if _l2_flush is not None: - _l2_flush.fill_(0.0) - # 3. Write hot tensors (simulates in_proj output landing in L2) - xbc_input_work.copy_(xbc_input0) - - show_kernel_col = baseline_fn is not None - - def _conv1d_split(xbc_in, conv_st, launch_dependent_kernels=False): - """Run conv1d update and split output into (x, B, C) views. - - The input tensor's strides are preserved through conv1d and the - transpose+view chain. With the production-matching layout - (contiguous (batch*T, conv_dim) viewed as (batch, conv_dim, T)), - the output after transpose+view has stride(-1)==1 and - stride(1)==dim, satisfying both our kernel and flashinfer. - """ - xbc_result = causal_conv1d_update( - xbc_in, - conv_st, - conv_weight, - conv_bias, - activation="silu", - launch_dependent_kernels=launch_dependent_kernels, - ) - xbc_flat = xbc_result.transpose(1, 2).view(batch * mtp_len, conv_dim) - x_flat, B_flat, C_flat = torch.split( - xbc_flat, [d_inner, ngroups * d_state, ngroups * d_state], dim=-1 - ) - x_conv = x_flat.view(batch, mtp_len, nheads, head_dim) - B_conv = B_flat.view(batch, mtp_len, ngroups, d_state) - C_conv = C_flat.view(batch, mtp_len, ngroups, d_state) - return x_conv, B_conv, C_conv - - # --- Baseline --- - if baseline_fn is not None: - tag = f"base_b{batch}_mtp{mtp_len}_s{state_dtype_name}_a{act_dtype_name}" - - philox_kwargs = {} - if rand_seed is not None and args.baseline == "flashinfer": - philox_kwargs = {"rand_seed": rand_seed, "philox_rounds": 10} - - if with_conv1d: - - def _run_baseline(): - x_conv, B_conv, C_conv = _conv1d_split(xbc_input_work, conv_state_work) - baseline_fn( - state_work, - x=x_conv, - dt=dt, - A=A, - B=B_conv, - C=C_conv, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - out=out_base, - disable_state_update=True, - intermediate_states_buffer=intermediate_states_buffer, - cache_steps=mtp_len, - **philox_kwargs, - ) - else: - - def _run_baseline(): - baseline_fn( - state_work, - x=x, - dt=dt, - A=A, - B=B, - C=C, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - out=out_base, - disable_state_update=True, - intermediate_states_buffer=intermediate_states_buffer, - cache_steps=mtp_len, - **philox_kwargs, - ) - - reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - median_us, p95_us, p99_us = _time_kernel(args, _run_baseline, reset_fn, tag) - - _print_row( - show_kernel_col, - args.baseline, - batch, - mtp_len, - "N/A", - state_dtype_name, - act_dtype_name, - median_us, - p95_us, - p99_us, - ) - - # --- Sweep parameter parsing (invariant across prev_k) --- - def _parse_sweep(val): - if val is None: - return [None] - return [int(v) for v in val.split(",")] - - block_size_m_values = _parse_sweep(args.block_size_m) - num_warps_values = _parse_sweep(args.num_warps) - num_stages_values = _parse_sweep(args.num_stages) - precompute_num_warps_values = _parse_sweep(args.precompute_num_warps) - precompute_num_stages_values = _parse_sweep(args.precompute_num_stages) - heads_per_block_values = _parse_sweep(args.heads_per_block) - - # --- Replay kernel, one row per prev_k --- - for prev_k in prev_ks: - prev_tokens.fill_(prev_k) - tag = f"incr_b{batch}_mtp{mtp_len}_k{prev_k}_s{state_dtype_name}_a{act_dtype_name}" - - for ( - block_size_m, - num_warps, - num_stages, - precompute_num_warps, - precompute_num_stages, - heads_per_block, - ) in itertools.product( - block_size_m_values, - num_warps_values, - num_stages_values, - precompute_num_warps_values, - precompute_num_stages_values, - heads_per_block_values, - ): - - def _run_incr( - prev_k=prev_k, - block_size_m=block_size_m, - num_warps=num_warps, - num_stages=num_stages, - precompute_num_warps=precompute_num_warps, - precompute_num_stages=precompute_num_stages, - heads_per_block=heads_per_block, - ): - if with_conv1d: - x_call, B_call, C_call = _conv1d_split( - xbc_input_work, conv_state_work, launch_dependent_kernels=args.external_pdl - ) - extra_kwargs = {"launch_with_pdl": args.external_pdl} - else: - x_call, B_call, C_call = x, B, C - extra_kwargs = {} - replay_selective_state_update( - state_work, - old_x_work, - old_B_work, - old_dt_work, - old_dA_cumsum_work, - cache_buf_idx_work, - prev_tokens, - x=x_call, - dt=dt, - A=A, - B=B_call, - C=C_call, - out=out_incr, - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=None, - rand_seed=rand_seed, - use_internal_pdl=args.internal_pdl, - _block_size_m=block_size_m, - _num_warps=num_warps, - _num_stages=num_stages, - _precompute_num_warps=precompute_num_warps, - _precompute_num_stages=precompute_num_stages, - _heads_per_block=heads_per_block, - **extra_kwargs, - ) - - parts = [] - if block_size_m is not None: - parts.append(f"M={block_size_m}") - if num_warps is not None: - parts.append(f"W={num_warps}") - if num_stages is not None: - parts.append(f"S={num_stages}") - if precompute_num_warps is not None: - parts.append(f"pW={precompute_num_warps}") - if precompute_num_stages is not None: - parts.append(f"pS={precompute_num_stages}") - if heads_per_block is not None: - parts.append(f"H={heads_per_block}") - sweep_suffix = (" " + ",".join(parts)) if parts else "" - sweep_tag = tag + sweep_suffix.replace(" ", "_").replace(",", "_") - - reset_fn = _reset_conv1d_realistic if with_conv1d else _reset - median_us, p95_us, p99_us = _time_kernel(args, _run_incr, reset_fn, sweep_tag) - - _print_row( - show_kernel_col, - "replay", - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - median_us, - p95_us, - p99_us, - sweep_suffix, - ) - - -def _print_row( - show_kernel_col, - kernel_name, - batch, - mtp_len, - prev_k, - state_dtype_name, - act_dtype_name, - median_us, - p95_us, - p99_us, - sweep_suffix="", -): - kernel_col = f"{kernel_name:>11} | " if show_kernel_col else "" - print( - f"| {kernel_col}{batch:>5} | {mtp_len:>7} | {str(prev_k):>6} | " - f"{state_dtype_name:>11} | {act_dtype_name:>9} | " - f"{median_us:>9.2f} | {p95_us:>7.2f} | {p99_us:>7.2f} |" - f"{sweep_suffix}" - ) - - -# Main benchmark loop - - -def _run_benchmark(args) -> None: - assert args.nheads % args.tp_size == 0, ( - f"nheads ({args.nheads}) must be divisible by tp_size ({args.tp_size})" - ) - assert args.ngroups % args.tp_size == 0, ( - f"ngroups ({args.ngroups}) must be divisible by tp_size ({args.tp_size})" - ) - args.tp_nheads = args.nheads // args.tp_size - args.tp_ngroups = args.ngroups // args.tp_size - - batch_sizes = [int(x) for x in args.batch_sizes.split(",")] - mtp_lengths = [int(x) for x in args.mtp_lengths.split(",")] - - dtype_map = {"bf16": torch.bfloat16, "fp32": torch.float32, "fp16": torch.float16} - state_dtypes = [dtype_map[s] for s in args.state_dtypes.split(",")] - act_dtypes = [dtype_map[s] for s in args.act_dtypes.split(",")] - - # Resolve baseline function - if args.baseline == "flashinfer": - from flashinfer.mamba import selective_state_update as baseline_fn - elif args.baseline == "triton": - baseline_fn = selective_state_update - else: - baseline_fn = None - - # --with-conv1d uses its own realistic L2 flush (cold cache flush then - # hot in_proj write). Override the generic l2_flush to avoid double-flushing. - if args.with_conv1d: - args.l2_flush = False - _init_l2_flush() # still needed for the realistic reset's flush step - elif args.l2_flush: - _init_l2_flush() - - if args.profile: - torch.cuda.cudart().cudaProfilerStart() - - # Print header - if baseline_fn is not None: - print( - f"| {'kernel':>11} | {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " - f"{'state_dtype':>11} | {'act_dtype':>9} | " - f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" - ) - print( - f"|{'-' * 13}|{'-' * 7}|{'-' * 9}|{'-' * 8}|" - f"{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" - ) - else: - print( - f"| {'batch':>5} | {'mtp_len':>7} | {'prev_k':>6} | " - f"{'state_dtype':>11} | {'act_dtype':>9} | " - f"{'median_us':>9} | {'p95_us':>7} | {'p99_us':>7} |" - ) - print( - f"|{'-' * 7}|{'-' * 9}|{'-' * 8}|{'-' * 13}|{'-' * 11}|{'-' * 11}|{'-' * 9}|{'-' * 9}|" - ) - - for batch in batch_sizes: - for mtp_len in mtp_lengths: - # Resolve prev_k fractions → clamped integers in [0, mtp_len] - prev_ks = sorted( - set(min(mtp_len, max(0, round(f * mtp_len))) for f in args.prev_tokens_fracs) - ) - for state_dtype in state_dtypes: - for act_dtype in act_dtypes: - _bench_config( - args, batch, mtp_len, prev_ks, state_dtype, act_dtype, baseline_fn - ) - - if args.profile: - torch.cuda.cudart().cudaProfilerStop() - - -# CLI - - -def _parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Benchmark replay_selective_state_update Triton kernel", - formatter_class=argparse.ArgumentDefaultsHelpFormatter, - ) - parser.add_argument( - "--nheads", - type=int, - default=NHEADS, - help="Full-model nheads (divided by --tp-size for per-GPU slice)", - ) - parser.add_argument( - "--ngroups", - type=int, - default=NGROUPS, - help="Full-model ngroups (divided by --tp-size for per-GPU slice)", - ) - parser.add_argument( - "--head-dim", type=int, default=HEAD_DIM, help="Head dimension (not TP-split)" - ) - parser.add_argument( - "--d-state", type=int, default=D_STATE, help="SSM state dimension (not TP-split)" - ) - parser.add_argument( - "--tp-size", - type=int, - default=TP_SIZE, - help="Tensor parallel size; divides nheads and ngroups", - ) - parser.add_argument( - "--batch-sizes", default="1,2,4,8", help="Comma-separated decode batch sizes" - ) - parser.add_argument( - "--mtp-lengths", - default="1,2,4,8", - help="Comma-separated per-request sequence lengths (num_draft_tokens + 1 target)", - ) - parser.add_argument( - "--state-dtypes", default="fp32", help="Comma-separated state dtypes: fp16,bf16,fp32" - ) - parser.add_argument( - "--act-dtypes", - default="bf16", - help="Comma-separated activation dtypes for x/B/C/dt: fp32,bf16", - ) - parser.add_argument("--warmup", type=int, default=20, help="Number of warmup iterations") - parser.add_argument("--iters", type=int, default=100, help="Number of timed iterations") - parser.add_argument( - "--profile", - action="store_true", - help="Wrap timed region in cudaProfilerStart/Stop (for ncu --target-processes all)", - ) - parser.add_argument( - "--l2-flush", - action=argparse.BooleanOptionalAction, - default=True, - help="L2 eviction between iterations", - ) - parser.add_argument( - "--cuda-graph", - action=argparse.BooleanOptionalAction, - default=True, - help="Capture all warmup + timed iterations in a " - "single CUDA graph with per-iteration events " - "inside the graph, eliminating all host overhead.", - ) - parser.add_argument( - "--prev-tokens-fracs", - default="0,0.5,1.0", - type=lambda s: [float(x) for x in s.split(",")], - help="Fractions of mtp_len to use as prev_num_accepted_tokens " - "for the replay kernel sweep. Values are rounded " - "and clamped to [0, mtp_len].", - ) - parser.add_argument( - "--baseline", - default=None, - nargs="?", - const="triton", - choices=[None, "triton", "flashinfer"], - help="Baseline to benchmark alongside the replay kernel. " - "'triton': native Triton selective_state_update. " - "'flashinfer': flashinfer selective_state_update (same signature). " - "Pass --baseline alone for 'triton'. Default: no baseline.", - ) - parser.add_argument( - "--output", - default=None, - help="Path to save results (file or directory). " - "If a directory, writes benchmark_replay_.txt inside it.", - ) - parser.add_argument( - "--block-size-m", - type=str, - default=None, - help="Override BLOCK_SIZE_M: single value or comma-separated sweep (e.g. '4,8,16,32').", - ) - parser.add_argument( - "--num-warps", - type=str, - default=None, - help="Override num_warps: single value or comma-separated sweep (e.g. '1,2,4').", - ) - parser.add_argument( - "--internal-pdl", - action=argparse.BooleanOptionalAction, - default=True, - help="Internal PDL between precompute and main kernels (default: on).", - ) - parser.add_argument( - "--num-stages", - type=str, - default=None, - help="Override num_stages for the main kernel (comma-separated sweep).", - ) - parser.add_argument( - "--precompute-num-warps", - type=str, - default=None, - help="Override num_warps for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--precompute-num-stages", - type=str, - default=None, - help="Override num_stages for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--with-conv1d", - action="store_true", - help="Include conv1d kernel before replay SSM. " - "Uses realistic L2 flush: cold caches flushed, hot in_proj output " - "kept warm. Measures conv1d → precompute → main span.", - ) - parser.add_argument( - "--external-pdl", - action=argparse.BooleanOptionalAction, - default=True, - help="External PDL: conv1d launches dependents, precompute waits. " - "Only relevant with --with-conv1d. --no-external-pdl disables.", - ) - parser.add_argument( - "--heads-per-block", - type=str, - default=None, - help="Override HEADS_PER_BLOCK for precompute kernel (comma-separated sweep).", - ) - parser.add_argument( - "--philox-rounding", - action="store_true", - help="Enable Philox stochastic rounding for fp16 state " - "(rand_seed generated per iteration, philox_rounds=10).", - ) - parser.add_argument( - "--full-import", - action="store_true", - help="Use standard tensorrt_llm import path instead of fast direct " - "module loading. Slower (~40s startup) but guaranteed correct " - "if the fast path breaks due to package changes.", - ) - return parser.parse_args() - - -class _Tee: - """Write to both stdout and a file simultaneously.""" - - def __init__(self, path: str): - parent = os.path.dirname(path) - if parent: - os.makedirs(parent, exist_ok=True) - self._file = open(path, "w") # noqa: SIM115 - self._stdout = sys.stdout - - def write(self, data): - self._stdout.write(data) - self._file.write(data) - - def flush(self): - self._stdout.flush() - self._file.flush() - - def close(self): - self._file.close() - - -if __name__ == "__main__": - _args = _parse_args() - - _out_path = None - if _args.output != "-": - _ts = datetime.now().strftime("%Y%m%d_%H%M%S") - _fname = f"benchmark_replay_{_ts}.txt" - if _args.output is None: - _out_path = os.path.expanduser(f"~/nemo_logs/{_fname}") - elif os.path.isdir(_args.output) or _args.output.endswith("/"): - _out_path = os.path.join(_args.output, _fname) - else: - _out_path = _args.output - - if _out_path: - _tee = _Tee(_out_path) - sys.stdout = _tee - print(f"# benchmark_replay_selective_state_update {datetime.now().isoformat()}") - print(f"# cmd: {' '.join(sys.argv)}") - - try: - _run_benchmark(_args) - finally: - if _out_path: - sys.stdout = _tee._stdout - _tee.close() - print(f"\nResults saved to: {_out_path}") diff --git a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py index c44e0d7f9f87..8aa2a847802b 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba2_metadata.py @@ -20,10 +20,18 @@ import torch from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, Mamba2Metadata, cu_seqlens_to_chunk_indices_offsets, cu_seqlens_to_chunk_indices_offsets_triton, ) +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import ( + MIN_REPLAY_HISTORY_SIZE, + ReplayStateUpdateMetadata, +) skip_no_cuda = pytest.mark.skipif( not torch.cuda.is_available(), @@ -89,6 +97,64 @@ def test_prepare_handles_tensor_cached_tokens(self): assert metadata.chunk_indices is not None assert metadata.chunk_offsets is not None + def test_prepare_replay_work_items_write_first(self): + class ReplayCacheManager: + use_replay_state_update = True + + def __init__(self): + self.state_indices = [0, 3, 1, 4, 2] + self.prev_num_accepted_tokens = torch.tensor( + [0, 4, 10, 11, 20], dtype=torch.int32, device="cuda" + ) + self.cache_buf_idx = torch.tensor([0, 1, 0, 1, 0], dtype=torch.int32, device="cuda") + + def get_state_indices(self, request_ids, is_padding): + return self.state_indices[: len(request_ids)] + + def get_replay_state_update_metadata(self): + return ReplayStateUpdateMetadata( + prev_num_accepted_tokens=self.prev_num_accepted_tokens, + cache_buf_idx=self.cache_buf_idx, + replay_step_width=6, + replay_history_size=MIN_REPLAY_HISTORY_SIZE, + ) + + metadata = Mamba2Metadata(max_batch_size=5, chunk_size=8) + seq_lens = torch.tensor([2, 7, 7, 7, 7], dtype=torch.int) + attn_metadata = SimpleNamespace( + seq_lens=seq_lens, + seq_lens_cuda=seq_lens.cuda(), + num_contexts=1, + num_ctx_tokens=2, + kv_cache_manager=ReplayCacheManager(), + request_ids=[10, 11, 12, 13, 14], + kv_cache_params=SimpleNamespace( + num_cached_tokens_per_seq=torch.tensor([0], dtype=torch.int), + ), + ) + + metadata.prepare(attn_metadata) + + expected = torch.tensor( + [ + [0, 3, 11, 1], + [2, 4, 20, 0], + [1, 1, 4, 1], + [3, 2, 10, 0], + ], + dtype=torch.int32, + device="cuda", + ) + actual = metadata.replay_work_items[:4] + torch.testing.assert_close(actual, expected) + torch.testing.assert_close( + metadata.replay_n_writes.cpu(), torch.tensor([2], dtype=torch.int32) + ) + assert actual[0, REPLAY_WORK_POSITION_IN_DECODE_BATCH] == 0 + assert actual[0, REPLAY_WORK_CACHE_SLOT] == 3 + assert actual[0, REPLAY_WORK_PNAT] == 11 + assert actual[0, REPLAY_WORK_CACHE_BUF_IDX] == 1 + def test_single_sequence_unaligned(self): """Test with a single sequence that doesn't align with chunk size.""" cu_seqlens = torch.tensor([0, 10], dtype=torch.int, device="cuda") diff --git a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py index f94a5f7a28cf..04f1ff00bcf3 100644 --- a/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py +++ b/tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py @@ -1,13 +1,11 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Unit tests for the Mamba SSM stochastic-rounding Philox seed plumbing. - -The Mamba SSM SR path previously generated `rand_seed` tensors via -`torch.randint(..., (1,))` on every decode forward. The cache manager now -owns a persistent per-cache-slot int64 buffer that is deterministically -initialized and rewritten on fresh request assignment. These tests pin the -contract: pure-function seed generation, deterministic allocation, and -per-slot reset without `torch.randint`. +"""Unit tests for Mamba SSM stochastic-rounding Philox seed plumbing. + +The cache manager owns a persistent per-cache-slot int64 seed buffer that is +deterministically initialized and rewritten on fresh request assignment. These +tests pin pure-function seed generation, deterministic allocation, and per-slot +reset without per-forward `torch.randint`. """ import pytest @@ -149,8 +147,7 @@ def test_padding_sentinel_does_not_churn_seeds(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") def test_replay_path_still_allocates_seed_buffer(): - # Backward-compatibility: the replay path used to allocate the seed - # buffer; the new wiring must not regress that. + # Replay stochastic rounding uses the persistent per-slot seed buffer. mgr = _make_python_manager(sr=False, replay=True) seed_buf = mgr.get_mamba_ssm_rand_seed() assert seed_buf is not None @@ -205,9 +202,7 @@ def test_cpp_hybrid_non_replay_mtp_layer_cache_carries_rand_seed(): mamba_ssm_rand_seed on the returned SpeculativeState. The mixer's non-replay MTP SR branch (mamba2_mixer.py) reads - `layer_cache.mamba_ssm_rand_seed` and asserts non-None. Iter5 review - caught the regression where the seed was only forwarded inside the - replay branch of mamba_layer_cache; this test pins both paths.""" + `layer_cache.mamba_ssm_rand_seed` and asserts non-None.""" from tensorrt_llm.llmapi.llm_args import MTPDecodingConfig spec_config = MTPDecodingConfig(max_draft_len=2) diff --git a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py index 7757f83ab2ad..dadebbd996db 100644 --- a/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py +++ b/tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py @@ -13,17 +13,63 @@ # See the License for the specific language governing permissions and # limitations under the License. +import math + import pytest import torch import torch.nn.functional as F +import triton +import triton.language as tl from einops import repeat +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( + _stochastic_round_int8_packed, + _stochastic_round_int16_packed, replay_selective_state_update, ) from tensorrt_llm._torch.modules.mamba.selective_state_update import selective_state_update from tensorrt_llm._utils import get_sm_version + +def _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + explicit_order=None, +): + """Build the replay metadata consumed by persistent_main.""" + position_in_decode_batch = torch.arange(batch, device=device, dtype=torch.int32) + cache_slot = state_batch_indices[:batch].to(torch.int32) + cache_slot_long = cache_slot.to(torch.long) + pnat = prev_tokens[cache_slot_long].to(torch.int32) + active_cache_buf_idx = cache_buf_idx[cache_slot_long].to(torch.int32) + write_mask = (pnat + T) > max_window + n_writes = write_mask.sum().to(torch.int32).reshape(1) + + if explicit_order is None: + order = torch.argsort((~write_mask).to(torch.int32), stable=True).to(torch.long) + else: + order = torch.tensor(explicit_order, device=device, dtype=torch.long) + + replay_work_items = torch.empty(batch, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32) + replay_work_items[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = position_in_decode_batch[order] + replay_work_items[:, REPLAY_WORK_CACHE_SLOT] = cache_slot[order] + replay_work_items[:, REPLAY_WORK_PNAT] = pnat[order] + replay_work_items[:, REPLAY_WORK_CACHE_BUF_IDX] = active_cache_buf_idx[order] + return n_writes, replay_work_items.contiguous() + + # Philox stochastic rounding uses PTX cvt.rs.f16x2.f32 which requires sm >= 100. _skip_pre_sm100 = pytest.mark.skipif( get_sm_version() < 100, reason="Philox stochastic rounding needs sm >= 100" @@ -38,16 +84,231 @@ (16, 64, 128, 1), # TP=8 production config (32, 64, 128, 2), # TP=4, ngroups>1 (more heads than B/C groups) ] +_HEADS_PER_BLOCK_CONFIGS = _CONFIGS + [ + (6, 64, 128, 2), # heads_per_group=3 exercises HPB divisor fallback +] + +# Quantized state dtypes and their representable-magnitude limits (== QUANT_MAX +# in the kernel). fp8_e4m3fn cells require SM 89+ for the fp32↔fp8 cvt PTX +# instructions; SR variants of fp16/fp8 additionally need SM 100+. +_QUANT_MAX_BY_DTYPE = { + torch.int8: 127.0, + torch.int16: 32767.0, + torch.float8_e4m3fn: 448.0, +} + + +def _quantize_state(state_fp32: torch.Tensor, state_dtype: torch.dtype, quant_max: float): + """Quantize fp32 state to (state_quant, decode_scale) using the same + per-(head, dim) channel scheme the kernel does on store. decode_scale = + max_abs_per_channel / quant_max (= 1/encode_scale). + """ + amax = state_fp32.abs().amax(dim=-1) # (cache, nheads, head_dim) + encode_scale = quant_max / amax.clamp(min=1e-30) + decode_scale = 1.0 / encode_scale + scaled = state_fp32 * encode_scale.unsqueeze(-1) + if state_dtype == torch.float8_e4m3fn: + # Native cast does RN at the fp8 grid; explicit round() would destroy + # sub-integer precision (matches the kernel's fp8 RN path). + state_quant = scaled.clamp(-quant_max, quant_max).to(state_dtype) + else: + state_quant = scaled.round().clamp(-quant_max, quant_max).to(state_dtype) + return state_quant, decode_scale + + +def _dequantize_state(state_quant: torch.Tensor, decode_scale: torch.Tensor): + return state_quant.to(torch.float32) * decode_scale.unsqueeze(-1) + + +def _make_strided_state_with_slot_gap(state: torch.Tensor, slot_gap_rows: int): + cache_size, nheads, head_dim, d_state = state.shape + ssm_rows_per_slot = nheads * head_dim + backing = torch.empty( + cache_size, + ssm_rows_per_slot + slot_gap_rows, + d_state, + device=state.device, + dtype=state.dtype, + ) + backing[:, ssm_rows_per_slot:].fill_(float("nan")) + strided_state = backing[:, :ssm_rows_per_slot].view_as(state) + strided_state.copy_(state) + return strided_state, backing + + +def _maybe_skip_dtype(state_dtype, use_sr): + """Skip on insufficient SM. fp8 e4m3fn (any) needs SM 89+; fp16/fp8 SR + needs SM 100+; int8/int16 (RN or SR) runs anywhere.""" + if state_dtype == torch.float8_e4m3fn and get_sm_version() < 89: + pytest.skip("fp8_e4m3fn requires SM 89+ (Ada Lovelace / Hopper / Blackwell)") + if use_sr and state_dtype in (torch.float16, torch.float8_e4m3fn) and get_sm_version() < 100: + pytest.skip(f"{state_dtype} stochastic rounding requires SM 100+ (Blackwell B200+)") + + +@pytest.mark.skipif(get_sm_version() < 90, reason="TMA descriptor path requires SM 90+") +@pytest.mark.parametrize("block_size_m", [8, 64], ids=["M8", "M64"]) +@pytest.mark.parametrize("rectangle_for_nowrite", [False, True], ids=["replay_nowrite", "rect"]) +def test_replay_tma_strided_state_layout_matches_contiguous(rectangle_for_nowrite, block_size_m): + """Block-reuse packs conv state after each slot's SSM state.""" + torch.manual_seed(123) + + cache_size = 4 + batch = 2 + T = 8 + max_window = 16 + nheads = 16 + head_dim = 64 + d_state = 128 + ngroups = 1 + device = "cuda" + dtype = torch.bfloat16 + state_dtype = torch.float32 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) + + state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + d_inner = nheads * head_dim + conv_dim = d_inner + 2 * ngroups * d_state + slot_gap_rows = conv_dim * 4 // d_state + assert (conv_dim * 4) % d_state == 0 + + old_x = torch.randn(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) + cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + prev_tokens = torch.zeros(cache_size, device=device, dtype=torch.int32) + prev_tokens[state_batch_indices[0]] = 4 + prev_tokens[state_batch_indices[1]] = 12 + n_writes, replay_work_items = _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + ) + + x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt = repeat(dt_base, "b t h -> b t h p", p=head_dim) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + B = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + D = repeat(torch.randn(nheads, device=device, dtype=dtype), "h -> h p", p=head_dim) + dt_bias = repeat(torch.randn(nheads, device=device, dtype=dtype), "h -> h p", p=head_dim) + + dense_state = state0.clone() + strided_state, strided_backing = _make_strided_state_with_slot_gap(state0, slot_gap_rows) + strided_gap = strided_backing[:, nheads * head_dim :] + dense_out = torch.empty(batch, T, nheads, head_dim, device=device, dtype=dtype) + strided_out = torch.empty_like(dense_out) + dense_old_x = old_x.clone() + dense_old_B = old_B.clone() + dense_old_dt = old_dt.clone() + dense_old_dA_cumsum = old_dA_cumsum.clone() + strided_old_x = old_x.clone() + strided_old_B = old_B.clone() + strided_old_dt = old_dt.clone() + strided_old_dA_cumsum = old_dA_cumsum.clone() + + common_kwargs = dict( + prev_num_accepted_tokens=prev_tokens, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + n_writes=n_writes, + replay_work_items=replay_work_items, + use_internal_pdl=False, + rectangle_for_nowrite=rectangle_for_nowrite, + mode="persistent_dynamic", + _block_size_m=block_size_m, + _use_tma_rect_load=True, + _use_tma_replay_write_load=True, + _use_tma_replay_write_store=True, + _use_tma_replay_nowrite_load=True, + _require_tma_state_layout=True, + ) + + replay_selective_state_update( + dense_state, + dense_old_x, + dense_old_B, + dense_old_dt, + dense_old_dA_cumsum, + cache_buf_idx.clone(), + out=dense_out, + **common_kwargs, + ) + replay_selective_state_update( + strided_state, + strided_old_x, + strided_old_B, + strided_old_dt, + strided_old_dA_cumsum, + cache_buf_idx.clone(), + out=strided_out, + **common_kwargs, + ) + + torch.testing.assert_close(strided_out, dense_out, rtol=0, atol=0) + torch.testing.assert_close(strided_state, dense_state, rtol=0, atol=0) + torch.testing.assert_close(strided_old_x, dense_old_x, rtol=0, atol=0) + torch.testing.assert_close(strided_old_B, dense_old_B, rtol=0, atol=0) + torch.testing.assert_close(strided_old_dt, dense_old_dt, rtol=0, atol=0) + torch.testing.assert_close(strided_old_dA_cumsum, dense_old_dA_cumsum, rtol=0, atol=0) + assert torch.isnan(strided_gap).all() @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize("state_dtype", [torch.float16, torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize( + "state_dtype", + [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.int8, + torch.int16, + torch.float8_e4m3fn, + ], + ids=["fp16", "bf16", "fp32", "int8", "int16", "fp8"], +) @pytest.mark.parametrize( "T", [6, 10, 16, 27, 32, 55], ids=["T6", "T10", "T16", "T27", "T32", "T55"] ) +@pytest.mark.parametrize( + "write_checkpoint,rectangle_for_nowrite", + [ + (True, False), # write path; rectangle_for_nowrite is ignored + (False, False), # nowrite path via replay-style kernels + (False, True), # nowrite path via dedicated rectangle kernels + ], + ids=["write", "no_write_replay", "no_write_rectangle"], +) +@pytest.mark.parametrize( + "mode", + ["persistent_dynamic", "persistent_main"], + ids=["persistent_dynamic", "persistent_main"], +) def test_replay_selective_state_update( - nheads, head_dim, d_state, ngroups, state_dtype, paged_cache, T + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + T, + write_checkpoint, + rectangle_for_nowrite, + mode, ): """ Verify that: @@ -55,18 +316,29 @@ def test_replay_selective_state_update( produces the same output as: selective_state_update(state_after_k_old_tokens, new_x, ...) and writes state_after_k_old_tokens back to the state tensor. + + Quantized state dtypes (int8/int16/fp8) follow the same flow with + a per-(head, dim) channel decode-scale tensor; comparison is done + via dequant(state, scales) against the fp32 reference. """ + _maybe_skip_dtype(state_dtype, use_sr=False) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + batch = 2 device = "cuda" dtype = torch.bfloat16 # input activations are bf16 assert nheads % ngroups == 0 - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = None + # Cache T-axis size (max_window). Use the kernel's BLOCK_SIZE_T as the + # ceiling — this is what the wrapper allows and enables PNAT-aware writes + # at [PNAT, PNAT+T) for no-replay-write mode. For T=6 that's 16 (production + # max_window); for larger T it scales with np2(T). + max_window = max(triton.next_power_of_2(T), 16) + + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) torch.manual_seed(42) @@ -82,28 +354,42 @@ def test_replay_selective_state_update( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - # Initial SSM state (cache_size slots) - state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + # Initial SSM state (cache_size slots). Quantized dtypes need a separate + # init: derive scales from a fp32 source so the quantized state isn't + # garbage on dequant. ref_input_state is what the fp32 reference run + # sees — for non-quant it's state0 (cast to fp32 inside reference); for + # quant it's the lossy dequant of state0 (matches what the kernel sees + # internally on load). + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + ref_input_state = _dequantize_state(state0, state0_scales) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None + ref_input_state = state0.float() - # Old inputs: T tokens per batch request - x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + # Seed enough history to cover every PNAT value swept below. + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) # stride(-1)=0 - B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) - # Capture intermediate SSM states using selective_state_update. + # Capture intermediate SSM states using selective_state_update across + # all step1_T positions — gives us reference states for k ∈ [0, step1_T]. states_buffer_f32 = torch.zeros( - cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 - ) - cache_idx_for_capture = ( - state_batch_indices - if paged_cache - else torch.arange(batch, device=device, dtype=torch.int32) + cache_size, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 ) - out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + cache_idx_for_capture = state_batch_indices + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) selective_state_update( - state0.clone(), + ref_input_state.clone(), x1, dt1, A, @@ -114,43 +400,59 @@ def test_replay_selective_state_update( dt_softplus=True, state_batch_indices=cache_idx_for_capture, intermediate_states_buffer=states_buffer_f32, - cache_steps=T, + cache_steps=step1_T, out=out1, disable_state_update=True, ) # Build cache tensors for the replay kernel. - # old_x: (cache, T, nheads, dim) bf16 — single-buffered - # old_B: (cache, 2, T, ngroups, dstate) bf16 — double-buffered - # old_dt: (cache, 2, nheads, T) fp32 — double-buffered, T contiguous - # old_dA_cumsum: (cache, 2, nheads, T) fp32 — double-buffered, T contiguous + # old_x: (cache, 2, max_window, nheads, dim) bf16 — double-buffered + # old_B: (cache, 2, max_window, ngroups, dstate) bf16 — double-buffered + # old_dt: (cache, 2, nheads, max_window) fp32 — double-buffered, window contiguous + # old_dA_cumsum: (cache, 2, nheads, max_window) fp32 — double-buffered, window contiguous # cache_buf_idx: random 0s and 1s to verify indexing correctness - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_x = torch.randn(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - # Fill each slot's READ buffer (indexed by cache_buf_idx) with step 1's data. - # The OTHER buffer has random garbage to catch indexing bugs. - slots = state_batch_indices if paged_cache else slice(None) - old_x[slots] = x1 + # Fill each slot's active buffer (= cache_buf_idx) with step 1's data at + # positions [0:step1_T) = [0:max_window). Whole buffer covered so PNAT + # values up to max_window are exercised. Inactive buffer has random + # garbage to catch indexing bugs. # Compute processed dt and dA_cumsum for step 1 dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) - # Write to each slot's read buffer based on its cache_buf_idx - slot_indices = state_batch_indices.tolist() if paged_cache else list(range(cache_size)) + # Write to each slot's active buffer based on its cache_buf_idx + slot_indices = state_batch_indices.tolist() for i, slot in enumerate(slot_indices): buf = cache_buf_idx[slot].item() batch_idx = i # maps slot back to the batch index - old_B[slot, buf] = B1[batch_idx] - old_dt[slot, buf] = dt1[batch_idx].T # (T, nheads) → (nheads, T) - old_dA_cumsum[slot, buf] = dA_cumsum1[batch_idx].T # (T, nheads) → (nheads, T) - - # Main loop: test each k (number of old tokens replayed) - for k in range(T + 1): + old_x[slot, buf, :step1_T] = x1[batch_idx] + old_B[slot, buf, :step1_T] = B1[batch_idx] + old_dt[slot, buf, :, :step1_T] = dt1[batch_idx].T # (step1_T, nheads) → (nheads, step1_T) + old_dA_cumsum[slot, buf, :, :step1_T] = dA_cumsum1[batch_idx].T + + # Main loop: test each k (number of old tokens replayed). + # write_checkpoint=False (nowrite): k ∈ [0, max_window-T] — new tokens + # append at [k, k+T) of the active buffer; need k+T ≤ max_window. + # write_checkpoint=True (write): k ∈ [max_window-T+1, max_window] — + # new tokens land in the staging buffer at [0, T); k > max_window-T + # captures the overflow case that triggers a checkpoint write in production. + # Combined sweep covers the full k ∈ [0, max_window] with the + # appropriate boundary handling per mode. + if write_checkpoint: + k_lo = max(0, max_window - T + 1) + k_hi = max_window + 1 # exclusive + else: + k_lo = 0 + k_hi = max_window - T + 1 # exclusive + for k in range(k_lo, k_hi): torch.manual_seed(k + 100) x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) @@ -159,8 +461,10 @@ def test_replay_selective_state_update( B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - # Reference - ref_state_f32 = state0.float().clone() + # Reference (fp32, starting from the same lossy-or-not state the + # kernel sees). + slots = state_batch_indices + ref_state_f32 = ref_input_state.clone() if k > 0: ref_state_f32[slots] = states_buffer_f32[slots, k - 1] @@ -175,22 +479,40 @@ def test_replay_selective_state_update( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=(state_batch_indices if paged_cache else None), + state_batch_indices=state_batch_indices, out=ref_out, ) - # Replay kernel + # Replay kernel — clone caches into mutable working copies that we + # can inspect AFTER the call to verify cache postconditions. test_state = state0.clone() + test_scales = state0_scales.clone() if is_quantized else None prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x_w = old_x.clone() + old_B_w = old_B.clone() + old_dt_w = old_dt.clone() + old_dA_cumsum_w = old_dA_cumsum.clone() # cache_buf_idx stays at its random values — each slot reads from its own buffer + # Persistent_main consumes write-first work items. For pure-write or + # pure-nowrite cases here, all slots have the same status, so the + # work-item order is identity. + n_writes_t, replay_work_items_t = _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + ) replay_selective_state_update( test_state, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), + old_x_w, + old_B_w, + old_dt_w, + old_dA_cumsum_w, cache_buf_idx.clone(), prev_tokens, x=x2, @@ -199,60 +521,694 @@ def test_replay_selective_state_update( B=B2, C=C2, out=test_out, + n_writes=n_writes_t, + replay_work_items=replay_work_items_t, D=D, dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, + state_scales=test_scales, + rectangle_for_nowrite=rectangle_for_nowrite, + mode=mode, ) - # Tolerance rationale: the replay kernel uses bf16 tl.dot for four - # matmuls (dB_scaled @ old_x, C @ state, CB_scaled @ x, and C @ B in - # precompute). The reference (selective_state_update) and flashinfer - # baseline use fp32 element-wise MACs. The bf16 input casts lose the - # dt_bias/A-derived bits that the baselines keep — per-element rounding, - # not accumulating. Prefill (ssd_chunk_scan) does identical bf16 tl.dot - # casts, so we match prefill precision exactly. Empirical: max ~1.0 at - # T<=16, ~2.0 at T=32-55; mean ~0.014; <0.02% of elements exceed 0.5. - # State dtype (fp16/bf16/fp32) doesn't shift the error — bf16 dot - # inputs dominate, not state storage. - torch.testing.assert_close( - test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch at k={k}" + # Tolerance rationale: replay uses bf16 tl.dot while the reference + # selective_state_update uses fp32 element-wise MACs. This matches + # prefill precision but needs a small absolute tolerance. Quantized + # states add decode-grid error that propagates through C @ state. + out_atol = ( + {torch.int8: 1.6, torch.int16: 1.05, torch.float8_e4m3fn: 4.0}[state_dtype] + if is_quantized + else 1.0 ) - - expected_state = ( - state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) - ) - torch.testing.assert_close( - test_state[slots], expected_state, rtol=2e-2, atol=1.0, msg=f"State mismatch at k={k}" + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized + else 2e-2 ) + out_diff = (test_out.float() - ref_out.float()).abs() + out_max = out_diff.max().item() + out_mean = out_diff.mean().item() + try: + torch.testing.assert_close( + test_out, + ref_out, + rtol=out_rtol, + atol=out_atol, + msg=f"Output mismatch at k={k}", + ) + except AssertionError: + print( + f"k={k} out: max={out_max:.4f} mean={out_mean:.4f} " + f"nan={torch.isnan(test_out).any().item()} " + f"inf={torch.isinf(test_out).any().item()}" + ) + raise + + # State expectation depends on write_checkpoint: + # True → kernel writes the post-replay SSM state; expect the + # selective_state_update reference's state at step k-1. + # False → kernel skips the HBM store; state must be UNCHANGED + # from the input (state0; for quant, scales also unchanged). + if is_quantized: + if write_checkpoint: + # Compare via dequant against the fp32 reference state. + expected_fp32 = ( + ref_input_state[slots] if k == 0 else states_buffer_f32[slots, k - 1] + ) + actual_fp32 = _dequantize_state(test_state[slots], test_scales[slots]) + # State tolerance covers bf16 replay dot error plus one + # per-element quantization step for the state dtype. + state_atol = { + torch.int8: 1.1, + torch.int16: 1.0, + torch.float8_e4m3fn: 2.5, + }[state_dtype] + state_rtol = { + torch.int8: 5e-2, + torch.int16: 2e-2, + torch.float8_e4m3fn: 1e-1, + }[state_dtype] + try: + torch.testing.assert_close( + actual_fp32, + expected_fp32, + rtol=state_rtol, + atol=state_atol, + msg=f"State mismatch at k={k} dtype={state_dtype}", + ) + except AssertionError: + diff = (actual_fp32 - expected_fp32).abs() + print( + f"k={k} state(dequant): max={diff.max().item():.4f} " + f"mean={diff.mean().item():.4f}" + ) + raise + # Scales sanity (fp32, finite, positive). + assert test_scales.dtype == torch.float32 + assert torch.isfinite(test_scales[slots]).all(), ( + f"state_scales has non-finite values at k={k}" + ) + assert (test_scales[slots] > 0).all(), ( + f"state_scales has non-positive values at k={k}" + ) + else: + # No write: raw quant state and scales unchanged. Use + # torch.equal for byte-level equality (dtype-agnostic; works + # for int8 / int16 / fp8 alike). + assert torch.equal(test_state[slots], state0[slots]), ( + f"Quant state changed at k={k} write_checkpoint=False" + ) + assert torch.equal(test_scales[slots], state0_scales[slots]), ( + f"State scales changed at k={k} write_checkpoint=False" + ) + else: + if write_checkpoint: + expected_state = ( + state0[slots] if k == 0 else states_buffer_f32[slots, k - 1].to(state_dtype) + ) + else: + expected_state = state0[slots] + state_diff = (test_state[slots].float() - expected_state.float()).abs() + state_max = state_diff.max().item() + state_mean = state_diff.mean().item() + try: + torch.testing.assert_close( + test_state[slots], + expected_state, + rtol=2e-2, + atol=1.0 if write_checkpoint else 0.0, + msg=f"State mismatch at k={k} (write_checkpoint={write_checkpoint})", + ) + except AssertionError: + print( + f"k={k} state: max={state_max:.4f} mean={state_mean:.4f} " + f"nan={torch.isnan(test_state).any().item()} " + f"inf={torch.isinf(test_state).any().item()}" + ) + raise + + # --- Cache postconditions --- + # Compute step 2's processed values (what the kernel should have + # stored at [write_offset : write_offset+T) of write_buf): + # write_buf = (1 - active_buf) if write_checkpoint else active_buf + # write_offset = 0 if write_checkpoint else k + # Untouched cache regions must equal their pre-call snapshots + # (old_x / old_B / old_dt / old_dA_cumsum captured before the call). + dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) # (B,T,H) + dA_cumsum2 = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) + write_offset = 0 if write_checkpoint else k + + for batch_idx, slot in enumerate(slot_indices): + active = cache_buf_idx[slot].item() + wb = (1 - active) if write_checkpoint else active + + # --- old_x (double-buffered): write at wb, [write_offset : +T) --- + written_x = old_x_w[slot, wb, write_offset : write_offset + T] + torch.testing.assert_close( + written_x, + x2[batch_idx], + rtol=0, + atol=0, + msg=f"old_x written region wrong at k={k} write={write_checkpoint}", + ) + # Untouched ranges of old_x[slot, wb] + if write_offset > 0: + torch.testing.assert_close( + old_x_w[slot, wb, :write_offset], + old_x[slot, wb, :write_offset], + rtol=0, + atol=0, + msg=f"old_x [0:{write_offset}) modified at k={k} write={write_checkpoint}", + ) + if write_offset + T < max_window: + torch.testing.assert_close( + old_x_w[slot, wb, write_offset + T :], + old_x[slot, wb, write_offset + T :], + rtol=0, + atol=0, + msg=f"old_x [{write_offset + T}:) modified at k={k} write={write_checkpoint}", + ) + # Other-buffer (= 1-wb) untouched + torch.testing.assert_close( + old_x_w[slot, 1 - wb], + old_x[slot, 1 - wb], + rtol=0, + atol=0, + msg=f"old_x inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_B (double-buffered): write at write_buf, [write_offset:+T) --- + torch.testing.assert_close( + old_B_w[slot, wb, write_offset : write_offset + T], + B2[batch_idx], + rtol=0, + atol=0, + msg=f"old_B written region wrong at k={k} write={write_checkpoint}", + ) + # Other-buffer (= 1-wb) untouched + torch.testing.assert_close( + old_B_w[slot, 1 - wb], + old_B[slot, 1 - wb], + rtol=0, + atol=0, + msg=f"old_B inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dt (double-buffered, fp32, layout (heads, T)): --- + torch.testing.assert_close( + old_dt_w[slot, wb, :, write_offset : write_offset + T], + dt2_proc[batch_idx].T, + rtol=1e-4, + atol=1e-4, + msg=f"old_dt written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dt_w[slot, 1 - wb], + old_dt[slot, 1 - wb], + rtol=0, + atol=0, + msg=f"old_dt inactive buffer modified at k={k} write={write_checkpoint}", + ) + + # --- old_dA_cumsum (double-buffered, fp32, layout (heads, T)): --- + # WRITE: fresh staging buf starts from 0, store per-step cumsum. + # NOWRITE: append at offset k of active buf — values are continuous + # from the start of the buffer, so add the prefix at position k-1 + # (matches the kernel's cross-step continuity fix). + if write_checkpoint or k == 0: + expected_dAcs = dA_cumsum2[batch_idx].T + else: + prefix = old_dA_cumsum[slot, wb, :, k - 1] # (heads,) + expected_dAcs = dA_cumsum2[batch_idx].T + prefix[:, None] + torch.testing.assert_close( + old_dA_cumsum_w[slot, wb, :, write_offset : write_offset + T], + expected_dAcs, + rtol=1e-4, + atol=1e-4, + msg=f"old_dA_cumsum written region wrong at k={k} write={write_checkpoint}", + ) + torch.testing.assert_close( + old_dA_cumsum_w[slot, 1 - wb], + old_dA_cumsum[slot, 1 - wb], + rtol=0, + atol=0, + msg=f"old_dA_cumsum inactive buf modified at k={k} write={write_checkpoint}", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,explicit_order,rectangle_for_nowrite", + [ + # All-write: every slot has PNAT triggering write + # (PNAT + T > max_window). No explicit order needed. + ("all_write", [12, 13, 14, 15], None, False), + # All-nowrite: every slot fits in the window. + ("all_nowrite", [3, 4, 5, 6], None, False), + # Mixed PNATs with hand-coded work-item order, independent of the + # _make_replay_work_items test helper. + ("mixed_explicit", [3, 10, 12, 16], [2, 3, 0, 1], False), + ("mixed_explicit_rect", [3, 10, 12, 16], [2, 3, 0, 1], True), + # Mixed PNATs with AUTO-COMPUTED work-item order via `_make_replay_work_items` + # — production-shaped flow. Different PNAT layout from the explicit + # case ([1, 3, 0, 2]) so the kernel sees a distinct order. + ("mixed_auto", [3, 12, 10, 15], None, False), + ("mixed_auto_rect", [3, 12, 10, 15], None, True), + ], + ids=[ + "all_write", + "all_nowrite", + "mixed_explicit", + "mixed_explicit_rect", + "mixed_auto", + "mixed_auto_rect", + ], +) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) +def test_replay_selective_state_update_scenarios( + scenario, + pnat_per_slot_list, + explicit_order, + rectangle_for_nowrite, + mode, +): + """ + Combined scenarios test covering both kernel modes (persistent_main, + persistent_dynamic) across a representative mix of write/nowrite + layouts: + + - all_write / all_nowrite: every slot on one branch — verifies the + empty-half early-return on pm and the all-uniform per-slot dispatch + on pd. + - mixed_explicit: hand-coded work-item order independent of the helper. + - mixed_auto: unsorted PNATs, work-item order auto-computed via the + helper — the production-shaped flow. + + Setup mirrors test_replay_selective_state_update_sorted_dispatch + (same fixed seeds, same input shapes) so the reference state evolution + is identical and we can compare per-slot output and HBM-state + postconditions to the same reference. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, + dt1_input, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.randn(batch, 2, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_x[i, buf, :step1_T] = x1[i] + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + n_writes_t, replay_work_items = _make_replay_work_items( + pnat_per_slot, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + explicit_order=explicit_order, + ) + replay_selective_state_update( + test_state, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + n_writes=n_writes_t, + replay_work_items=replay_work_items, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + mode=mode, + rectangle_for_nowrite=rectangle_for_nowrite, + ) + + torch.testing.assert_close( + test_out.float(), + ref_out.float(), + atol=1.0, + rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), + ref_state_after_replay[i].float(), + atol=1.0, + rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], + state0[i], + rtol=0, + atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) + + +@pytest.mark.parametrize( + "scenario,pnat_per_slot_list,n_writes_expected,work_item_order", + [ + # Mixed batch with write-first work-item order. + ("mixed_sorted", [3, 10, 12, 16], 2, [2, 3, 0, 1]), + # Boundary: n_writes=batch. The nowrite half has an empty slot range + # and the persistent loop must do no work. + ("all_write_noskip", [12, 13, 14, 15], 4, [0, 1, 2, 3]), + # Boundary: n_writes=0. The write half has an empty slot range. + ("all_nowrite_noskip", [3, 4, 5, 6], 0, [0, 1, 2, 3]), + ], + ids=["mixed_sorted", "all_write_noskip", "all_nowrite_noskip"], +) +@pytest.mark.parametrize("rectangle_for_nowrite", [True, False], ids=["rect", "norect"]) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) +def test_replay_selective_state_update_persistent_main_device_n_writes( + scenario, + pnat_per_slot_list, + n_writes_expected, + work_item_order, + rectangle_for_nowrite, + mode, +): + """ + Persistent_main with the device-tensor n_writes plumbing. + + The kernel reads `n_writes` from device memory at entry, so CUDA graphs + can reuse the same captured pointer while updating the value between + replays. Both persistent_main halves launch even when one half has no + slots. + + Verifies: + 1. Kernel reads device n_writes correctly (output matches reference). + 2. Empty-half launches don't corrupt state (n_writes=0 / =batch). + 3. replay_work_items still work through the device-n_writes path. + """ + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 + T = 6 + max_window = 16 + batch = 4 + device = "cuda" + dtype = torch.bfloat16 + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) + + pnat_per_slot = torch.tensor(pnat_per_slot_list, device=device, dtype=torch.int32) + pnat_means_write = (pnat_per_slot + T > max_window).tolist() + write_count = sum(pnat_means_write) + assert write_count == n_writes_expected, ( + f"test setup error: expected {n_writes_expected} writes, got {write_count}" + ) + + # Device-tensor n_writes. Caller mutates between iters in CUDA-graph + # benchmarking; we only run one iter here so a single fill is enough. + n_writes = torch.tensor([n_writes_expected], device=device, dtype=torch.int32) + + torch.manual_seed(42) + A_base = -torch.rand(nheads, device=device) - 0.5 + A = repeat(A_base, "h -> h p n", p=head_dim, n=d_state) + dt_bias_base = torch.randn(nheads, device=device, dtype=dtype) + dt_bias = repeat(dt_bias_base, "h -> h p", p=head_dim) + D_base = torch.randn(nheads, device=device, dtype=dtype) + D = repeat(D_base, "h -> h p", p=head_dim) + + state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=dtype) + ref_input_state = state0.float() + + step1_T = max_window + x1 = torch.randn(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, step1_T, nheads, device=device, dtype=dtype) + dt1_input = repeat(dt1_base, "b t h -> b t h p", p=head_dim) + B1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, step1_T, ngroups, d_state, device=device, dtype=dtype) + + states_buffer_f32 = torch.zeros( + batch, step1_T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) + out1 = torch.zeros(batch, step1_T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_input_state.clone(), + x1, + dt1_input, + A, + B1, + C1, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=cache_idx_for_capture, + intermediate_states_buffer=states_buffer_f32, + cache_steps=step1_T, + out=out1, + disable_state_update=True, + ) + + old_x = torch.randn(batch, 2, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.randn(batch, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.randn(batch, 2, nheads, max_window, device=device, dtype=torch.float32) + cache_buf_idx = torch.randint(0, 2, (batch,), device=device, dtype=torch.int32) + _n_writes_check, replay_work_items = _make_replay_work_items( + pnat_per_slot, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + explicit_order=work_item_order, + ) + torch.testing.assert_close(_n_writes_check, n_writes) + + dt1_processed = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_processed, dim=1) + for i in range(batch): + buf = cache_buf_idx[i].item() + old_x[i, buf, :step1_T] = x1[i] + old_B[i, buf, :step1_T] = B1[i] + old_dt[i, buf, :, :step1_T] = dt1_processed[i].T + old_dA_cumsum[i, buf, :, :step1_T] = dA_cumsum1[i].T + + torch.manual_seed(123) + x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) + B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + + ref_state_f32 = ref_input_state.clone() + for i in range(batch): + k_i = pnat_per_slot[i].item() + if k_i > 0: + ref_state_f32[i] = states_buffer_f32[i, k_i - 1] + ref_state_after_replay = ref_state_f32.clone() + + ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + ref_state_f32, + x2, + dt2, + A, + B2, + C2, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + out=ref_out, + ) + + test_state = state0.clone() + test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + replay_selective_state_update( + test_state, + old_x.clone(), + old_B.clone(), + old_dt.clone(), + old_dA_cumsum.clone(), + cache_buf_idx.clone(), + pnat_per_slot, + x=x2, + dt=dt2, + A=A, + B=B2, + C=C2, + out=test_out, + n_writes=n_writes, + replay_work_items=replay_work_items, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_batch_indices, + mode=mode, + rectangle_for_nowrite=rectangle_for_nowrite, + ) + + torch.testing.assert_close( + test_out.float(), + ref_out.float(), + atol=1.0, + rtol=0.05, + msg=f"Output mismatch (scenario={scenario})", + ) + + for i in range(batch): + if pnat_means_write[i]: + torch.testing.assert_close( + test_state[i].float(), + ref_state_after_replay[i].float(), + atol=1.0, + rtol=0.05, + msg=f"Write slot {i}: state mismatch (scenario={scenario})", + ) + else: + torch.testing.assert_close( + test_state[i], + state0[i], + rtol=0, + atol=0, + msg=f"Nowrite slot {i}: state HBM modified (scenario={scenario})", + ) -@_skip_pre_sm100 @pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["no_cache_indices", "paged_cache"]) +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) -def test_replay_selective_state_update_philox(nheads, head_dim, d_state, ngroups, paged_cache, T): +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) +def test_replay_selective_state_update_philox( + state_dtype, + nheads, + head_dim, + d_state, + ngroups, + T, + mode, +): """ - Verify that Philox stochastic rounding produces correct results. - - Runs our kernel twice with identical inputs: once without rounding - (fp16 state, deterministic), once with rounding (fp16 state, Philox). - The outputs should be nearly identical — stochastic rounding only - perturbs the state by ±1 fp16 ULP, which barely affects output. - Also verifies the state dtype remains fp16. + Verify that Philox stochastic rounding produces correct results across + all SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). + + Runs the replay kernel twice with identical inputs — once without rand_seed + (deterministic RN), once with rand_seed (Philox SR) — and confirms: + - Outputs are within bf16-dot tolerance (state perturbation ≤ 1 ULP). + - State dtype is preserved. + - State difference is bounded by ~1 ULP of the chosen grid. """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + batch = 2 device = "cuda" dtype = torch.bfloat16 - state_dtype = torch.float16 assert nheads % ngroups == 0 - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - else: - cache_size = batch - state_batch_indices = None + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) torch.manual_seed(42) @@ -263,10 +1219,19 @@ def test_replay_selective_state_update_philox(nheads, head_dim, d_state, ngroups D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) + if is_quantized: + state0_fp32 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=torch.float32 + ) + state0, state0_scales = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state0 = torch.randn( + cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype + ) + state0_scales = None - # Cache tensors - old_x = torch.randn(cache_size, T, nheads, head_dim, device=device, dtype=dtype) + # Replay history cache tensors. + old_x = torch.randn(cache_size, 2, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) @@ -281,6 +1246,19 @@ def test_replay_selective_state_update_philox(nheads, head_dim, d_state, ngroups prev_tokens = torch.full((cache_size,), T // 2, device=device, dtype=torch.int32) + # max_window is old_x.shape[2] per the wrapper convention (after dbuf); + # the philox test sets old_x's window axis = T, so max_window = T here. + _max_window_philox = T + _n_writes_philox, _replay_work_items_philox = _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + _max_window_philox, + batch, + state_batch_indices, + device, + ) + common_kwargs = dict( x=x, dt=dt, @@ -291,10 +1269,14 @@ def test_replay_selective_state_update_philox(nheads, head_dim, d_state, ngroups dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, + n_writes=_n_writes_philox, + replay_work_items=_replay_work_items_philox, + mode=mode, ) - # --- Run without rounding (deterministic fp16 state store) --- + # --- Run without rounding (deterministic RN store) --- state_no_round = state0.clone() + scales_no_round = state0_scales.clone() if is_quantized else None out_no_round = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( state_no_round, @@ -305,12 +1287,14 @@ def test_replay_selective_state_update_philox(nheads, head_dim, d_state, ngroups cache_buf_idx.clone(), prev_tokens, out=out_no_round, + state_scales=scales_no_round, **common_kwargs, ) # --- Run with Philox rounding --- rand_seed = torch.tensor([12345], device=device, dtype=torch.int64) state_rounded = state0.clone() + scales_rounded = state0_scales.clone() if is_quantized else None out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( state_rounded, @@ -323,47 +1307,101 @@ def test_replay_selective_state_update_philox(nheads, head_dim, d_state, ngroups out=out_rounded, rand_seed=rand_seed, philox_rounds=10, + state_scales=scales_rounded, **common_kwargs, ) # Outputs should be nearly identical — rounding only perturbs the - # post-replay state by ±1 ULP before the output phase reads it. - torch.testing.assert_close( - out_rounded, out_no_round, rtol=2e-2, atol=1.0, msg="Output diverged with Philox rounding" + # post-replay SSM state by ±1 ULP before the output phase reads it. + # Out_atol = bf16_baseline + 6.5 * per_elem_ULP_after_dequant: + # non-quant fp16: fp16 ULP at typical magnitude is tiny → 1.0 + # int8: amax/127 ≈ 23/127 → 6.5*0.18 ≈ 1.2 + bf16_baseline + # int16: amax/32767 ≈ 7e-4 → ~bf16_baseline only + # fp8: amax/14 ≈ 23/14 → 6.5*1.6 ≈ 10.7 + bf16_baseline + out_atol = ( + {torch.int8: 1.5, torch.int16: 1.0, torch.float8_e4m3fn: 6.0}[state_dtype] + if is_quantized + else 1.0 + ) + out_rtol = ( + {torch.int8: 2e-2, torch.int16: 2e-2, torch.float8_e4m3fn: 5e-2}[state_dtype] + if is_quantized + else 2e-2 ) - - # State should remain fp16 - assert state_rounded.dtype == torch.float16 - - # States should differ by at most 1 fp16 ULP per element. - # fp16 ULP depends on magnitude: up to 0.5 for values near 512. - # Use rtol to account for magnitude-dependent ULP. - slots = state_batch_indices if paged_cache else slice(None) torch.testing.assert_close( - state_rounded[slots], - state_no_round[slots], - rtol=2e-3, - atol=0.2, - msg="State diverged with Philox rounding", + out_rounded, + out_no_round, + rtol=out_rtol, + atol=out_atol, + msg=f"Output diverged with Philox rounding ({state_dtype})", ) + # State dtype preserved. + assert state_rounded.dtype == state_dtype + + # State diff between RN and SR is bounded by 1 quant cell per element. + # Per-channel decode_scale varies by 10x+ across channels (amax depends + # on randn extremes), so a single flat atol can't bound it accurately — + # use per-channel ULP-aware comparison. + slots = state_batch_indices + if is_quantized: + rounded_fp32 = _dequantize_state(state_rounded[slots], scales_rounded[slots]) + no_round_fp32 = _dequantize_state(state_no_round[slots], scales_no_round[slots]) + diff = (rounded_fp32 - no_round_fp32).abs() + # Per-element bound = max(decode_scale_no_round, decode_scale_rounded). + # decode_scale is shape (cache, nheads, dim); broadcast over dstate. + scale_bound = torch.maximum(scales_no_round[slots], scales_rounded[slots]).unsqueeze(-1) + # int8 / int16: 1 cell after dequant = decode_scale exactly. + # fp8_e4m3: variable grid; the largest cell within a channel scaled + # to fit ±448 is at the channel's max-magnitude element, where the + # cell is ~32x larger than the average. Bound = decode_scale * 32. + # Apply a 1.5x slack pad for floating-point compare quirks at the + # exact-cell boundary. + cell_pad = 32.0 if state_dtype == torch.float8_e4m3fn else 1.0 + bound = scale_bound * (cell_pad * 1.5) + if not (diff <= bound).all(): + offenders = (diff > bound).sum().item() + n_total = diff.numel() + pytest.fail( + f"State RN-SR diff exceeds 1 cell per element for " + f"{offenders}/{n_total} elements ({state_dtype}). " + f"max_diff={diff.max().item():.4g}, " + f"max_bound={bound.max().item():.4g}." + ) + else: + # fp16 ULP depends on magnitude — rtol absorbs that. + torch.testing.assert_close( + state_rounded[slots], + state_no_round[slots], + rtol=2e-3, + atol=0.2, + msg=f"State diverged with Philox rounding ({state_dtype})", + ) -@_skip_pre_sm100 -def test_philox_rounding_unbiased(): - """ - Verify that Philox stochastic rounding is unbiased. - Runs the replay kernel with fp32 state (capturing the true fp32 - post-replay state) and with fp16 state + Philox rounding. Compares the - rounding residual (fp16_state.float() - fp32_state) against deterministic - rounding (fp32_state.to(fp16).float() - fp32_state). +@pytest.mark.parametrize( + "state_dtype", + [torch.float16, torch.int8, torch.int16, torch.float8_e4m3fn], + ids=["fp16", "int8", "int16", "fp8"], +) +def test_philox_rounding_unbiased(state_dtype): + """ + Verify that Philox stochastic rounding is unbiased across all + SR-supported state dtypes (fp16, int8, int16, fp8_e4m3fn). - Deterministic round-to-nearest-even has a systematic positive bias on - the residual. Philox stochastic rounding should be unbiased: the mean - residual should be near zero. + Captures the true fp32 post-replay SSM state by running with fp32 storage, + then runs the kernel with the target dtype + Philox SR. Compares the + SR rounding residual against the deterministic-RN residual: SR should + have mean residual closer to zero than RN, since RN has a systematic + round-to-nearest-even bias and SR is unbiased by construction. Uses a large batch (16) for ~2M state elements — plenty of statistics. """ + _maybe_skip_dtype(state_dtype, use_sr=True) + + quant_max = _QUANT_MAX_BY_DTYPE.get(state_dtype, 0.0) + is_quantized = quant_max > 0.0 + nheads, head_dim, d_state, ngroups = 16, 64, 128, 1 batch, T = 16, 6 device = "cuda" @@ -377,10 +1415,11 @@ def test_philox_rounding_unbiased(): D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - # Use fp32 initial state so replay produces non-fp16-representable values - state0 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=torch.float32) + # fp32 reference state — replay produces values that don't fit cleanly + # in the target dtype's grid, exposing the rounding bias. + state0_fp32 = torch.randn(batch, nheads, head_dim, d_state, device=device, dtype=torch.float32) - old_x = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) + old_x = torch.randn(batch, 2, T, nheads, head_dim, device=device, dtype=dtype) old_B = torch.randn(batch, 2, T, ngroups, d_state, device=device, dtype=dtype) old_dt = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) old_dA_cumsum = torch.randn(batch, 2, nheads, T, device=device, dtype=torch.float32) @@ -393,7 +1432,18 @@ def test_philox_rounding_unbiased(): C = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) prev_tokens = torch.full((batch,), T, device=device, dtype=torch.int32) + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) + # max_window = old_x.shape[2] = T (after dbuf at axis 1) + _n_writes_unb, _replay_work_items_unb = _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + T, + batch, + state_batch_indices, + device, + ) common_kwargs = dict( x=x, dt=dt_val, @@ -403,10 +1453,13 @@ def test_philox_rounding_unbiased(): D=D, dt_bias=dt_bias, dt_softplus=True, + state_batch_indices=state_batch_indices, + n_writes=_n_writes_unb, + replay_work_items=_replay_work_items_unb, ) - # 1. fp32 state — captures true post-replay state - state_fp32 = state0.clone() + # 1. fp32 state — captures true post-replay fp32 state. + state_fp32 = state0_fp32.clone() out_fp32 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( state_fp32, @@ -420,9 +1473,14 @@ def test_philox_rounding_unbiased(): **common_kwargs, ) - # 2. fp16 state with Philox rounding + # 2. Target dtype + Philox SR. For quant we also need scales (derived + # from the same per-channel amax used by the kernel on store). rand_seed = torch.tensor([99999], device=device, dtype=torch.int64) - state_rounded = state0.to(torch.float16).clone() + if is_quantized: + state_rounded, scales_rounded = _quantize_state(state0_fp32, state_dtype, quant_max) + else: + state_rounded = state0_fp32.to(state_dtype) + scales_rounded = None out_rounded = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) replay_selective_state_update( state_rounded, @@ -435,44 +1493,62 @@ def test_philox_rounding_unbiased(): out=out_rounded, rand_seed=rand_seed, philox_rounds=10, + state_scales=scales_rounded, **common_kwargs, ) - # Compute rounding residuals where fp32 state has non-zero values - fp32_vals = state_fp32.flatten() - stochastic_residual = state_rounded.float().flatten() - fp32_vals - deterministic_residual = fp32_vals.to(torch.float16).float() - fp32_vals + # Compute residuals. For non-quant: stochastic_residual = SR(fp32) - + # fp32, deterministic_residual = RN(fp32) - fp32. For quant: dequant + # both, comparing in fp32. + if is_quantized: + fp32_vals = state_fp32.flatten() + stochastic_residual = _dequantize_state(state_rounded, scales_rounded).flatten() - fp32_vals + # Deterministic reference: do the same per-channel quant on the + # captured fp32 state, then dequant. This is what the kernel would + # have produced with rand_seed=None. + det_quant, det_scales = _quantize_state(state_fp32, state_dtype, quant_max) + deterministic_residual = _dequantize_state(det_quant, det_scales).flatten() - fp32_vals + else: + fp32_vals = state_fp32.flatten() + stochastic_residual = state_rounded.float().flatten() - fp32_vals + deterministic_residual = fp32_vals.to(state_dtype).float() - fp32_vals - # Only consider elements where rounding matters (non-zero residual possible) + # Only consider elements where rounding matters (non-zero residual possible). nonzero_mask = deterministic_residual.abs() > 0 num_nonzero = nonzero_mask.sum().item() assert num_nonzero > 1000, f"Too few roundable elements: {num_nonzero}" stochastic_mean = stochastic_residual[nonzero_mask].mean().item() + stochastic_std = stochastic_residual[nonzero_mask].std().item() deterministic_mean = deterministic_residual[nonzero_mask].mean().item() - # Stochastic rounding should be less biased than deterministic. - # With ~millions of elements, the stochastic mean should be very close to 0. - # Deterministic round-to-nearest-even has a small but systematic bias. - assert abs(stochastic_mean) < abs(deterministic_mean) or abs(stochastic_mean) < 1e-5, ( - f"Stochastic rounding appears biased: stochastic_mean={stochastic_mean:.6f}, " - f"deterministic_mean={deterministic_mean:.6f}, n_elements={num_nonzero}" + # SE-based bias check. An unbiased estimator's sample mean has standard + # error SE = std / sqrt(n). We require |sr_mean| < K*SE (K=4 ≈ ~3.2e-5 + # one-sided false-positive rate). This auto-calibrates per dtype: + # * int16: residual std ~1e-4 → SE ~9e-8 (very tight bound) + # * int8: residual std ~3e-2 → SE ~2e-5 + # * fp8: residual std ~1e-1 → SE ~9e-5 (loosest, magnitude-driven) + # A fixed absolute threshold is below SE for int8/fp8. Gaussian inputs + # also make RN nearly unbiased, so |sr| < |det| is not reliable here. + se_sr = stochastic_std / (num_nonzero**0.5) + K = 4 + assert abs(stochastic_mean) < K * se_sr, ( + f"SR mean exceeds {K}*SE (likely biased) ({state_dtype}): " + f"stochastic_mean={stochastic_mean:.3e}, " + f"SE={se_sr:.3e} (K*SE={K * se_sr:.3e}), " + f"deterministic_mean={deterministic_mean:.3e} (for reference), " + f"n_elements={num_nonzero}" ) -# HEADS_PER_BLOCK > 1 test. The default heuristic only picks HPB > 1 at large -# total_heads (>= 256-512), which the main test with batch=2 never reaches. -# This test overrides _heads_per_block to exercise the two-loop structure in -# the precompute kernel (store-then-reload of per-head dt/dA_cumsum). -# Configs: (nheads=16, ngroups=1) and (nheads=32, ngroups=2) both have -# heads_per_group=16. The heuristic caps HPB at min(2|4, hpg), so HPB=2, 4. -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +# HEADS_PER_BLOCK coverage for multi-head precompute tiles and non-power +# heads_per_group, where the wrapper must keep each tile within one B/C group. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _HEADS_PER_BLOCK_CONFIGS) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16, 32], ids=["T6", "T16", "T32"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -@pytest.mark.parametrize("launch_with_pdl", [False, True], ids=["no_ext_pdl", "ext_pdl"]) -@pytest.mark.parametrize("use_internal_pdl", [False, True], ids=["no_int_pdl", "int_pdl"]) -@pytest.mark.parametrize("batch", [1, 2, 8, 16], ids=["B1", "B2", "B8", "B16"]) +@pytest.mark.parametrize("rectangle_nowrite", [True, False], ids=["rect", "norect"]) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_heads_per_block( nheads, head_dim, @@ -481,25 +1557,27 @@ def test_replay_heads_per_block( state_dtype, T, heads_per_block, - launch_with_pdl, - use_internal_pdl, - batch, + rectangle_nowrite, + mode, ): """ Verify replay_selective_state_update produces correct results when - _heads_per_block > 1, exercising the precompute kernel's two-loop - structure (store per-head dt/dA_cumsum in loop 1, reload in loop 2). + _heads_per_block > 1. + + In addition to the output + state checks, this verifies the full write + contract: per-slot the kernel touches the correct staging/active buffer + at the correct offset for old_x, old_B, old_dt, old_dA_cumsum; leaves + untouched regions and the other buffer byte-identical to pre-call; and + (for nowrite) preserves the dA_cumsum prefix continuity by adding the + pre-call old_dA_cumsum[slot, active_buf, head, PNAT-1] value. """ + # PDL flags use wrapper defaults; trimming the parametrize keeps this + # suite fast. Coverage of {launch_with_pdl, use_internal_pdl} variations + # lives in the dedicated correctness tests above (test_replay_selective_state_update). + batch = 8 device = "cuda" dtype = torch.bfloat16 - if nheads % heads_per_block != 0: - pytest.skip(f"nheads ({nheads}) not divisible by heads_per_block ({heads_per_block})") - if heads_per_block > nheads // ngroups: - pytest.skip( - f"heads_per_block ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})" - ) - torch.manual_seed(42) A_base = -torch.rand(nheads, device=device) - 0.5 @@ -509,20 +1587,27 @@ def test_replay_heads_per_block( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) + # max_window = 2*next_pow2(T) so the PNAT=T nowrite case is always + # valid (requires max_window >= 2T) and we have headroom for the + # full PNAT sweep below. + max_window = max(2 * triton.next_power_of_2(T), 16) + cache_size = batch state0 = torch.randn(cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype) - x1 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) - dt1_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) + # Generate enough fill data to cover max_window steps (needed for the + # write-slot replay reference, which walks up to max_window-1 steps). + x1 = torch.randn(batch, max_window, nheads, head_dim, device=device, dtype=dtype) + dt1_base = torch.randn(batch, max_window, nheads, device=device, dtype=dtype) dt1 = repeat(dt1_base, "b t h -> b t h p", p=head_dim) - B1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) - C1 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + B1 = torch.randn(batch, max_window, ngroups, d_state, device=device, dtype=dtype) + C1 = torch.randn(batch, max_window, ngroups, d_state, device=device, dtype=dtype) states_buffer_f32 = torch.zeros( - cache_size, T, nheads, head_dim, d_state, device=device, dtype=torch.float32 + cache_size, max_window, nheads, head_dim, d_state, device=device, dtype=torch.float32 ) cache_idx_for_capture = torch.arange(batch, device=device, dtype=torch.int32) - out1 = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + out1 = torch.zeros(batch, max_window, nheads, head_dim, device=device, dtype=dtype) selective_state_update( state0.clone(), x1, @@ -535,39 +1620,93 @@ def test_replay_heads_per_block( dt_softplus=True, state_batch_indices=cache_idx_for_capture, intermediate_states_buffer=states_buffer_f32, - cache_steps=T, + cache_steps=max_window, out=out1, disable_state_update=True, ) - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.randn(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.randn(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + # Pre-fill cache buffers. All four (old_x, old_B, old_dt, old_dA_cumsum) + # are double-buffered. Initialize BOTH buffers with controlled random + # data so "outside write range / other buffer unchanged" assertions have + # well-defined expected values for both. + old_x_init = torch.randn( + cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype + ) + old_B_init = torch.randn( + cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype + ) + old_dt_init = torch.randn(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum_init = torch.randn( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) cache_buf_idx = torch.randint(0, 2, (cache_size,), device=device, dtype=torch.int32) - old_x[:] = x1 - dt1 = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) - dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1, dim=1) + # Capture-step data populates ONE buffer per slot (the active one selected + # by cache_buf_idx). Mirrors production: previous step wrote into the + # now-active buffer; the other buffer holds stale data the kernel must + # not touch on a nowrite call. + dt1_proc = F.softplus(dt1_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum1 = torch.cumsum(A_base.float()[None, None, :] * dt1_proc, dim=1) for slot in range(cache_size): - buf = cache_buf_idx[slot].item() - old_B[slot, buf] = B1[slot] - old_dt[slot, buf] = dt1[slot].T - old_dA_cumsum[slot, buf] = dA_cumsum1[slot].T + buf = int(cache_buf_idx[slot].item()) + old_x_init[slot, buf] = x1[slot] + old_B_init[slot, buf] = B1[slot] + old_dt_init[slot, buf] = dt1_proc[slot].T # (nheads, max_window) + old_dA_cumsum_init[slot, buf] = dA_cumsum1[slot].T # (nheads, max_window) + + # --- PNAT sweep -------------------------------------------------------- + # Cover nowrite (PNAT+T <= max_window) and write (PNAT+T > max_window) + # paths plus the boundary, with both PNAT=0 (no prefix) and PNAT=T (the + # smallest prefix-load case the kernel cares about). + candidate_pnats = [ + 0, # nowrite, no prefix + 1, # nowrite, smallest nontrivial prefix + T, # nowrite, prefix length one stored step + max_window - T - 1, # nowrite, largest PNAT just below threshold + max_window - T, # nowrite, exactly at threshold + max_window - T + 1, # write, smallest PNAT above threshold + max_window - 1, # write, maximum + ] + seen = set() + pnat_list = [] + for p in candidate_pnats: + if 0 <= p < max_window and p not in seen: + seen.add(p) + pnat_list.append(p) + while len(pnat_list) < batch: + pnat_list.append(0) + pnat_list = pnat_list[:batch] + + has_write = any((p + T) > max_window for p in pnat_list) + has_nowrite = any((p + T) <= max_window for p in pnat_list) + assert has_write and has_nowrite, ( + f"PNAT sweep must cover both write and nowrite: {pnat_list}, T={T}, max_window={max_window}" + ) - k = T - torch.manual_seed(123) + prev_tokens = torch.tensor(pnat_list, device=device, dtype=torch.int32) + pnat_means_write = [(pnat_list[i] + T) > max_window for i in range(batch)] + torch.manual_seed(123) x2 = torch.randn(batch, T, nheads, head_dim, device=device, dtype=dtype) dt2_base = torch.randn(batch, T, nheads, device=device, dtype=dtype) dt2 = repeat(dt2_base, "b t h -> b t h p", p=head_dim) B2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) C2 = torch.randn(batch, T, ngroups, d_state, device=device, dtype=dtype) + # Per-step processed dt and per-step cumsum — what the kernel writes to + # old_dt and the cumsum-from-zero portion of old_dA_cumsum. + dt2_proc = F.softplus(dt2_base.float() + dt_bias_base.float()[None, None, :]) + dA_cumsum2_step = torch.cumsum(A_base.float()[None, None, :] * dt2_proc, dim=1) + + # Build reference by replaying old history, then this step's tokens. ref_state_f32 = state0.float().clone() - ref_state_f32[:] = states_buffer_f32[:, k - 1] + for slot in range(batch): + if pnat_list[slot] > 0: + ref_state_f32[slot] = states_buffer_f32[slot, pnat_list[slot] - 1] + ref_state_after_replay = ref_state_f32.clone() ref_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + state_batch_indices = torch.arange(batch, device=device, dtype=torch.int32) selective_state_update( ref_state_f32, x2, @@ -578,21 +1717,44 @@ def test_replay_heads_per_block( D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, out=ref_out, ) + # Pre-call snapshots double as the "expected" baseline for untouched + # regions / untouched buffer. Kernel operates on the _test copies. test_state = state0.clone() - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + state_pre = test_state.clone() + old_x_pre = old_x_init.clone() + old_B_pre = old_B_init.clone() + old_dt_pre = old_dt_init.clone() + old_dA_cumsum_pre = old_dA_cumsum_init.clone() + cache_buf_idx_pre = cache_buf_idx.clone() + + old_x_test = old_x_pre.clone() + old_B_test = old_B_pre.clone() + old_dt_test = old_dt_pre.clone() + old_dA_cumsum_test = old_dA_cumsum_pre.clone() + cache_buf_idx_test = cache_buf_idx_pre.clone() + + n_writes_t, replay_work_items_t = _make_replay_work_items( + prev_tokens, + cache_buf_idx_test, + T, + max_window, + batch, + state_batch_indices, + device, + ) replay_selective_state_update( test_state, - old_x.clone(), - old_B.clone(), - old_dt.clone(), - old_dA_cumsum.clone(), - cache_buf_idx.clone(), + old_x_test, + old_B_test, + old_dt_test, + old_dA_cumsum_test, + cache_buf_idx_test, prev_tokens, x=x2, dt=dt2, @@ -600,62 +1762,201 @@ def test_replay_heads_per_block( B=B2, C=C2, out=test_out, + n_writes=n_writes_t, + replay_work_items=replay_work_items_t, D=D, dt_bias=dt_bias, dt_softplus=True, - state_batch_indices=None, + state_batch_indices=state_batch_indices, + mode=mode, + rectangle_for_nowrite=rectangle_nowrite, _heads_per_block=heads_per_block, - launch_with_pdl=launch_with_pdl, - use_internal_pdl=use_internal_pdl, ) + # ---------------- Output + state checks --------------------------------- torch.testing.assert_close( test_out, ref_out, rtol=2e-2, atol=1.0, msg=f"Output mismatch with HPB={heads_per_block}, T={T}, " - f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}, " + f"rect={rectangle_nowrite}, pnats={pnat_list}", ) - expected_state = states_buffer_f32[:, k - 1].to(state_dtype) - torch.testing.assert_close( - test_state, - expected_state, - rtol=2e-2, - atol=1.0, - msg=f"State mismatch with HPB={heads_per_block}, T={T}, " - f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}", - ) + for slot in range(batch): + if pnat_means_write[slot]: + torch.testing.assert_close( + test_state[slot].float(), + ref_state_after_replay[slot].float(), + rtol=2e-2, + atol=1.0, + msg=( + f"Write slot {slot} (PNAT={pnat_list[slot]}): state mismatch " + f"(HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, state_dtype={state_dtype}, " + f"rect={rectangle_nowrite})" + ), + ) + else: + torch.testing.assert_close( + test_state[slot], + state_pre[slot], + rtol=0, + atol=0, + msg=( + f"Nowrite slot {slot} (PNAT={pnat_list[slot]}): state HBM " + f"modified (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, state_dtype={state_dtype}, " + f"rect={rectangle_nowrite})" + ), + ) + + # ---------------- New checks: kernel write contract --------------------- + # For each slot: decide active/staging buffer, write offset, build the + # expected per-buffer tensor element-wise, compare against the actual + # kernel-modified tensor. The expected_* tensors are clones of the + # pre-call snapshot with only [target_buf, write_offset:write_end] + # overwritten — so the full-slot equality compares implicitly assert + # the "other buffer untouched" and "outside write range untouched" + # contracts. + for slot in range(batch): + pnat = pnat_list[slot] + is_write = pnat_means_write[slot] + active_buf = int(cache_buf_idx_pre[slot].item()) + staging_buf = 1 - active_buf + + if is_write: + target_buf = staging_buf + write_offset = 0 + else: + target_buf = active_buf + write_offset = pnat + write_end = write_offset + T + + # ----- old_x (double-buffer (cache, 2, max_window, nheads, dim)) ----- + expected_old_x_slot = old_x_pre[slot].clone() + expected_old_x_slot[target_buf, write_offset:write_end] = x2[slot] + torch.testing.assert_close( + old_x_test[slot], + expected_old_x_slot, + rtol=0, + atol=0, + msg=( + f"old_x slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) + + # ----- old_B (double-buffer) ----- + expected_old_B_slot = old_B_pre[slot].clone() + expected_old_B_slot[target_buf, write_offset:write_end] = B2[slot] + torch.testing.assert_close( + old_B_test[slot], + expected_old_B_slot, + rtol=0, + atol=0, + msg=( + f"old_B slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) + + # ----- old_dt (double-buffer (cache, 2, nheads, max_window)) ----- + # Kernel writes per-head processed dt at [target_buf, :, write_offset:write_end]. + # softplus on chip vs F.softplus host: small ULP diff possible; use + # tight but non-zero tolerance. + expected_old_dt_slot = old_dt_pre[slot].clone() + expected_old_dt_slot[target_buf, :, write_offset:write_end] = dt2_proc[slot].T + torch.testing.assert_close( + old_dt_test[slot], + expected_old_dt_slot, + rtol=1e-5, + atol=1e-5, + msg=( + f"old_dt slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) + + # ----- old_dA_cumsum (double-buffer (cache, 2, nheads, max_window)) ----- + # WRITE: per-step cumsum starting from 0 (fresh staging buf). + # NOWRITE: continuous — cumsum offset by the prefix value at + # old_dA_cumsum_pre[slot, active_buf, head, PNAT-1] (or 0 if PNAT=0). + # Verifies dA_cumsum continuity across no-write appends. + expected_old_dAcs_slot = old_dA_cumsum_pre[slot].clone() + step_cumsum = dA_cumsum2_step[slot].T # (nheads, T) + if is_write: + expected_old_dAcs_slot[target_buf, :, write_offset:write_end] = step_cumsum + else: + if pnat > 0: + prefix = old_dA_cumsum_pre[slot, active_buf, :, pnat - 1] + else: + prefix = torch.zeros(nheads, device=device, dtype=torch.float32) + expected_old_dAcs_slot[target_buf, :, write_offset:write_end] = ( + step_cumsum + prefix[:, None] + ) + torch.testing.assert_close( + old_dA_cumsum_test[slot], + expected_old_dAcs_slot, + rtol=1e-5, + atol=1e-5, + msg=( + f"old_dA_cumsum slot {slot} (PNAT={pnat}, is_write={is_write}, " + f"target_buf={target_buf}, write_offset={write_offset}): " + f"mismatch (HPB={heads_per_block}, T={T}, nheads={nheads}, " + f"ngroups={ngroups}, rect={rectangle_nowrite})" + ), + ) # HPB > 1 multi-step test. Production chains decode steps; bugs in # buffer ordering or stale cache values accumulate across steps and can # be invisible in a single-step test. -@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _CONFIGS) +# +# Divergent per-slot acceptance: slot 0 accepts more tokens each step, +# slot 1 accepts a smaller fixed count. This forces the write/nowrite +# mask to differ between slots on multiple steps (n_writes in {0, 1, 2} +# within an 8-step run and the work-item order hits both identity [0,1] +# and swapped [1,0] — exercising the kernel's per-slot dispatch). +# `rectangle_nowrite` forces the rectangle vs non-rectangle nowrite path +# (via `mode="persistent_main"` + `rectangle_for_nowrite=…` kwargs) +# rather than relying on the tuning table's mode pick. +@pytest.mark.parametrize("nheads,head_dim,d_state,ngroups", _HEADS_PER_BLOCK_CONFIGS) @pytest.mark.parametrize("state_dtype", [torch.bfloat16, torch.float32]) @pytest.mark.parametrize("T", [6, 16], ids=["T6", "T16"]) @pytest.mark.parametrize("heads_per_block", [2, 4], ids=["HPB2", "HPB4"]) -@pytest.mark.parametrize("paged_cache", [False, True], ids=["contig", "paged"]) +@pytest.mark.parametrize("rectangle_nowrite", [True, False], ids=["rect", "norect"]) +@pytest.mark.parametrize("mode", ["persistent_main", "persistent_dynamic"], ids=["pm", "pd"]) def test_replay_heads_per_block_multistep( - nheads, head_dim, d_state, ngroups, state_dtype, T, heads_per_block, paged_cache + nheads, + head_dim, + d_state, + ngroups, + state_dtype, + T, + heads_per_block, + rectangle_nowrite, + mode, ): """ Chain N decode steps with HPB > 1 and verify each step's output matches a fresh reference. A bug that mixes up WRITE/READ buffers, writes wrong - data to cache, or races in the two-loop structure would accumulate - across steps. + data to cache, or reuses stale cache values would accumulate across steps. + Per-slot acceptance diverges so write/nowrite masks differ across slots; + the n_writes=1 case (mixed batch) is exercised. """ batch = 2 device = "cuda" dtype = torch.bfloat16 n_steps = 8 - if nheads % heads_per_block != 0: - pytest.skip(f"nheads ({nheads}) not divisible by HPB ({heads_per_block})") - if heads_per_block > nheads // ngroups: - pytest.skip(f"HPB ({heads_per_block}) exceeds heads_per_group ({nheads // ngroups})") - torch.manual_seed(42) A_base = -torch.rand(nheads, device=device) - 0.5 @@ -665,14 +1966,8 @@ def test_replay_heads_per_block_multistep( D_base = torch.randn(nheads, device=device, dtype=dtype) D = repeat(D_base, "h -> h p", p=head_dim) - if paged_cache: - cache_size = 4 - state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) - slots = state_batch_indices - else: - cache_size = batch - state_batch_indices = None - slots = slice(None) + cache_size = 4 + state_batch_indices = torch.tensor([1, 3], device=device, dtype=torch.int32) all_x = [] all_dt = [] @@ -691,41 +1986,81 @@ def test_replay_heads_per_block_multistep( cache_size, nheads, head_dim, d_state, device=device, dtype=state_dtype ) + # max_window = 2*np2(T) so the buffer has slack for several nowrite + # steps before overflow — required for the divergent-acceptance pattern + # to produce a chain of mixed write/nowrite steps within n_steps=8. + # For T=6 this is 16, for T=16 this is 32. + max_window = max(triton.next_power_of_2(2 * T), 16) + + # Per-slot acceptance counts. Slot 0 advances by 6 per step, slot 1 by + # 4 — divergence (PNAT trajectories desync, write_mask varies between + # slots, n_writes hits 0/1/2 within the 8-step run). Values constant + # across T because they exercise the kernel's per-slot dispatch + # independent of T; acc <= T is the only requirement. + accepted_per_slot = [6, 4] + accepted_tensor = torch.tensor(accepted_per_slot, device=device, dtype=torch.int32) + + # Per-slot reference: each slot's reference state advances by only its + # `accepted` tokens per step, not all T. selective_state_update doesn't + # natively support per-slot variable T, so run it once per slot per step + # with a single-slot batch view. ref_state = state_init.float().clone() ref_outs = [] - ref_slots = ( - state_batch_indices - if paged_cache - else torch.arange(batch, device=device, dtype=torch.int32) - ) for step in range(n_steps): out_step = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) - selective_state_update( - ref_state, - all_x[step], - all_dt[step], - A, - all_B[step], - all_C[step], - D=D, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=ref_slots, - out=out_step, - ) + for s_local in range(batch): + acc = accepted_per_slot[s_local] + if acc == 0: + continue + c_idx = state_batch_indices[s_local].item() + s_state = ref_state[c_idx : c_idx + 1].clone() + s_x = all_x[step][s_local : s_local + 1, :acc].contiguous() + s_dt = all_dt[step][s_local : s_local + 1, :acc].contiguous() + s_B = all_B[step][s_local : s_local + 1, :acc].contiguous() + s_C = all_C[step][s_local : s_local + 1, :acc].contiguous() + s_out = torch.zeros(1, acc, nheads, head_dim, device=device, dtype=dtype) + selective_state_update( + s_state, + s_x, + s_dt, + A, + s_B, + s_C, + D=D, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=torch.tensor([0], device=device, dtype=torch.int32), + out=s_out, + ) + out_step[s_local, :acc] = s_out[0] + ref_state[c_idx] = s_state[0] ref_outs.append(out_step) test_state = state_init.clone() - old_x = torch.zeros(cache_size, T, nheads, head_dim, device=device, dtype=dtype) - old_B = torch.zeros(cache_size, 2, T, ngroups, d_state, device=device, dtype=dtype) - old_dt = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) - old_dA_cumsum = torch.zeros(cache_size, 2, nheads, T, device=device, dtype=torch.float32) + old_x = torch.zeros(cache_size, 2, max_window, nheads, head_dim, device=device, dtype=dtype) + old_B = torch.zeros(cache_size, 2, max_window, ngroups, d_state, device=device, dtype=dtype) + old_dt = torch.zeros(cache_size, 2, nheads, max_window, device=device, dtype=torch.float32) + old_dA_cumsum = torch.zeros( + cache_size, 2, nheads, max_window, device=device, dtype=torch.float32 + ) cache_buf_idx = torch.zeros(cache_size, device=device, dtype=torch.int32) + # Per-active-slot PNAT tracker, advanced by `accepted` (not T) per step. + pnat_active = torch.zeros(batch, device=device, dtype=torch.int32) + for step in range(n_steps): - k = T if step > 0 else 0 - prev_tokens = torch.full((cache_size,), k, device=device, dtype=torch.int32) + prev_tokens = torch.zeros(cache_size, device=device, dtype=torch.int32) + prev_tokens[state_batch_indices.long()] = pnat_active test_out = torch.zeros(batch, T, nheads, head_dim, device=device, dtype=dtype) + n_writes_t, replay_work_items_t = _make_replay_work_items( + prev_tokens, + cache_buf_idx, + T, + max_window, + batch, + state_batch_indices, + device, + ) replay_selective_state_update( test_state, @@ -741,24 +2076,282 @@ def test_replay_heads_per_block_multistep( B=all_B[step], C=all_C[step], out=test_out, + n_writes=n_writes_t, + replay_work_items=replay_work_items_t, D=D, dt_bias=dt_bias, dt_softplus=True, state_batch_indices=state_batch_indices, _heads_per_block=heads_per_block, + mode=mode, + rectangle_for_nowrite=rectangle_nowrite, + ) + + # PNAT update uses `accepted` (per-slot), not T. Write step resets + # PNAT to `accepted` of this step (new buffer starts fresh); nowrite + # appends `accepted` to current PNAT. + write_mask = (pnat_active + T) > max_window + new_pnat_active = torch.where( + write_mask, + accepted_tensor, + pnat_active + accepted_tensor, ) + pnat_active = new_pnat_active + cache_active_idx = state_batch_indices.long() + write_slots = cache_active_idx[write_mask] + cache_buf_idx[write_slots] = 1 - cache_buf_idx[write_slots] + + # Per-slot output comparison: only the first `accepted` output tokens + # of each slot are meaningful (the rest are produced from "candidate" + # tokens that wouldn't be accepted in production). + for s_local in range(batch): + acc = accepted_per_slot[s_local] + if acc == 0: + continue + torch.testing.assert_close( + test_out[s_local, :acc], + ref_outs[step][s_local, :acc], + rtol=2e-2, + atol=2.0, + msg=f"Output mismatch at step {step}, slot {s_local} " + f"(acc={acc}) with HPB={heads_per_block}, T={T}, " + f"nheads={nheads}, ngroups={ngroups}, state_dtype={state_dtype}, " + f"rectangle={rectangle_nowrite}", + ) + + +# ----- SR grid-bracket tests (fp8 and fp16) ----- +# +# Verify that each PTX SR output lands on the destination dtype's grid as +# a bracket neighbour of the fp32 input. Catches byte-order traps in the +# inline-asm source-register specifier: +# * fp8: cvt.rs.satfinite.e4m3x4.f32 with pack=4, asm "{$4,$3,$2,$1}" +# * fp16: cvt.rs.f16x2.f32 with pack=2, asm "$0, $2, $1, $3" +# The unbiased test (test_philox_rounding_unbiased) wouldn't catch a +# shuffle: outputs that are still on-grid but swapped within a pack still +# average correctly. Only the per-element bracket check exposes it. +# +# Both kernels are inline copies of the production helpers — kept here so +# the test exercises the exact PTX form independent of wrapper changes. + + +@triton.jit +def _packed_int8_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 4)) + y = _stochastic_round_int8_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int8)) + + +@triton.jit +def _packed_int16_sr_kernel(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + (offs // 2)) + y = _stochastic_round_int16_packed(x, rand, offs) + tl.store(out_ptr + offs, y.to(tl.int16)) + - if paged_cache: - cache_buf_idx[slots] = 1 - cache_buf_idx[slots] +def _bitrev_int(x: int, bits: int) -> int: + out = 0 + for _ in range(bits): + out = (out << 1) | (x & 1) + x >>= 1 + return out + + +def _rand_words(rand: torch.Tensor) -> list[int]: + return [int(v) & 0xFFFFFFFF for v in rand.cpu().tolist()] + + +def test_packed_int_sr_matches_reference(): + device = "cuda" + n = 1024 + offs = torch.arange(n, device=device, dtype=torch.float32) + x = ((offs % 37) - 18.0) + (((offs * 13.0) % 97.0) + 0.3) / 128.0 + + torch.manual_seed(42) + rand_i8 = torch.randint(-(2**31), 2**31, (n // 4,), device=device, dtype=torch.int32) + out_i8 = torch.empty(n, device=device, dtype=torch.int8) + _packed_int8_sr_kernel[(1,)](x, rand_i8, out_i8, BLOCK=n) + + x_cpu = x.cpu().tolist() + rand_i8_words = _rand_words(rand_i8) + ref_i8 = [] + for i, value in enumerate(x_cpu): + word = rand_i8_words[i // 4] + low = word & 0x0000FFFF + high = (word >> 16) & 0x0000FFFF + pos = i & 3 + if pos == 0: + rand16 = low + elif pos == 1: + rand16 = _bitrev_int(low, 16) + elif pos == 2: + rand16 = high else: - cache_buf_idx[:] = 1 - cache_buf_idx + rand16 = _bitrev_int(high, 16) + ref_i8.append(math.floor(value + rand16 / float(1 << 16))) - torch.testing.assert_close( - test_out, - ref_outs[step], - rtol=2e-2, - atol=2.0, - msg=f"Output mismatch at step {step} with HPB={heads_per_block}, " - f"T={T}, nheads={nheads}, ngroups={ngroups}, " - f"state_dtype={state_dtype}, paged_cache={paged_cache}", - ) + torch.testing.assert_close( + out_i8.cpu().to(torch.int16), + torch.tensor(ref_i8, dtype=torch.int16), + rtol=0, + atol=0, + ) + + rand_i16 = torch.randint(-(2**31), 2**31, (n // 2,), device=device, dtype=torch.int32) + out_i16 = torch.empty(n, device=device, dtype=torch.int16) + _packed_int16_sr_kernel[(1,)](x, rand_i16, out_i16, BLOCK=n) + + rand_i16_words = _rand_words(rand_i16) + ref_i16 = [] + for i, value in enumerate(x_cpu): + word = rand_i16_words[i // 2] + rand_bits = word if (i & 1) == 0 else _bitrev_int(word, 32) + rand24 = rand_bits & 0x00FFFFFF + ref_i16.append(math.floor(value + rand24 / float(1 << 24))) + + torch.testing.assert_close( + out_i16.cpu(), + torch.tensor(ref_i16, dtype=torch.int16), + rtol=0, + atol=0, + ) + + +@triton.jit +def _bracket_kernel_fp8(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="cvt.rs.satfinite.e4m3x4.f32 $0, {$4, $3, $2, $1}, $5;", + constraints="=r,r,r,r,r,r,r,r,r", + args=(x, rand), + dtype=tl.float8e4nv, + is_pure=True, + pack=4, + ) + tl.store(out_ptr + offs, y) + + +@triton.jit +def _bracket_kernel_fp16(x_ptr, rand_ptr, out_ptr, BLOCK: tl.constexpr): + offs = tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs) + rand = tl.load(rand_ptr + offs) + y = tl.inline_asm_elementwise( + asm="""{ + cvt.rs.f16x2.f32 $0, $2, $1, $3; + }""", + constraints=("=r,r,r,r,r"), + args=(x, rand), + dtype=tl.float16, + is_pure=True, + pack=2, + ) + tl.store(out_ptr + offs, y) + + +_BRACKET_KERNEL = { + torch.float8_e4m3fn: _bracket_kernel_fp8, + torch.float16: _bracket_kernel_fp16, +} + + +def _build_finite_grid(dtype: torch.dtype, device: str) -> torch.Tensor: + """Reinterpret all bit patterns of ``dtype`` as floats; return sorted + unique finite values (drops ±inf, NaNs).""" + if dtype == torch.float8_e4m3fn: + ints = torch.arange(256, dtype=torch.uint8, device=device) + full = ints.view(torch.float8_e4m3fn).to(torch.float32) + elif dtype == torch.float16: + # int16 view of all 65536 patterns (covers fp16 normals + subnormals + # + ±inf + NaN; we filter to finite below). + ints = torch.arange(65536, dtype=torch.int32, device=device).to(torch.int16) + full = ints.view(torch.float16).to(torch.float32) + else: + raise ValueError(f"Unsupported bracket-test dtype: {dtype}") + return full[torch.isfinite(full)].sort()[0].unique() + + +def _build_bracket_inputs(dtype: torch.dtype, n: int, device: str) -> torch.Tensor: + """Test inputs spanning the dtype's grid range. Includes on-grid points + so we exercise the no-rounding case; for fp8 also includes overflow to + test saturation (PTX `cvt.rs.satfinite.e4m3x4.f32` clamps in-op). + + fp16 inputs are kept inside the finite range — `cvt.rs.f16x2.f32` does + NOT have a `satfinite` modifier and produces ±inf for OOR inputs (not + a saturate-to-±max). The kernel only ever sees in-range fp32 state in + practice (state_amax is always ≪ fp16_max), so the test mirrors that. + """ + grid = _build_finite_grid(dtype, device) + g_min, g_max = grid[0].item(), grid[-1].item() + x = torch.empty(n, device=device, dtype=torch.float32) + if dtype == torch.float8_e4m3fn: + # 1.5x range exercises saturation; satfinite handles it in-op. + x.uniform_(g_min * 1.5, g_max * 1.5) + else: # fp16: four magnitude bands, all within finite range. + x[: n // 4].uniform_(-1.0, 1.0) + x[n // 4 : n // 2].uniform_(-100, 100) + x[n // 2 : 3 * n // 4].uniform_(-1000, 1000) + x[3 * n // 4 :].uniform_(g_min * 0.99, g_max * 0.99) + return x, grid + + +@_skip_pre_sm100 +@pytest.mark.parametrize( + "state_dtype", + [torch.float8_e4m3fn, torch.float16], + ids=["fp8", "fp16"], +) +def test_sr_grid_bracket(state_dtype): + """Verify SR PTX outputs each lie on the destination grid as a bracket + neighbour of the fp32 input.""" + device = "cuda" + n = 1024 # multiple of both pack=4 (fp8) and pack=2 (fp16) + + torch.manual_seed(42) + x, grid_finite = _build_bracket_inputs(state_dtype, n, device) + g_min, g_max = grid_finite[0].item(), grid_finite[-1].item() + + # Bracket [lo, hi] in the destination grid for each input. For + # out-of-range inputs the bracket is the saturating endpoint pair. + x_clamped = x.clamp(g_min, g_max) + idx = torch.searchsorted(grid_finite, x_clamped, right=False).clamp( + min=1, max=len(grid_finite) - 1 + ) + lo = grid_finite[idx - 1] + hi = grid_finite[idx] + # For x exactly on grid, idx points at it; lo = grid[i-1], hi = x — the + # bracket allows out==hi (=x) which is what RN-on-grid produces. + + kernel = _BRACKET_KERNEL[state_dtype] + + for seed in range(4): + torch.manual_seed(seed) + # int32 for raw random bits — PTX takes the bit pattern, sign + # interpretation doesn't matter. + rand = torch.randint(-(2**31), 2**31, (n,), device=device, dtype=torch.int32) + out = torch.empty(n, device=device, dtype=state_dtype) + kernel[(1,)](x, rand, out, BLOCK=n) + out_fp32 = out.to(torch.float32) + + on_grid = (out_fp32 == lo) | (out_fp32 == hi) + if not on_grid.all(): + offenders = ~on_grid + n_off = offenders.sum().item() + sample = ( + x[offenders][:5].tolist(), + lo[offenders][:5].tolist(), + hi[offenders][:5].tolist(), + out_fp32[offenders][:5].tolist(), + ) + pytest.fail( + f"{state_dtype} SR output not on grid bracket for {n_off}/{n} " + f"elements (seed={seed}). x={sample[0]} lo={sample[1]} " + f"hi={sample[2]} out={sample[3]}. Likely the PTX byte-order " + "bug (cvt.rs source-register order)." + ) diff --git a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py index b8db6132dc0b..89beba153c45 100644 --- a/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py +++ b/tests/unittest/auto_deploy/singlegpu/custom_ops/mamba/test_flashinfer_mamba_cached_op.py @@ -20,6 +20,13 @@ import tensorrt_llm._torch.auto_deploy # noqa: F401 from tensorrt_llm._torch.auto_deploy.custom_ops.attention_interface import BatchInfo +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from tensorrt_llm._torch.modules.mamba.replay_selective_state_update import ( replay_selective_state_update as _real_replay_selective_state_update, ) @@ -116,6 +123,8 @@ def test_flashinfer_decode_matches_triton(mamba_env): None, # replay_old_da_cumsum None, # replay_cache_buf_idx None, # replay_prev_num_accepted + None, # replay_work_items + None, # replay_n_writes # CONSTANTS time_step_limit, chunk_size, @@ -166,26 +175,41 @@ def test_flashinfer_extend_replay_calls_replay_kernel(mamba_env, head_dim): slot_idx = torch.tensor([0], device=device, dtype=torch.int32) # Replay buffers: all zeros (first step, nothing cached yet; kernel still runs). + replay_history_size = 16 replay_old_x = torch.zeros( - max_batch_size, tokens_per_extend, num_heads, head_dim, device=device, dtype=torch.bfloat16 + max_batch_size, + 2, + replay_history_size, + num_heads, + head_dim, + device=device, + dtype=torch.bfloat16, ) replay_old_b = torch.zeros( max_batch_size, 2, - tokens_per_extend, + replay_history_size, n_groups, ssm_state_size, device=device, dtype=torch.bfloat16, ) replay_old_dt = torch.zeros( - max_batch_size, 2, num_heads, tokens_per_extend, device=device, dtype=torch.float32 + max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 ) replay_old_da_cumsum = torch.zeros( - max_batch_size, 2, num_heads, tokens_per_extend, device=device, dtype=torch.float32 + max_batch_size, 2, num_heads, replay_history_size, device=device, dtype=torch.float32 ) replay_cache_buf_idx = torch.zeros(max_batch_size, device=device, dtype=torch.int32) replay_prev_num_accepted = torch.zeros(max_batch_size, device=device, dtype=torch.int32) + replay_work_items = torch.zeros( + max_batch_size, REPLAY_WORK_ITEM_WIDTH, device=device, dtype=torch.int32 + ) + replay_work_items[0, REPLAY_WORK_POSITION_IN_DECODE_BATCH] = 0 + replay_work_items[0, REPLAY_WORK_CACHE_SLOT] = slot_idx[0] + replay_work_items[0, REPLAY_WORK_PNAT] = 0 + replay_work_items[0, REPLAY_WORK_CACHE_BUF_IDX] = 0 + replay_n_writes = torch.zeros(1, device=device, dtype=torch.int32) # Extend-only batch with replay mode enabled. _bi = BatchInfo() @@ -229,6 +253,8 @@ def test_flashinfer_extend_replay_calls_replay_kernel(mamba_env, head_dim): replay_old_da_cumsum, replay_cache_buf_idx, replay_prev_num_accepted, + replay_work_items, + replay_n_writes, # CONSTANTS time_step_limit, chunk_size, diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py index 2dab50197afc..57ee76d03d10 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_cached_sequence_interface.py @@ -33,11 +33,13 @@ IntermediateSSMStateHandler, KVPagedResourceHandler, ReplayCacheBufIdxHandler, + ReplayNWritesHandler, ReplayOldBHandler, ReplayOldDAcumsumHandler, ReplayOldDtHandler, ReplayOldXHandler, ReplayPrevNumAcceptedHandler, + ReplayWorkItemsHandler, SequenceInfo, SSMResourceHandler, StateResourceHandler, @@ -1390,6 +1392,10 @@ def _add_managed_spec_replay_resources(interface, num_layers=2): replay_names.append( interface.add_resource(f"replay_prev_num_accepted_{i}", ReplayPrevNumAcceptedHandler()) ) + replay_names.append( + interface.add_resource(f"replay_work_items_{i}", ReplayWorkItemsHandler()) + ) + replay_names.append(interface.add_resource(f"replay_n_writes_{i}", ReplayNWritesHandler())) return replay_names diff --git a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py index 2bb5d7fcba3f..b9d97ee45c46 100644 --- a/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py +++ b/tests/unittest/auto_deploy/singlegpu/shim/test_engine.py @@ -12,6 +12,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace from typing import List, Optional, Type import pytest @@ -25,6 +26,13 @@ from tensorrt_llm._torch.auto_deploy.shim.ad_executor import ADEngine from tensorrt_llm._torch.auto_deploy.shim.demollm import DemoEngine from tensorrt_llm._torch.auto_deploy.shim.interface import CachedSequenceInterface +from tensorrt_llm._torch.modules.mamba.mamba2_metadata import ( + REPLAY_WORK_CACHE_BUF_IDX, + REPLAY_WORK_CACHE_SLOT, + REPLAY_WORK_ITEM_WIDTH, + REPLAY_WORK_PNAT, + REPLAY_WORK_POSITION_IN_DECODE_BATCH, +) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.llmapi import AttentionDpConfig @@ -680,6 +688,72 @@ def get_tokens(self, _beam: int) -> List[int]: return self._tokens +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cached_sequence_interface_prepare_replay_metadata_write_first(): + device = torch.device("cuda") + cache_seq_interface = CachedSequenceInterface( + max_seq_len=64, + max_batch_size=4, + max_num_tokens=64, + device=device, + kv_cache_config=KvCacheConfig(tokens_per_block=16), + ) + cache_seq_interface.to(device) + + prev_num_accepted_tokens = torch.zeros(6, dtype=torch.int32, device=device) + cache_buf_idx = torch.zeros(6, dtype=torch.int32, device=device) + slot_idx = torch.tensor([3, 1, 5, 0], dtype=torch.long, device=device) + prev_num_accepted_tokens[slot_idx] = torch.tensor( + [13, 7, 14, 2], dtype=torch.int32, device=device + ) + cache_buf_idx[slot_idx] = torch.tensor([1, 0, 1, 0], dtype=torch.int32, device=device) + + cache_seq_interface._replay_work_items = torch.empty( + cache_seq_interface.info.max_num_state_slots, + REPLAY_WORK_ITEM_WIDTH, + dtype=torch.int32, + device=device, + ) + cache_seq_interface._replay_n_writes = torch.zeros(1, dtype=torch.int32, device=device) + cache_seq_interface._kv_cache_manager = SimpleNamespace( + shutdown=lambda: None, + get_replay_state_update_metadata=lambda: SimpleNamespace( + prev_num_accepted_tokens=prev_num_accepted_tokens, + cache_buf_idx=cache_buf_idx, + replay_step_width=6, + replay_history_size=16, + ), + ) + + cache_seq_interface.info.batch_info.update([0, 0, 4, 4, 0, 0]) + cache_seq_interface.info.batch_info.update_use_replay(True) + cache_seq_interface.info._input_buffer.copy_("slot_idx", slot_idx) + + cache_seq_interface.prepare_replay_metadata() + + expected = torch.tensor( + [ + [0, 3, 13, 1], + [2, 5, 14, 1], + [1, 1, 7, 0], + [3, 0, 2, 0], + ], + dtype=torch.int32, + device=device, + ) + actual = cache_seq_interface._replay_work_items[:4] + assert cache_seq_interface._replay_n_writes.item() == 2 + assert torch.equal( + actual[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH], + expected[:, REPLAY_WORK_POSITION_IN_DECODE_BATCH], + ) + assert torch.equal(actual[:, REPLAY_WORK_CACHE_SLOT], expected[:, REPLAY_WORK_CACHE_SLOT]) + assert torch.equal(actual[:, REPLAY_WORK_PNAT], expected[:, REPLAY_WORK_PNAT]) + assert torch.equal(actual[:, REPLAY_WORK_CACHE_BUF_IDX], expected[:, REPLAY_WORK_CACHE_BUF_IDX]) + + cache_seq_interface.shutdown() + + def test_ad_engine_prepare_inputs_with_hybrid_cache_manager(): """Test ADEngine _prepare_inputs uses mamba_cache_index when available.""" seed = 42