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