From 2bd3f1e209d557fe9d765b6eb836394f20d828de Mon Sep 17 00:00:00 2001 From: Jin Li <59594262+liji-nv@users.noreply.github.com> Date: Wed, 22 Jul 2026 01:19:44 -0700 Subject: [PATCH] [None][perf] avoid GDN state reset host synchronization GDN prefill reset currently selects cache slots with a CUDA boolean mask before assigning zeros. The boolean-indexed selection has a data-dependent output shape, which forces the host to obtain the selected element count and introduces a stream synchronization for every GDN layer. Replace the two indexed assignments with a fixed-shape Triton launch. Each program reads the request's cache index and initialization flag on device, validates the cache slot, and clears both the SSM and convolution state rows in one kernel. Requests with initialized state or invalid sentinel indices remain untouched. Compute the state row offsets explicitly in int64. Recurrent-state views stride across the interleaved SSM and convolution pool, so multiplying an int32 cache index by either stride is not guaranteed to fit in int32 for a large pool. Add a CUDA unit test covering reset slots, preserved initialized slots, untouched unrelated slots, and the negative invalid-index sentinel. Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- .../_torch/modules/mamba/gdn_mixer.py | 74 +++++++++++++++++-- .../mamba/test_gdn_kernel_optimizations.py | 34 +++++++++ 2 files changed, 103 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py index 3d29e6b05cfc..f786417a4e46 100644 --- a/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py +++ b/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py @@ -89,6 +89,70 @@ def _extract_gdn_extra_attrs(layer_idx: str): return metadata, gdn_layer, extra_attrs.get("spec_metadata", None) +@triton.jit +def _reset_gdn_states_kernel( + ssm_states, + conv_states, + state_indices, + has_initial_states, + ssm_state_stride, + conv_state_stride, + NUM_CACHE_LINES: tl.constexpr, + SSM_STATE_SIZE: tl.constexpr, + CONV_STATE_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + request_idx = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(state_indices + request_idx).to(tl.int64) + needs_reset = ~tl.load(has_initial_states + request_idx).to(tl.int1) + valid_state = (state_idx >= 0) & (state_idx < NUM_CACHE_LINES) + ssm_row_offset = state_idx * ssm_state_stride.to(tl.int64) + conv_row_offset = state_idx * conv_state_stride.to(tl.int64) + + tl.store( + ssm_states + ssm_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < SSM_STATE_SIZE), + ) + tl.store( + conv_states + conv_row_offset + offsets, + 0.0, + mask=needs_reset & valid_state & (offsets < CONV_STATE_SIZE), + ) + + +def _reset_gdn_states( + ssm_states: torch.Tensor, + conv_states: torch.Tensor, + state_indices: torch.Tensor, + has_initial_states: torch.Tensor, +) -> None: + num_requests = state_indices.shape[0] + if num_requests == 0: + return + + ssm_state_size = ssm_states.numel() // ssm_states.shape[0] + conv_state_size = conv_states.numel() // conv_states.shape[0] + block_size = 256 + grid = ( + num_requests, + triton.cdiv(max(ssm_state_size, conv_state_size), block_size), + ) + _reset_gdn_states_kernel[grid]( + ssm_states, + conv_states, + state_indices, + has_initial_states, + ssm_states.stride(0), + conv_states.stride(0), + ssm_states.shape[0], + ssm_state_size, + conv_state_size, + block_size, + ) + + @torch.library.custom_op("trtllm::gdn_custom_op_inplace", mutates_args=("output",)) def gdn_custom_op_inplace( mixed_qkv: torch.Tensor, @@ -1038,11 +1102,11 @@ def forward_core( if num_prefills > 0: # PyExecutor guarantees prefill requests are placed before decode requests has_initial_states_p = has_initial_states[:num_prefills] - ssm_states[state_indices_p[~has_initial_states_p]] = torch.zeros( - (), dtype=ssm_states.dtype, device=ssm_states.device - ) - conv_states[state_indices_p[~has_initial_states_p]] = torch.zeros( - (), dtype=conv_states.dtype, device=conv_states.device + _reset_gdn_states( + ssm_states, + conv_states, + state_indices_p, + has_initial_states_p, ) is_target_verify = ( 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..33bc67aeea6a 100644 --- a/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py +++ b/tests/unittest/_torch/modules/mamba/test_gdn_kernel_optimizations.py @@ -481,6 +481,40 @@ def test_pack_gdn_decode_qkv( # ---- Tests for the GDN compile boundary and derived state ---- +@skip_no_cuda +def test_reset_gdn_states_preserves_initialized_and_invalid_slots(): + from tensorrt_llm._torch.modules.mamba.gdn_mixer import _reset_gdn_states + + num_slots = 4 + ssm_state_size = 6 + conv_state_size = 8 + state_pool = ( + torch.arange( + num_slots * (ssm_state_size + conv_state_size), + device="cuda", + dtype=torch.float32, + ) + .to(torch.bfloat16) + .reshape(num_slots, -1) + ) + state_pool_before = state_pool.clone() + ssm_states = state_pool[:, :ssm_state_size].view(num_slots, 2, 3) + conv_states = state_pool[:, ssm_state_size:].view(num_slots, 2, 4) + state_indices = torch.tensor([1, 3, -1], device="cuda", dtype=torch.int32) + has_initial_states = torch.tensor([False, True, False], device="cuda", dtype=torch.bool) + + _reset_gdn_states( + ssm_states, + conv_states, + state_indices, + has_initial_states, + ) + + torch.testing.assert_close(state_pool[1], torch.zeros_like(state_pool[1])) + for state_idx in (0, 2, 3): + torch.testing.assert_close(state_pool[state_idx], state_pool_before[state_idx]) + + def test_gdn_custom_op_forwards_split_inputs(monkeypatch): """The compile boundary consumes split inputs and mutates its output.""" import tensorrt_llm._torch.modules.mamba.gdn_mixer as gdn_mixer