diff --git a/tensorrt_llm/_torch/modules/fla/replay_gated_delta_rule_update.py b/tensorrt_llm/_torch/modules/fla/replay_gated_delta_rule_update.py new file mode 100644 index 000000000000..83f04a8c2ba0 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fla/replay_gated_delta_rule_update.py @@ -0,0 +1,314 @@ +# 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. +# +# Gated-delta-rule (GDN) replay state update for MTP speculative decoding. +# +# This is the GDN analogue of the Mamba SSM +# ``replay_selective_state_update`` kernel. Instead of caching a full +# intermediate recurrent state for every draft position (the "legacy" MTP +# path), it keeps a compact double-buffered *history* of the raw per-token +# inputs needed to advance the gated-delta-rule recurrence (k, v, g, beta) +# and "replays" (re-runs) that history from a checkpoint to reconstruct the +# committed state at the start of each verification step. +# +# Contract shared with the Mamba replay path (see +# ``pyexecutor/mamba_cache_manager.py`` ``update_mamba_states`` and +# ``modules/mamba/replay_selective_state_update.py``): +# * ``temporal`` (``ssm_states``) holds a *checkpoint* recurrent state. +# * ``prev_num_accepted_tokens[slot]`` (PNAT) counts committed-but-unfolded +# tokens whose inputs live in the active history buffer at ``[0, PNAT)``. +# * ``cache_buf_idx[slot]`` selects the active history buffer (double +# buffered along a size-2 dimension). +# * ``wrote_checkpoint = PNAT + T > HISTORY`` decides whether this step +# folds the PNAT tokens into ``temporal`` and stages the current tokens in +# the other buffer (write path) or appends the current tokens to the active +# buffer (nowrite path). ``update_mamba_states`` bumps PNAT / flips +# ``cache_buf_idx`` after acceptance using the *same* predicate. + +from typing import Optional + +import torch +import triton +import triton.language as tl + +from tensorrt_llm._torch.modules.fla.op import exp + + +@triton.jit +def _gdc_wait(): + tl.inline_asm_elementwise( + "griddepcontrol.wait; // dummy $0", + "=r,~{memory}", + [], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + + +@triton.jit(do_not_specialize=["T"]) +def _replay_gated_delta_rule_update_kernel( + q, + k, + v, + g, + beta, + o, + ssm_states, + old_k, + old_v, + old_g, + old_beta, + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + scale, + T, + s_state_slot, + s_state_hv, + s_state_v, + s_state_k, + HISTORY: tl.constexpr, + H: tl.constexpr, + HV: tl.constexpr, + K: tl.constexpr, + V: tl.constexpr, + BK: tl.constexpr, + BV: tl.constexpr, + USE_QK_L2NORM_IN_KERNEL: tl.constexpr, + LAUNCH_WITH_PDL: tl.constexpr, +): + if LAUNCH_WITH_PDL: + _gdc_wait() + + i_v = tl.program_id(0) + i_ndhv = tl.program_id(1) + i_nd = i_ndhv // HV + i_hv = i_ndhv % HV + # GVA: map value head to its key/query head group. + i_h = i_hv // (HV // H) + + slot = tl.load(state_batch_indices + i_nd).to(tl.int64) + pnat = tl.load(prev_num_accepted_tokens + slot).to(tl.int32) + buf = tl.load(cache_buf_idx + slot).to(tl.int32) + + wrote_checkpoint = (pnat + T) > HISTORY + write_buf = tl.where(wrote_checkpoint, 1 - buf, buf).to(tl.int64) + write_off = tl.where(wrote_checkpoint, 0, pnat) + + o_k = tl.arange(0, BK) + o_v = i_v * BV + tl.arange(0, BV) + mask_k = o_k < K + mask_v = o_v < V + mask_h = mask_k[:, None] & mask_v[None, :] + + # Checkpoint state from the pool: logical layout [slots, HV, V, K]. + # Strides are explicit so this works for both the dense Mixed-manager pool + # and the strided unified C++ pool view. + p_h = (ssm_states + slot * s_state_slot + i_hv * s_state_hv + + o_v[None, :] * s_state_v + o_k[:, None] * s_state_k) + b_h = tl.load(p_h, mask=mask_h, other=0.0).to(tl.float32) + + # History buffer base pointers for the active (read) buffer. + # old_k: [cache, 2, HISTORY, H, K] + # old_v: [cache, 2, HISTORY, HV, V] + # old_g: [cache, 2, HISTORY, HV] + # old_beta: [cache, 2, HISTORY, HV] + buf64 = buf.to(tl.int64) + rbase_k = old_k + (slot * 2 + buf64) * HISTORY * H * K + rbase_v = old_v + (slot * 2 + buf64) * HISTORY * HV * V + rbase_g = old_g + (slot * 2 + buf64) * HISTORY * HV + rbase_beta = old_beta + (slot * 2 + buf64) * HISTORY * HV + + # --- Replay committed-but-unfolded tokens [0, PNAT) to reconstruct the + # state at the start of this verification step (state only, no output). + for t in range(0, pnat): + b_k = tl.load(rbase_k + t * H * K + i_h * K + o_k, mask=mask_k, + other=0.0).to(tl.float32) + b_v = tl.load(rbase_v + t * HV * V + i_hv * V + o_v, mask=mask_v, + other=0.0).to(tl.float32) + b_g = tl.load(rbase_g + t * HV + i_hv).to(tl.float32) + b_beta = tl.load(rbase_beta + t * HV + i_hv).to(tl.float32) + if USE_QK_L2NORM_IN_KERNEL: + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) + b_h *= exp(b_g) + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= b_beta + b_h += b_k[:, None] * b_v[None, :] + + # State immediately before the current tokens = new checkpoint (write path). + b_h_checkpoint = b_h + + # History buffer base pointers for the write buffer. + wbase_k = old_k + (slot * 2 + write_buf) * HISTORY * H * K + wbase_v = old_v + (slot * 2 + write_buf) * HISTORY * HV * V + wbase_g = old_g + (slot * 2 + write_buf) * HISTORY * HV + wbase_beta = old_beta + (slot * 2 + write_buf) * HISTORY * HV + + # k / g / beta are shared across V blocks (and k across the GVA group), so + # only one program writes each to avoid redundant traffic. + is_first_hv_in_group = (i_hv % (HV // H)) == 0 + + for t in range(0, T): + b_q = tl.load(q + ((i_nd * T + t) * H + i_h) * K + o_k, mask=mask_k, + other=0.0).to(tl.float32) + b_k = tl.load(k + ((i_nd * T + t) * H + i_h) * K + o_k, mask=mask_k, + other=0.0).to(tl.float32) + b_v = tl.load(v + ((i_nd * T + t) * HV + i_hv) * V + o_v, mask=mask_v, + other=0.0).to(tl.float32) + b_g = tl.load(g + (i_nd * T + t) * HV + i_hv).to(tl.float32) + b_beta = tl.load(beta + (i_nd * T + t) * HV + i_hv).to(tl.float32) + + # Append raw inputs to the write buffer at write_off + t. + wpos = write_off + t + if i_v == 0: + if is_first_hv_in_group: + tl.store(wbase_k + wpos * H * K + i_h * K + o_k, + b_k.to(old_k.dtype.element_ty), + mask=mask_k) + tl.store(wbase_g + wpos * HV + i_hv, b_g.to(old_g.dtype.element_ty)) + tl.store(wbase_beta + wpos * HV + i_hv, + b_beta.to(old_beta.dtype.element_ty)) + tl.store(wbase_v + wpos * HV * V + i_hv * V + o_v, + b_v.to(old_v.dtype.element_ty), + mask=mask_v) + + if USE_QK_L2NORM_IN_KERNEL: + b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q)) + 1e-6) + b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k)) + 1e-6) + b_q = b_q * scale + b_h *= exp(b_g) + b_v -= tl.sum(b_h * b_k[:, None], 0) + b_v *= b_beta + b_h += b_k[:, None] * b_v[None, :] + b_o = tl.sum(b_h * b_q[:, None], 0) + tl.store(o + ((i_nd * T + t) * HV + i_hv) * V + o_v, + b_o.to(o.dtype.element_ty), + mask=mask_v) + + # On the write path, fold the PNAT tokens into the pool checkpoint. On the + # nowrite path the checkpoint is left untouched (temporal keeps lagging). + if wrote_checkpoint: + tl.store(p_h, + b_h_checkpoint.to(ssm_states.dtype.element_ty), + mask=mask_h) + + +def replay_gated_delta_rule_update( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + ssm_states: torch.Tensor, + old_k: torch.Tensor, + old_v: torch.Tensor, + old_g: torch.Tensor, + old_beta: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + state_batch_indices: torch.Tensor, + replay_step_width: int, + replay_history_size: int, + scale: Optional[float] = None, + use_qk_l2norm_in_kernel: bool = True, + output: Optional[torch.Tensor] = None, + launch_with_pdl: bool = False, +) -> torch.Tensor: + """Fused GDN replay state update for the MTP target-verify decode path. + + Args: + q, k: current-step queries/keys, shape ``[num_decodes, T, H, K]``. + v: current-step values, shape ``[num_decodes, T, HV, V]``. + g, beta: per-head decay / beta, shape ``[num_decodes, T, HV]``. + ssm_states: recurrent state pool ``[slots, HV, V, K]`` (checkpoint, + updated in place on the write path). + old_k, old_v, old_g, old_beta: this layer's double-buffered history + caches with a leading ``[cache, 2, HISTORY, ...]`` layout. + cache_buf_idx: ``[cache]`` int active-buffer selector (read only here; + flipped by ``update_mamba_states`` after acceptance). + prev_num_accepted_tokens: ``[cache]`` int PNAT (read only here). + state_batch_indices: ``[num_decodes]`` int cache slot per decode. + replay_step_width: fixed ``T`` (``runtime_draft_len + 1``). + replay_history_size: history buffer capacity along the window dim. + output: optional preallocated ``[num_decodes, T, HV, V]`` output. + + Returns: + Output tensor ``[num_decodes, T, HV, V]``. + """ + num_decodes, T, H, K = q.shape + HV = v.shape[2] + V = v.shape[3] + assert T == replay_step_width, ( + f"runtime token width {T} must match fixed replay step width " + f"{replay_step_width}") + assert k.shape == (num_decodes, T, H, K) + assert v.shape == (num_decodes, T, HV, V) + assert g.shape[-1] == HV and beta.shape[-1] == HV + + if scale is None: + scale = K**-0.5 + + BK = triton.next_power_of_2(K) + BV = min(triton.next_power_of_2(V), 8) + NV = triton.cdiv(V, BV) + assert triton.cdiv(K, BK) == 1, "K larger than one tensor tile is unsupported" + + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + g = g.contiguous() + beta = beta.contiguous() + + if output is None: + o = q.new_empty(num_decodes, T, HV, V) + else: + o = output + + grid = (NV, num_decodes * HV) + _replay_gated_delta_rule_update_kernel[grid]( + q, + k, + v, + g, + beta, + o, + ssm_states, + old_k, + old_v, + old_g, + old_beta, + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + scale, + T, + ssm_states.stride(0), + ssm_states.stride(1), + ssm_states.stride(2), + ssm_states.stride(3), + HISTORY=replay_history_size, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + LAUNCH_WITH_PDL=launch_with_pdl, + num_warps=1, + num_stages=3, + ) + return o diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index f81f92cbf945..0418b45d52d6 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -20,6 +20,9 @@ _flashinfer_gdn_verify, fused_sigmoid_gating_delta_rule_update, ) +from tensorrt_llm._torch.modules.fla.replay_gated_delta_rule_update import ( + replay_gated_delta_rule_update, +) from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import use_cpp_mamba_cache_manager from tensorrt_llm._utils import is_flashinfer_gdn_supported_arch from tensorrt_llm.mapping import Mapping @@ -602,6 +605,61 @@ def _postprocess_gdn_output( attn_out = attn_out.reshape(*attn_out.shape[:-2], -1) return self.out_proj(attn_out, all_reduce_params=all_reduce_params) + def _replay_gdn_verify( + self, + query, + key, + value, + a, + b, + ssm_states, + state_batch_indices, + num_decodes, + draft_token_num, + replay_ctx, + output_d, + ): + """MTP target-verify decode using the GDN replay state update. + + query/key are ``[num_decodes, T, H, K]``, value is + ``[num_decodes, T, HV, V]`` and a/b are ``[num_decodes, T, HV]``. + Returns the flattened output ``[1, num_decodes * T, HV, V]``. + """ + beta = b.sigmoid() + g = fused_gdn_gating( + self.A_log, + a.reshape(num_decodes * draft_token_num, -1), + self.dt_bias, + ).reshape(num_decodes, draft_token_num, -1) + + attn_out = replay_gated_delta_rule_update( + q=query, + k=key, + v=value, + g=g, + beta=beta, + ssm_states=ssm_states, + old_k=replay_ctx["old_k"], + old_v=replay_ctx["old_v"], + old_g=replay_ctx["old_g"], + old_beta=replay_ctx["old_beta"], + cache_buf_idx=replay_ctx["cache_buf_idx"], + prev_num_accepted_tokens=replay_ctx["prev_num_accepted_tokens"], + state_batch_indices=state_batch_indices[:num_decodes].to( + torch.int32), + replay_step_width=replay_ctx["replay_step_width"], + replay_history_size=replay_ctx["replay_history_size"], + scale=self.head_k_dim**-0.5, + use_qk_l2norm_in_kernel=True, + output=output_d, + ) + return attn_out.view( + 1, + num_decodes * draft_token_num, + self.num_v_heads // self.attn_tp_size, + self.head_v_dim, + ) + def forward_decode( self, conv_states, @@ -619,6 +677,8 @@ def forward_decode( b = kwargs["b"] cache_indices = kwargs["cache_indices"] num_decodes = kwargs["num_decodes"] + use_replay = kwargs.get("use_replay", False) + replay_ctx = kwargs.get("replay_ctx", None) if is_target_verify: draft_token_num = spec_metadata.runtime_draft_len + 1 @@ -627,7 +687,9 @@ def forward_decode( assert a.shape[0] == num_decodes * draft_token_num assert b.shape[0] == num_decodes * draft_token_num assert intermediate_conv_states is not None - assert intermediate_ssm_states is not None + # Replay reconstructs state from its double-buffered history and + # does not allocate the legacy intermediate SSM buffer. + assert use_replay or intermediate_ssm_states is not None # Speculative verification path: # 1. run conv update with per-step intermediate cache writes @@ -670,6 +732,31 @@ def forward_decode( a = a.reshape(num_decodes, draft_token_num, -1) b = b.reshape(num_decodes, draft_token_num, -1) + # Replay path: reconstruct the committed state from the compact + # double-buffered history and update the pool checkpoint in place. + if use_replay: + output_d = None + if output is not None: + output_d = output.view( + num_decodes, + draft_token_num, + self.num_v_heads // self.attn_tp_size, + self.head_v_dim, + ) + return self._replay_gdn_verify( + query=query, + key=key, + value=value, + a=a, + b=b, + ssm_states=ssm_states, + state_batch_indices=cache_indices, + num_decodes=num_decodes, + draft_token_num=draft_token_num, + replay_ctx=replay_ctx, + output_d=output_d, + ) + # Prefer the FlashInfer MTP kernel (raw a/b gating in-kernel, # initial state gathered from the pool via cache indices, per-step # intermediate states written to the batch-scoped [:num_decodes] @@ -818,6 +905,8 @@ def forward_extend( state_indices_d = kwargs["state_indices_d"] num_prefill = kwargs["num_prefill"] num_decodes = kwargs["num_decodes"] + use_replay = kwargs.get("use_replay", False) + replay_ctx = kwargs.get("replay_ctx", None) conv_states_to_use = conv_states @@ -853,7 +942,9 @@ def forward_extend( assert a_d.shape[0] == num_decodes * draft_token_num assert b_d.shape[0] == num_decodes * draft_token_num assert intermediate_conv_states is not None - assert intermediate_ssm_states is not None + # Replay reconstructs state from its double-buffered history and + # does not allocate the legacy intermediate SSM buffer. + assert use_replay or intermediate_ssm_states is not None intermediate_state_indices = torch.arange( num_decodes, dtype=torch.int32, device=state_indices_d.device @@ -974,7 +1065,21 @@ def forward_extend( self.head_v_dim, ) - if _can_use_flashinfer_gdn_verify( + if use_replay: + attn_out_decode = self._replay_gdn_verify( + query=query_d, + key=key_d, + value=value_d, + a=a_d, + b=b_d, + ssm_states=ssm_states, + state_batch_indices=state_indices_d, + num_decodes=num_decodes, + draft_token_num=draft_token_num, + replay_ctx=replay_ctx, + output_d=output_d, + ).reshape(1, num_decode_tokens, out_v_heads, self.head_v_dim) + elif _can_use_flashinfer_gdn_verify( ssm_states, self.head_k_dim, self.head_v_dim, draft_token_num ): # FI gathers the initial state from the pool via state_indices_d @@ -1136,7 +1241,32 @@ def forward_core( intermediate_conv_states = ( layer_cache.intermediate_conv_window if is_target_verify else None ) - intermediate_ssm_states = layer_cache.intermediate_ssm if is_target_verify else None + # Replay MTP path: reconstruct the committed recurrent state from a + # compact double-buffered history instead of caching full per-draft-step + # intermediate states. When active, intermediate_ssm is not allocated. + use_replay = is_target_verify and getattr( + attn_metadata.kv_cache_manager, "use_replay_state_update", False) + replay_ctx = None + intermediate_ssm_states = None + if use_replay: + replay_meta = ( + attn_metadata.kv_cache_manager.get_replay_state_update_metadata()) + assert replay_meta is not None, ( + "GDN replay state update is enabled but replay metadata was " + "not allocated.") + replay_ctx = { + "old_k": layer_cache.old_k, + "old_v": layer_cache.old_v, + "old_g": layer_cache.old_g, + "old_beta": layer_cache.old_beta, + "cache_buf_idx": layer_cache.cache_buf_idx, + "prev_num_accepted_tokens": + layer_cache.prev_num_accepted_tokens, + "replay_step_width": replay_meta.replay_step_width, + "replay_history_size": replay_meta.replay_history_size, + } + elif is_target_verify: + intermediate_ssm_states = layer_cache.intermediate_ssm kwargs = { "mixed_qkv": mixed_qkv, @@ -1153,6 +1283,8 @@ def forward_core( "state_indices_d": state_indices_d, "num_prefill": num_prefills, "num_decodes": num_decodes, + "use_replay": use_replay, + "replay_ctx": replay_ctx, } if num_prefills > 0: attn_out = self.forward_extend( diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 48d0c0dccdc4..f3bec52afa25 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1952,6 +1952,30 @@ def _create_kv_cache_manager( spec_config=spec_config, quant_config=quant_config, ) + + # GDN replay state update kernel for MTP: default on for sm >= 80. + # bf16 recurrent state only; stochastic rounding is not supported on + # this path, so (unlike Nemotron) there is no fp16/Philox gate here. + sm = get_sm_version() + use_replay = spec_config is not None and sm >= 80 + if spec_config is None: + logger.info( + "GDN replay kernel requires speculative decoding; using " + "non-replay path") + + # Tree attention: replay assumes a linear token sequence. + if (spec_config is not None + and (getattr(spec_config, 'eagle_choices', None) is not None + or getattr(spec_config, 'use_dynamic_tree', False))): + logger.info("GDN replay kernel incompatible with tree attention; " + "using legacy MTP path") + use_replay = False + + if os.environ.get('TRTLLM_USE_MAMBA_REPLAY', '1') == '0': + logger.info( + "GDN replay kernel is disabled by TRTLLM_USE_MAMBA_REPLAY=0") + use_replay = False + kv_cache_manager = kv_cache_manager_cls( # mamba cache parameters mamba_params.state_size, @@ -1980,6 +2004,7 @@ def _create_kv_cache_manager( is_estimating_kv_cache=estimating_kv_cache, execution_stream=execution_stream, model_type="qwen3_next", + use_replay_state_update=use_replay, **manager_extra_kwargs, ) else: diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 54be9fd150d5..11474bbc4eae 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -377,7 +377,9 @@ class SpeculativeState(State): Supports two SSM update paths (only one set of tensors is allocated): - Legacy: caches full intermediate SSM states (intermediate_ssm) - - Replay: compact double-buffered cache (old_x, old_B, old_dt, old_dA_cumsum) + - Replay: compact double-buffered history cache. The Mamba2 selective + scan stores (old_x, old_B, old_dt, old_dA_cumsum); the GDN gated + delta rule stores (old_k, old_v, old_g, old_beta). """ _SHARED_FIELDS = frozenset({ "prev_num_accepted_tokens", "cache_buf_idx", "mamba_ssm_rand_seed" @@ -397,11 +399,17 @@ 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 + # Mamba2 selective-scan replay history. 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, history) fp32 old_dA_cumsum: torch.Tensor | None = None # (layers, cache, 2, nheads, history) fp32 + # GDN gated-delta-rule replay history (raw per-token inputs). + old_k: torch.Tensor | None = None # (layers, cache, 2, history, num_k_heads, head_k_dim) + old_v: torch.Tensor | None = None # (layers, cache, 2, history, num_v_heads, head_v_dim) + old_g: torch.Tensor | None = None # (layers, cache, 2, history, num_v_heads) fp32 + old_beta: torch.Tensor | None = None # (layers, cache, 2, history, num_v_heads) fp32 def __init__( self, @@ -426,6 +434,7 @@ def __init__( self.mamba_ssm_cache_dtype = ssm_cache_dtype self.speculative_num_draft_tokens = speculative_num_draft_tokens self.spec_state_size = spec_state_size + self._model_type = model_type self._use_replay_state_update = use_replay_state_update self.replay_history_size: Optional[int] = None self.replay_step_width: Optional[int] = None @@ -535,47 +544,95 @@ def __init__( n_groups_per_rank = n_groups // tp_size self.replay_history_size = max(MIN_REPLAY_HISTORY_SIZE, T) - # Compact replay cache. + # Shared double-buffering bookkeeping (both architectures). spec_kwargs['prev_num_accepted_tokens'] = torch.zeros( 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, - 2, - self.replay_history_size, - nheads, - head_dim, - dtype=dtype, - device=device) - spec_kwargs['old_B'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - self.replay_history_size, - n_groups_per_rank, - d_state, - dtype=dtype, - device=device) - spec_kwargs['old_dt'] = torch.zeros(num_local_layers, - max_batch_size, - 2, - nheads, - self.replay_history_size, - 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'] - ] + + if model_type == "qwen3_next": + # GDN gated-delta-rule history: raw per-token k/v/g/beta. + # Here nheads == num_v_heads (HV), n_groups == num_k_heads + # (H), head_dim == head_v_dim (V), d_state == head_k_dim (K). + spec_kwargs['old_k'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + n_groups_per_rank, + d_state, + dtype=dtype, + device=device) + spec_kwargs['old_v'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + head_dim, + dtype=dtype, + device=device) + spec_kwargs['old_g'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + dtype=torch.float32, + device=device) + spec_kwargs['old_beta'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + dtype=torch.float32, + device=device) + ssm_spec_cache = [ + spec_kwargs['old_k'], spec_kwargs['old_v'], + spec_kwargs['old_g'], spec_kwargs['old_beta'] + ] + else: + # Mamba2 selective-scan history. + spec_kwargs['old_x'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + nheads, + head_dim, + dtype=dtype, + device=device) + spec_kwargs['old_B'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + self.replay_history_size, + n_groups_per_rank, + d_state, + dtype=dtype, + device=device) + spec_kwargs['old_dt'] = torch.zeros( + num_local_layers, + max_batch_size, + 2, + nheads, + self.replay_history_size, + 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'] + ] spec_path_label = "replay" else: # Legacy: full intermediate SSM states at each step @@ -838,6 +895,30 @@ def get_replay_old_dA_cumsum(self, layer_offset = self.mamba_layer_offsets[layer_idx] return self.mamba_cache.old_dA_cumsum[layer_offset] + def get_replay_old_k(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_k[layer_offset] + + def get_replay_old_v(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_v[layer_offset] + + def get_replay_old_g(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_g[layer_offset] + + def get_replay_old_beta(self, layer_idx: int) -> Optional[torch.Tensor]: + if not self._use_replay_state_update: + return None + layer_offset = self.mamba_layer_offsets[layer_idx] + return self.mamba_cache.old_beta[layer_offset] + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: if not self._use_replay_state_update: return None @@ -911,6 +992,10 @@ def _drop(tensor): old_B=_drop(self.mamba_cache.old_B), old_dt=_drop(self.mamba_cache.old_dt), old_dA_cumsum=_drop(self.mamba_cache.old_dA_cumsum), + old_k=_drop(self.mamba_cache.old_k), + old_v=_drop(self.mamba_cache.old_v), + old_g=_drop(self.mamba_cache.old_g), + old_beta=_drop(self.mamba_cache.old_beta), ) else: self.mamba_cache = self.State(conv=empty, temporal=empty) @@ -1132,6 +1217,22 @@ def get_replay_old_dA_cumsum(self, assert not self._use_cpp, "get_replay_old_dA_cumsum is not supported in CppMambaCacheManager" return self._impl.get_replay_old_dA_cumsum(layer_idx) + def get_replay_old_k(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_k is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_k(layer_idx) + + def get_replay_old_v(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_v is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_v(layer_idx) + + def get_replay_old_g(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_g is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_g(layer_idx) + + def get_replay_old_beta(self, layer_idx: int) -> Optional[torch.Tensor]: + assert not self._use_cpp, "get_replay_old_beta is not supported in CppMambaCacheManager" + return self._impl.get_replay_old_beta(layer_idx) + def get_replay_cache_buf_idx(self) -> Optional[torch.Tensor]: assert not self._use_cpp, "get_replay_cache_buf_idx is not supported in CppMambaCacheManager" return self._impl.get_replay_cache_buf_idx() @@ -1875,6 +1976,14 @@ def _prepare_resources(self, scheduled_batch: ScheduledRequests): self.old_dt[:, ctx_slots] = 0 if self.old_dA_cumsum is not None: self.old_dA_cumsum[:, ctx_slots] = 0 + if self.old_k is not None: + self.old_k[:, ctx_slots] = 0 + if self.old_v is not None: + self.old_v[:, ctx_slots] = 0 + if self.old_g is not None: + self.old_g[:, ctx_slots] = 0 + if self.old_beta is not None: + self.old_beta[:, ctx_slots] = 0 # Deterministic per-context-slot seed rotation. Runs whenever # the seed buffer exists, including the non-replay SR path. # Bump the host counter once per batch and write one new seed @@ -2067,10 +2176,17 @@ def mamba_layer_cache( # Per-layer slices for the replay kernel; shared 1D tensors # (cache_buf_idx, prev_num_accepted_tokens) are passed # untouched via the SpeculativeState._SHARED_FIELDS contract. - spec_kwargs['old_x'] = self.old_x[layer_offset] - spec_kwargs['old_B'] = self.old_B[layer_offset] - spec_kwargs['old_dt'] = self.old_dt[layer_offset] - spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[layer_offset] + if self._rnn_conv_section_layout == "qwen3_next": + spec_kwargs['old_k'] = self.old_k[layer_offset] + spec_kwargs['old_v'] = self.old_v[layer_offset] + spec_kwargs['old_g'] = self.old_g[layer_offset] + spec_kwargs['old_beta'] = self.old_beta[layer_offset] + else: + spec_kwargs['old_x'] = self.old_x[layer_offset] + spec_kwargs['old_B'] = self.old_B[layer_offset] + spec_kwargs['old_dt'] = self.old_dt[layer_offset] + spec_kwargs['old_dA_cumsum'] = self.old_dA_cumsum[ + layer_offset] spec_kwargs['cache_buf_idx'] = self.cache_buf_idx spec_kwargs['prev_num_accepted_tokens'] = ( self.prev_num_accepted_tokens) @@ -2287,6 +2403,10 @@ def _setup_replay_buffers(self, spec_config) -> None: self.old_B = None self.old_dt = None self.old_dA_cumsum = None + self.old_k = None + self.old_v = None + self.old_g = None + self.old_beta = None if (not self._use_replay_state_update and not self._mamba_ssm_stochastic_rounding): @@ -2310,6 +2430,10 @@ def _setup_replay_buffers(self, spec_config) -> None: self.old_B = None self.old_dt = None self.old_dA_cumsum = None + self.old_k = None + self.old_v = None + self.old_g = None + self.old_beta = None self._dummy_request_mask = None self._dummy_request_mask_host = None return @@ -2332,37 +2456,72 @@ def _setup_replay_buffers(self, spec_config) -> None: 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, - 2, - history_size, - nheads, - head_dim, - dtype=self.conv_state_dtype, - device=device) - # Per-layer double-buffered caches. - self.old_B = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - history_size, - n_groups_per_rank, - d_state, - dtype=self.conv_state_dtype, - device=device) - self.old_dt = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - nheads, - history_size, - dtype=torch.float32, - device=device) - self.old_dA_cumsum = torch.zeros(num_local_mamba_layers, - cache_size, - 2, - nheads, - history_size, - dtype=torch.float32, - device=device) + if self._rnn_conv_section_layout == "qwen3_next": + # GDN gated-delta-rule history: raw per-token k/v/g/beta. Here + # nheads == num_v_heads (HV), n_groups == num_k_heads (H), + # head_dim == head_v_dim (V), d_state == head_k_dim (K). + self.old_k = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + history_size, + n_groups_per_rank, + d_state, + dtype=self.conv_state_dtype, + device=device) + self.old_v = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + history_size, + nheads, + head_dim, + dtype=self.conv_state_dtype, + device=device) + self.old_g = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + history_size, + nheads, + dtype=torch.float32, + device=device) + self.old_beta = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + history_size, + nheads, + dtype=torch.float32, + device=device) + else: + self.old_x = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + history_size, + nheads, + head_dim, + dtype=self.conv_state_dtype, + device=device) + # Per-layer double-buffered caches. + self.old_B = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + history_size, + n_groups_per_rank, + d_state, + dtype=self.conv_state_dtype, + device=device) + self.old_dt = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + nheads, + history_size, + dtype=torch.float32, + device=device) + self.old_dA_cumsum = torch.zeros(num_local_mamba_layers, + cache_size, + 2, + nheads, + history_size, + dtype=torch.float32, + device=device) @property def use_replay_state_update(self) -> bool: diff --git a/tests/unittest/_torch/modules/mamba/test_replay_gated_delta_rule_update.py b/tests/unittest/_torch/modules/mamba/test_replay_gated_delta_rule_update.py new file mode 100644 index 000000000000..6e65c4db9c77 --- /dev/null +++ b/tests/unittest/_torch/modules/mamba/test_replay_gated_delta_rule_update.py @@ -0,0 +1,314 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parity tests for the GDN MTP *replay* state update. + +The replay path reconstructs the committed recurrent state from a compact +double-buffered history (old_k/old_v/old_g/old_beta) instead of caching a full +intermediate state per draft position. These tests assert that: + +1. A single verify step (PNAT == 0) matches the legacy + ``fused_recurrent_gated_delta_rule_update`` reference exactly (fp32). +2. A multi-step MTP protocol (with buffer flips / checkpoints) keeps producing + outputs that match a golden recurrence tracked over the committed tokens. +3. The kernel is CUDA-graph capturable and replays deterministically. +""" + +import pytest +import torch + +skip_no_cuda = pytest.mark.skipif(not torch.cuda.is_available(), + reason="Requires CUDA") + + +def _gating(A_log, a, dt_bias, N, T, HV): + from tensorrt_llm._torch.modules.mamba.gdn_mixer import fused_gdn_gating + return fused_gdn_gating(A_log, a.reshape(N * T, HV), dt_bias).reshape( + N, T, HV) + + +def _rand_step(N, T, H, HV, K, V, dev, dtype, scale_amp=0.1): + q = (torch.randn(N, T, H, K, device=dev) * scale_amp).to(dtype) + k = (torch.randn(N, T, H, K, device=dev) * scale_amp).to(dtype) + v = (torch.randn(N, T, HV, V, device=dev) * scale_amp).to(dtype) + a = torch.randn(N, T, HV, device=dev) * scale_amp + b = torch.randn(N, T, HV, device=dev) * scale_amp + return q, k, v, a, b + + +@skip_no_cuda +@pytest.mark.parametrize("draft_token_num", [1, 2, 4, 5]) +@pytest.mark.parametrize("num_decodes", [1, 3]) +@pytest.mark.parametrize("H,HV", [(4, 8), (2, 2)]) +def test_replay_single_step_matches_recurrent(draft_token_num, num_decodes, H, + HV): + """PNAT == 0: replay output equals the fused_recurrent reference (fp32).""" + from tensorrt_llm._torch.modules.fla.fused_recurrent import ( + fused_recurrent_gated_delta_rule_update, ) + from tensorrt_llm._torch.modules.fla.replay_gated_delta_rule_update import ( + replay_gated_delta_rule_update, ) + + torch.manual_seed(0) + dev = "cuda" + N, T, K, V = num_decodes, draft_token_num, 128, 128 + HISTORY = max(16, T) + scale = K**-0.5 + dtype = torch.float32 + + q, k, v, a, b = _rand_step(N, T, H, HV, K, V, dev, dtype) + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + g = _gating(A_log, a, dt_bias, N, T, HV) + beta = b.sigmoid() + + state_pool = torch.randn(N, HV, V, K, device=dev, dtype=dtype) * 0.1 + idx = torch.arange(N, device=dev, dtype=torch.int32) + + # Reference from the committed (initial) state. + ref = fused_recurrent_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state_source=state_pool.clone(), + initial_state_indices=idx, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + ) + + old_k = torch.zeros(N, 2, HISTORY, H, K, device=dev, dtype=dtype) + old_v = torch.zeros(N, 2, HISTORY, HV, V, device=dev, dtype=dtype) + old_g = torch.zeros(N, 2, HISTORY, HV, device=dev, dtype=torch.float32) + old_beta = torch.zeros(N, 2, HISTORY, HV, device=dev, dtype=torch.float32) + cache_buf_idx = torch.zeros(N, device=dev, dtype=torch.int32) + pnat = torch.zeros(N, device=dev, dtype=torch.int32) + state_replay = state_pool.clone() + + out = replay_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + ssm_states=state_replay, + old_k=old_k, + old_v=old_v, + old_g=old_g, + old_beta=old_beta, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=pnat, + state_batch_indices=idx, + replay_step_width=T, + replay_history_size=HISTORY, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + + torch.testing.assert_close(out.view(N, T, HV, V), + ref.view(N, T, HV, V), + rtol=1e-4, + atol=1e-4) + # PNAT == 0 leaves the checkpoint untouched (replaying zero history tokens). + torch.testing.assert_close(state_replay, state_pool, rtol=1e-4, atol=1e-4) + + +def _replay_protocol(dtype, rtol, atol, accepted_pattern, T, HISTORY, H=4, HV=8): + """Run the full MTP replay protocol and compare per-step outputs to a golden + recurrence tracked over the committed tokens.""" + from tensorrt_llm._torch.modules.fla.fused_recurrent import ( + fused_recurrent_gated_delta_rule_update, ) + from tensorrt_llm._torch.modules.fla.replay_gated_delta_rule_update import ( + replay_gated_delta_rule_update, ) + + torch.manual_seed(1) + dev = "cuda" + N, K, V = 3, 128, 128 + scale = K**-0.5 + + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + idx = torch.arange(N, device=dev, dtype=torch.int32) + + # Golden committed state (advanced by accepted tokens each step). + golden_state = torch.randn(N, HV, V, K, device=dev, dtype=dtype) * 0.1 + # Replay checkpoint starts equal to the committed state. + state_replay = golden_state.clone() + + old_k = torch.zeros(N, 2, HISTORY, H, K, device=dev, dtype=dtype) + old_v = torch.zeros(N, 2, HISTORY, HV, V, device=dev, dtype=dtype) + old_g = torch.zeros(N, 2, HISTORY, HV, device=dev, dtype=torch.float32) + old_beta = torch.zeros(N, 2, HISTORY, HV, device=dev, dtype=torch.float32) + cache_buf_idx = torch.zeros(N, device=dev, dtype=torch.int32) + pnat = torch.zeros(N, device=dev, dtype=torch.int32) + + for accepted in accepted_pattern: + q, k, v, a, b = _rand_step(N, T, H, HV, K, V, dev, dtype) + g = _gating(A_log, a, dt_bias, N, T, HV) + beta = b.sigmoid() + + # Golden: outputs for all T tokens from the committed state. + ref = fused_recurrent_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + scale=scale, + initial_state_source=golden_state.clone(), + initial_state_indices=idx, + use_qk_l2norm_in_kernel=True, + disable_state_update=True, + ) + + out = replay_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + ssm_states=state_replay, + old_k=old_k, + old_v=old_v, + old_g=old_g, + old_beta=old_beta, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=pnat, + state_batch_indices=idx, + replay_step_width=T, + replay_history_size=HISTORY, + scale=scale, + use_qk_l2norm_in_kernel=True, + ) + torch.testing.assert_close(out.view(N, T, HV, V), + ref.view(N, T, HV, V), + rtol=rtol, + atol=atol) + + # Advance the golden committed state by the accepted prefix. + fused_recurrent_gated_delta_rule_update( + q=q[:, :accepted].contiguous(), + k=k[:, :accepted].contiguous(), + v=v[:, :accepted].contiguous(), + g=g[:, :accepted].contiguous(), + beta=beta[:, :accepted].contiguous(), + scale=scale, + initial_state_source=golden_state, + initial_state_indices=idx, + use_qk_l2norm_in_kernel=True, + disable_state_update=False, + disable_output_calculation=True, + ) + + # Emulate update_mamba_states PNAT / buffer bookkeeping. + acc = torch.full((N, ), accepted, device=dev, dtype=torch.int32) + wrote = (pnat + T) > HISTORY + pnat = torch.where(wrote, acc, pnat + acc) + cache_buf_idx = torch.where(wrote, 1 - cache_buf_idx, cache_buf_idx) + + +@skip_no_cuda +def test_replay_protocol_nowrite_only_fp32(): + """Large history => never checkpoints; pure nowrite accumulation.""" + _replay_protocol(torch.float32, + rtol=2e-4, + atol=2e-4, + accepted_pattern=[2, 3, 1, 2], + T=4, + HISTORY=16) + + +@skip_no_cuda +def test_replay_protocol_with_checkpoints_fp32(): + """Small history forces buffer flips / checkpoint writes.""" + _replay_protocol(torch.float32, + rtol=2e-4, + atol=2e-4, + accepted_pattern=[3, 3, 2, 3, 1], + T=4, + HISTORY=8) + + +@skip_no_cuda +def test_replay_protocol_bf16_loose(): + """bf16 checkpoint/history: same protocol with bf16-appropriate tolerance.""" + _replay_protocol(torch.bfloat16, + rtol=3e-2, + atol=3e-2, + accepted_pattern=[3, 2, 3, 2], + T=4, + HISTORY=8) + + +@skip_no_cuda +def test_replay_cuda_graph_capturable(): + """The replay kernel captures into a CUDA graph and replays deterministically.""" + from tensorrt_llm._torch.modules.fla.replay_gated_delta_rule_update import ( + replay_gated_delta_rule_update, ) + + torch.manual_seed(2) + dev = "cuda" + N, T, H, HV, K, V = 2, 4, 4, 8, 128, 128 + HISTORY = 16 + dtype = torch.float32 + scale = K**-0.5 + + q, k, v, a, b = _rand_step(N, T, H, HV, K, V, dev, dtype) + A_log = torch.empty(HV, device=dev).uniform_(1.0, 16.0).log() + dt_bias = torch.randn(HV, device=dev) * 0.1 + g = _gating(A_log, a, dt_bias, N, T, HV) + beta = b.sigmoid() + + state_pool = torch.randn(N, HV, V, K, device=dev, dtype=dtype) * 0.1 + idx = torch.arange(N, device=dev, dtype=torch.int32) + old_k = torch.zeros(N, 2, HISTORY, H, K, device=dev, dtype=dtype) + old_v = torch.zeros(N, 2, HISTORY, HV, V, device=dev, dtype=dtype) + old_g = torch.zeros(N, 2, HISTORY, HV, device=dev, dtype=torch.float32) + old_beta = torch.zeros(N, 2, HISTORY, HV, device=dev, dtype=torch.float32) + cache_buf_idx = torch.zeros(N, device=dev, dtype=torch.int32) + pnat = torch.zeros(N, device=dev, dtype=torch.int32) + out = torch.zeros(N, T, HV, V, device=dev, dtype=dtype) + + def run(): + replay_gated_delta_rule_update( + q=q, + k=k, + v=v, + g=g, + beta=beta, + ssm_states=state_pool, + old_k=old_k, + old_v=old_v, + old_g=old_g, + old_beta=old_beta, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=pnat, + state_batch_indices=idx, + replay_step_width=T, + replay_history_size=HISTORY, + scale=scale, + use_qk_l2norm_in_kernel=True, + output=out, + ) + + # Warmup on a side stream (required before capture). + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(s): + run() + torch.cuda.current_stream().wait_stream(s) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + + # Reset inputs/state and replay: PNAT==0 so state stays; output determinstic. + state_pool_ref = state_pool.clone() + graph.replay() + torch.cuda.synchronize() + out_first = out.clone() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(out, out_first, rtol=0, atol=0) + # Idempotent for PNAT==0: checkpoint unchanged. + torch.testing.assert_close(state_pool, state_pool_ref, rtol=1e-4, atol=1e-4)