From b43bd9ef98e972c2eb95d5780ba92b48728abacd Mon Sep 17 00:00:00 2001 From: jiant <107457950+JadoTu@users.noreply.github.com> Date: Thu, 23 Jul 2026 01:59:28 +0000 Subject: [PATCH 1/2] fintune the replay kernel under low batch size Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com> --- .../_torch/modules/fla/cached_replay.py | 49 +++++++++++++++---- 1 file changed, 40 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fla/cached_replay.py b/tensorrt_llm/_torch/modules/fla/cached_replay.py index 1eba282eaba4..3c926e8adb24 100644 --- a/tensorrt_llm/_torch/modules/fla/cached_replay.py +++ b/tensorrt_llm/_torch/modules/fla/cached_replay.py @@ -13,6 +13,9 @@ from tensorrt_llm._utils import get_sm_version _SMALL_GRID_HEAD_TILES = 512 +_BV16_MAX_HEAD_TILES = 64 +_BV32_MAX_HEAD_TILES = 128 +_RATIO2_FINE_MAPPING_MAX_HEAD_TILES = 256 _EIGHT_WARP_COMMIT_HEAD_TILES = 1024 _PIPELINED_COMMIT_HEAD_TILES = 2048 _TWO_STAGE_REPLAY_HEAD_TILES = 4096 @@ -21,6 +24,33 @@ CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE = 16 +def _default_cached_replay_block_v( + use_tuned_bf16_mapping: bool, + head_tiles: int, + value_dim: int, +) -> int: + """Select the replay value tile without changing its cache layout.""" + if use_tuned_bf16_mapping: + if head_tiles <= _BV16_MAX_HEAD_TILES: + return 16 + if head_tiles <= _BV32_MAX_HEAD_TILES: + return 32 + if head_tiles <= _SMALL_GRID_HEAD_TILES: + return 64 + return triton.next_power_of_2(value_dim) + + +def _supports_fine_grained_replay_tiling( + num_key_heads: int, + num_value_heads: int, + head_tiles: int, +) -> bool: + """Return whether the measured fine-grained value tiling applies.""" + return num_value_heads == 4 * num_key_heads or ( + num_value_heads == 2 * num_key_heads and head_tiles <= _RATIO2_FINE_MAPPING_MAX_HEAD_TILES + ) + + @triton.jit def _gdc_wait_with_memory_clobber(): tl.inline_asm_elementwise( @@ -669,25 +699,26 @@ def fused_recurrent_gated_delta_rule_cached_replay_update( v = v.contiguous() packed_qkv = q BK = triton.next_power_of_2(K) - # GB200 dispatch for the production Qwen3.5 MTP per-CTA shape. Balanced - # DEP and TEP runs at the same global batch have the same N * HV head-tile - # count, so use that workload measure instead of topology-specific H/HV - # values or the per-rank batch alone. - use_tuned_bf16_mapping = ( + # The complete workload mapping is retained for the 4:1 shape; the + # 2:1 fine tiling is enabled only through the measured low-batch range. + use_production_bf16_shape = ( T == 4 and history_size <= 16 - and HV == 4 * H and K == 128 and V == 128 and ssm_states.dtype == torch.bfloat16 ) head_tiles = N * HV - use_small_grid_mapping = use_tuned_bf16_mapping and head_tiles <= _SMALL_GRID_HEAD_TILES + use_tuned_bf16_mapping = use_production_bf16_shape and HV == 4 * H + use_fine_grained_mapping = use_production_bf16_shape and _supports_fine_grained_replay_tiling( + H, HV, head_tiles + ) + use_small_grid_mapping = use_fine_grained_mapping and head_tiles <= _SMALL_GRID_HEAD_TILES if block_v is None: - block_v = 64 if use_small_grid_mapping else triton.next_power_of_2(V) + block_v = _default_cached_replay_block_v(use_fine_grained_mapping, head_tiles, V) if num_warps is None: num_warps = ( - 2 if use_tuned_bf16_mapping and (use_small_grid_mapping or launch_with_pdl) else 4 + 2 if use_fine_grained_mapping and (use_small_grid_mapping or launch_with_pdl) else 4 ) BV = block_v use_large_workload_mapping = ( From f954542571d11e3cafe4fa419cd47929276a10c6 Mon Sep 17 00:00:00 2001 From: jiant <107457950+JadoTu@users.noreply.github.com> Date: Fri, 24 Jul 2026 07:21:33 +0000 Subject: [PATCH 2/2] optimizer runtime, remove the tensor copy ops, replay is now defaultly on Signed-off-by: jiant <107457950+JadoTu@users.noreply.github.com> --- .../_torch/modules/fla/cached_replay.py | 16 +- .../_torch/modules/mamba/gdn_mixer.py | 4 +- .../_torch/modules/mamba/mamba2_metadata.py | 13 +- tensorrt_llm/_torch/pyexecutor/_util.py | 4 +- .../_torch/pyexecutor/mamba_cache_manager.py | 187 +++++++++++++----- tensorrt_llm/_torch/utils.py | 7 +- .../executor/test_mamba_cache_manager.py | 53 +++++ .../mamba/test_gdn_kernel_optimizations.py | 4 +- .../mamba/test_gdn_replay_recurrent.py | 7 +- 9 files changed, 229 insertions(+), 66 deletions(-) diff --git a/tensorrt_llm/_torch/modules/fla/cached_replay.py b/tensorrt_llm/_torch/modules/fla/cached_replay.py index 3c926e8adb24..0811491a8503 100644 --- a/tensorrt_llm/_torch/modules/fla/cached_replay.py +++ b/tensorrt_llm/_torch/modules/fla/cached_replay.py @@ -81,6 +81,8 @@ def _cached_replay_kernel( pnat, scale, pool_stride_slot, + g_stride_row, + beta_stride_row, A_log, dt_bias, T: tl.constexpr, @@ -222,12 +224,12 @@ def _cached_replay_kernel( other=0, ) b_g = tl.load( - g + (i_n * T + o_t) * HV + i_hv, + g + (i_n * T + o_t) * g_stride_row + i_hv, mask=mask_t, other=0.0, ).to(tl.float32) b_beta = tl.load( - beta + (i_n * T + o_t) * HV + i_hv, + beta + (i_n * T + o_t) * beta_stride_row + i_hv, mask=mask_t, other=0.0, ).to(tl.float32) @@ -634,6 +636,8 @@ def commit_gdn_cached_replay_history_layers( "q", "k", "v", + "g", + "beta", "packed_qkv", "ssm_states", "old_u", @@ -685,6 +689,12 @@ def fused_recurrent_gated_delta_rule_cached_replay_update( HV, V = v.shape[2], v.shape[3] assert q.shape == k.shape assert v.shape[:2] == (N, T) + assert g.shape == (N, T, HV) + assert beta.shape == (N, T, HV) + assert g.stride(2) == 1, "g head dimension must be contiguous" + assert beta.stride(2) == 1, "beta head dimension must be contiguous" + assert g.stride(0) == T * g.stride(1), "g token rows must have a uniform stride" + assert beta.stride(0) == T * beta.stride(1), "beta token rows must have a uniform stride" use_packed_qkv = packed_qkv is not None if use_packed_qkv: qkv_width = 2 * H * K + HV * V @@ -804,6 +814,8 @@ def launch( pnat=prev_num_accepted_tokens, scale=scale, pool_stride_slot=s_h0_0, + g_stride_row=g.stride(1), + beta_stride_row=beta.stride(1), A_log=A_log, dt_bias=dt_bias, T=T, diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 3d29e6b05cfc..0fccbbf2896d 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -181,12 +181,12 @@ def __init__( replay_enabled = is_gdn_replay_enabled() if replay_enabled: logger.info_once( - "Configured GDN MTP replay implementation: cached", + "GDN MTP replay is enabled; set TRTLLM_USE_GDN_REPLAY=0 to disable it", key="gdn_mtp_replay_cached", ) else: logger.info_once( - "GDN MTP replay disabled; set TRTLLM_USE_GDN_REPLAY=1 to enable it", + "GDN MTP replay is disabled; set TRTLLM_USE_GDN_REPLAY=1 to enable it", key="gdn_mtp_replay_disabled", ) diff --git a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py index 5b97e32a6ee3..faa73250b774 100644 --- a/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py +++ b/tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py @@ -289,9 +289,20 @@ def _prepare_replay_work_items(self, kv_cache_manager, batch_size: int, return num_decodes = batch_size - num_contexts self.replay_num_decodes = num_decodes - self.replay_n_writes.zero_() if num_decodes == 0: return + if getattr(kv_cache_manager, 'use_gdn_cached_replay_all_layer_commit', + False): + from tensorrt_llm._torch.modules.fla.cached_replay import \ + CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE + + # The fused small-batch GDN kernel commits its checkpoint in-layer + # and indexes cache metadata directly. Work items are only consumed + # by the partitioned replay + all-layer commit path. + if num_decodes < CACHED_REPLAY_PARTITION_MIN_BATCH_SIZE: + return + + self.replay_n_writes.zero_() if not hasattr(kv_cache_manager, 'get_replay_state_update_metadata'): raise RuntimeError( "Replay state update is enabled, but the KV cache manager " diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 0149ba7ae03a..86722ad66d68 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -2021,10 +2021,8 @@ def _create_kv_cache_manager( "MTP path") use_replay = False - # Replay is opt-in because its end-to-end benefit is workload-dependent. + # Replay is enabled by default for eligible GDN MTP workloads. if not is_gdn_replay_enabled(): - logger.info("GDN replay kernel is disabled; set " - "TRTLLM_USE_GDN_REPLAY=1 to enable it") use_replay = False logger.info("GDN replay state update: " + ("ENABLED" if use_replay else "DISABLED")) diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 16fe22d1c777..0b2cda183726 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -1150,8 +1150,12 @@ def _promote_mamba_state_kernel( src_ptr, dst_ptr, src_idx_ptr, - acc_ptr, + accepted_position_source_ptr, + num_accepted_tokens_ptr, blk_ptr, + replay_pnat_ptr, + replay_cache_buf_idx_ptr, + dummy_request_mask_ptr, num_gens, count, src_s_layer, @@ -1159,6 +1163,10 @@ def _promote_mamba_state_kernel( src_s_step, dst_s_layer, dst_s_block, + POSITION_SOURCE_IS_TOKEN_COUNT: tl.constexpr, + REPLAY_STEP_WIDTH: tl.constexpr, + REPLAY_HISTORY_SIZE: tl.constexpr, + UPDATE_REPLAY_BOOKKEEPING: tl.constexpr, BLOCK: tl.constexpr, ): # One program copies a BLOCK-sized tile of the contiguous inner state for a @@ -1168,8 +1176,31 @@ def _promote_mamba_state_kernel( layer = (pair // num_gens).to(tl.int64) g = pair % num_gens row = tl.load(src_idx_ptr + g).to(tl.int64) - acc = tl.load(acc_ptr + g).to(tl.int64) + accepted_position = tl.load(accepted_position_source_ptr + g) + if POSITION_SOURCE_IS_TOKEN_COUNT: + accepted_position -= 1 + acc = accepted_position.to(tl.int64) blk = tl.load(blk_ptr + g).to(tl.int64) + + # The replay metadata update is batch-scoped, so layer 0's first copy tile + # owns it. This folds PNAT and active-buffer maintenance into the existing + # accepted-convolution-state promotion launch. + if UPDATE_REPLAY_BOOKKEEPING: + if (layer == 0) & (tile == 0): + num_accepted_tokens = tl.load(num_accepted_tokens_ptr + g) + is_dummy_request = tl.load(dummy_request_mask_ptr + g) + pnat = tl.load(replay_pnat_ptr + blk) + cache_buf_idx = tl.load(replay_cache_buf_idx_ptr + blk) + wrote_checkpoint = pnat + REPLAY_STEP_WIDTH > REPLAY_HISTORY_SIZE + next_pnat = tl.where(wrote_checkpoint, num_accepted_tokens, + pnat + num_accepted_tokens) + next_cache_buf_idx = tl.where(wrote_checkpoint, 1 - cache_buf_idx, + cache_buf_idx) + tl.store(replay_pnat_ptr + blk, next_pnat, mask=~is_dummy_request) + tl.store(replay_cache_buf_idx_ptr + blk, + next_cache_buf_idx, + mask=~is_dummy_request) + # int64 throughout: per-layer strides are O(1e8), so layer*stride overflows # int32 and would corrupt addresses / fault. src_base = layer * src_s_layer + row * src_s_row + acc * src_s_step @@ -1180,12 +1211,20 @@ def _promote_mamba_state_kernel( tl.store(dst_ptr + dst_base + offs, v, mask=mask) -def _promote_mamba_state_triton(dst: torch.Tensor, - intermediate: torch.Tensor, - src_state_indices: torch.Tensor, - accepted_draft: torch.Tensor, - dst_state_indices: torch.Tensor, - BLOCK: int = 2048) -> None: +def _promote_mamba_state_triton( + dst: torch.Tensor, + intermediate: torch.Tensor, + src_state_indices: torch.Tensor, + accepted_position_source: torch.Tensor, + dst_state_indices: torch.Tensor, + num_accepted_tokens: Optional[torch.Tensor] = None, + position_source_is_token_count: bool = False, + replay_pnat: Optional[torch.Tensor] = None, + replay_cache_buf_idx: Optional[torch.Tensor] = None, + dummy_request_mask: Optional[torch.Tensor] = None, + replay_step_width: int = 0, + replay_history_size: int = 0, + BLOCK: int = 2048) -> None: """Scatter each generation request's accepted draft-step recurrent state from the per-request intermediate buffer into the unified C++ pool view. @@ -1198,11 +1237,28 @@ def _promote_mamba_state_triton(dst: torch.Tensor, (dtype-view mutation) and Inductor codegen (``XBLOCK`` on the uint8 pool). intermediate: ``[num_layers, max_batch_size, T, *state_shape]`` dense. - src_state_indices, accepted_draft, dst_state_indices: ``[num_gens]`` int. + src_state_indices, accepted_position_source, dst_state_indices: + ``[num_gens]`` int. + num_accepted_tokens: ``[num_gens]`` accepted-token counts. Required + when replay bookkeeping is enabled. + position_source_is_token_count: Whether ``accepted_position_source`` + contains accepted-token counts instead of zero-based intermediate + state positions. When true, the kernel subtracts one while + computing the source position. + replay_pnat, replay_cache_buf_idx, dummy_request_mask: Optional replay + metadata. When supplied together, the kernel advances PNAT and the + active history buffer while promoting the convolution state. For each gen ``g``:: - dst[:, dst_state_indices[g]] = intermediate[:, src_state_indices[g], accepted_draft[g]] + accepted_position = ( + accepted_position_source[g] - 1 + if position_source_is_token_count + else accepted_position_source[g] + ) + dst[:, dst_state_indices[g]] = intermediate[ + :, src_state_indices[g], accepted_position + ] Pure gather->scatter copy; bandwidth-bound (~85% of HBM peak), one launch. """ @@ -1210,6 +1266,25 @@ def _promote_mamba_state_triton(dst: torch.Tensor, num_gens = src_state_indices.shape[0] if num_gens == 0: return + update_replay_bookkeeping = replay_pnat is not None + if update_replay_bookkeeping: + assert num_accepted_tokens is not None + assert replay_cache_buf_idx is not None + assert dummy_request_mask is not None + assert replay_step_width > 0 + assert replay_history_size >= replay_step_width + assert replay_pnat.dtype == torch.int32 + assert replay_cache_buf_idx.dtype == torch.int32 + assert dummy_request_mask.dtype == torch.bool + assert dummy_request_mask.numel() >= num_gens + else: + # Compile-time-disabled pointer arguments still require tensors. + num_accepted_tokens = accepted_position_source + assert replay_cache_buf_idx is None + assert dummy_request_mask is None + replay_pnat = dst_state_indices + replay_cache_buf_idx = dst_state_indices + dummy_request_mask = dst_state_indices count = 1 for s in dst.shape[2:]: count *= s @@ -1223,8 +1298,12 @@ def _promote_mamba_state_triton(dst: torch.Tensor, intermediate, dst, src_state_indices, - accepted_draft, + accepted_position_source, + num_accepted_tokens, dst_state_indices, + replay_pnat, + replay_cache_buf_idx, + dummy_request_mask, num_gens, count, intermediate.stride(0), @@ -1232,6 +1311,10 @@ def _promote_mamba_state_triton(dst: torch.Tensor, intermediate.stride(2), dst.stride(0), dst.stride(1), + POSITION_SOURCE_IS_TOKEN_COUNT=position_source_is_token_count, + REPLAY_STEP_WIDTH=replay_step_width, + REPLAY_HISTORY_SIZE=replay_history_size, + UPDATE_REPLAY_BOOKKEEPING=update_replay_bookkeeping, BLOCK=BLOCK, ) @@ -1811,13 +1894,17 @@ def update_mamba_states( batch_size = attn_metadata.num_seqs num_contexts = attn_metadata.num_contexts num_gens = batch_size - num_contexts - num_accepted_draft_tokens = ( - num_accepted_tokens[num_contexts:num_contexts + num_gens] - 1).to( - torch.int32) - # Dynamic tree passes tree-node leaf positions; linear MTP uses depth. - accepted_positions = (accepted_leaf_positions.to(torch.int32) - if accepted_leaf_positions is not None else - num_accepted_draft_tokens) + accepted_tokens = num_accepted_tokens[num_contexts:num_contexts + + num_gens] + # Dynamic tree supplies explicit tree-node leaf positions. Linear MTP + # passes accepted-token counts and lets the promotion kernel subtract + # one, avoiding a separate GPU subtraction/cast. + if accepted_leaf_positions is None: + accepted_position_source = accepted_tokens + position_source_is_token_count = True + else: + accepted_position_source = accepted_leaf_positions.to(torch.int32) + position_source_is_token_count = False # Match the API of MambaCacheManager.update_mamba_states: callers # may pass per-request state slot indices explicitly (e.g. MTP via # attn_metadata.mamba_metadata.state_indices). Fall back to this @@ -1840,45 +1927,41 @@ def update_mamba_states( # in one launch before PNAT and the active history buffer change. self._commit_gdn_cached_replay_history_layers( attn_metadata, num_gens) - # 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].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, - self.intermediate_ssm_states, - src_state_indices, accepted_positions, - state_indices_d) + _promote_mamba_state_triton( + self.all_ssm_states, + self.intermediate_ssm_states, + src_state_indices, + accepted_position_source, + state_indices_d, + position_source_is_token_count=position_source_is_token_count) # Conv: both paths save all intermediate conv windows, carry over the - # accepted one. - _promote_mamba_state_triton(self.all_conv_states, - self.intermediate_conv_states, - src_state_indices, accepted_positions, - state_indices_d) + # accepted one. The replay path also advances PNAT and the active cache + # buffer in this same launch. + replay_dummy_request_mask = None + if self._use_replay_state_update: + replay_dummy_request_mask = self._dummy_request_mask[ + num_contexts:num_contexts + num_gens] + _promote_mamba_state_triton( + self.all_conv_states, + self.intermediate_conv_states, + src_state_indices, + accepted_position_source, + state_indices_d, + num_accepted_tokens=accepted_tokens, + position_source_is_token_count=position_source_is_token_count, + replay_pnat=(self.prev_num_accepted_tokens + if self._use_replay_state_update else None), + replay_cache_buf_idx=(self.cache_buf_idx + if self._use_replay_state_update else None), + dummy_request_mask=replay_dummy_request_mask, + replay_step_width=(self.replay_step_width + if self._use_replay_state_update else 0), + replay_history_size=(self.replay_history_size + if self._use_replay_state_update else 0)) @torch.inference_mode() def _refresh_dummy_request_mask(self, is_dummy: List[bool]) -> None: diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index ba0453fed790..a7b69e22b9bc 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -1,3 +1,6 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import contextlib import functools import os @@ -39,8 +42,8 @@ def is_gdn_replay_enabled() -> bool: - """Return whether GDN replay was explicitly enabled.""" - return os.environ.get("TRTLLM_USE_GDN_REPLAY", "0") == "1" + """Return whether GDN replay is enabled (default: enabled).""" + return os.environ.get("TRTLLM_USE_GDN_REPLAY", "1") == "1" # IMPORTANT: Keep the same order of activation functions in this enum and the enum in diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 0291ef3e74dd..757a0d208f50 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -20,6 +20,7 @@ CppMambaHybridCacheManager, PythonMambaCacheManager, _get_mamba_hybrid_pool_size, + _promote_mamba_state_triton, ) from tensorrt_llm._torch.pyexecutor.resource_manager import CacheTypeCpp, DataType, KVCacheManager from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests @@ -225,6 +226,58 @@ def test_replay_update_mamba_states_skips_dummy_slots(): assert mgr.mamba_cache.cache_buf_idx[dummy_slot].item() == 1 +@skip_no_cuda +@pytest.mark.parametrize( + ( + "position_source_values", + "position_source_is_token_count", + "expected_positions", + ), + [ + ([3, 2], True, [2, 1]), + ([1, 3], False, [1, 3]), + ], +) +def test_promote_mamba_state_fuses_replay_bookkeeping( + position_source_values, + position_source_is_token_count, + expected_positions, +): + intermediate = torch.arange(2 * 2 * 4 * 3, dtype=torch.float32, device="cuda").reshape( + 2, 2, 4, 3 + ) + dst = torch.zeros(2, 5, 3, dtype=torch.float32, device="cuda") + src_state_indices = torch.tensor([0, 1], dtype=torch.int32, device="cuda") + accepted_position_source = torch.tensor( + position_source_values, dtype=torch.int32, device="cuda" + ) + num_accepted_tokens = torch.tensor([3, 2], dtype=torch.int32, device="cuda") + dst_state_indices = torch.tensor([1, 3], dtype=torch.int32, device="cuda") + replay_pnat = torch.tensor([0, 13, 0, 7, 0], dtype=torch.int32, device="cuda") + replay_cache_buf_idx = torch.tensor([0, 1, 0, 1, 0], dtype=torch.int32, device="cuda") + dummy_request_mask = torch.tensor([False, True], dtype=torch.bool, device="cuda") + + _promote_mamba_state_triton( + dst, + intermediate, + src_state_indices, + accepted_position_source, + dst_state_indices, + num_accepted_tokens=num_accepted_tokens, + position_source_is_token_count=position_source_is_token_count, + replay_pnat=replay_pnat, + replay_cache_buf_idx=replay_cache_buf_idx, + dummy_request_mask=dummy_request_mask, + replay_step_width=4, + replay_history_size=MIN_REPLAY_HISTORY_SIZE, + ) + + torch.testing.assert_close(dst[:, 1], intermediate[:, 0, expected_positions[0]]) + torch.testing.assert_close(dst[:, 3], intermediate[:, 1, expected_positions[1]]) + assert replay_pnat.tolist() == [0, 3, 0, 7, 0] + assert replay_cache_buf_idx.tolist() == [0, 0, 0, 1, 0] + + @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/test_gdn_kernel_optimizations.py b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py index a80fc0d083ec..9c54cfc1cc9b 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py @@ -30,9 +30,9 @@ @pytest.mark.parametrize( "value,expected", - [(None, False), ("", False), ("0", False), ("false", False), ("typo", False), ("1", True)], + [(None, True), ("", False), ("0", False), ("false", False), ("typo", False), ("1", True)], ) -def test_gdn_replay_env_requires_one(monkeypatch, value, expected): +def test_gdn_replay_env_defaults_on_and_requires_one_when_set(monkeypatch, value, expected): from tensorrt_llm._torch.utils import is_gdn_replay_enabled if value is None: diff --git a/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py b/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py index dab201710b67..8d89e12b9f7d 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_replay_recurrent.py @@ -146,8 +146,11 @@ def test_gdn_replay_vs_legacy_and_ref(H, HV, K, V, T, HIST, iters, pool_dtype, f assert packed_qkv.is_contiguous() assert not q.is_contiguous() and not k.is_contiguous() and not v.is_contiguous() if fused_gating: - a_raw = torch.randn(N, T, HV, device=device, dtype=dtype) * 0.5 - b_raw = torch.randn(N, T, HV, device=device, dtype=dtype) + packed_ba = torch.randn(N * T, 2 * HV, device=device, dtype=dtype) + packed_ba[:, HV:].mul_(0.5) + b_raw = packed_ba[:, :HV].view(N, T, HV) + a_raw = packed_ba[:, HV:].view(N, T, HV) + assert not a_raw.is_contiguous() and not b_raw.is_contiguous() # Host reference of the in-kernel gating (fp32, same inputs). g = -(A_log_t.exp() * torch.nn.functional.softplus(a_raw.float() + dt_bias_t)) beta = torch.sigmoid(b_raw.float())