---
id: pr-tensorrt-llm-13453
repo: NVIDIA/TensorRT-LLM
pr: 13453
title: '[None][feat] Use a replay method for state rollback in Mamba-2 speculative decoding'
author: hnover-nv
date: '2026-04-30'
url: https://github.com/NVIDIA/TensorRT-LLM/pull/13453
source_category: upstream-code
architectures:
- sm100
tags:
- attention kernel
- benchmark kernel
- triton kernel
techniques: []
hardware_features: []
kernel_types:
- attention
languages:
- cuda-cpp
- triton
captured_at: '2026-05-16'
status: merged
merge_sha: 5221c7bf1c4cd6c4bd97621a0cb0d8a1ad219f37
inclusion_reason: Search-matched CUDA kernel optimization PR; materialize diff/source before synthesis.
changed_paths:
- tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py
- tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py
- tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py
- tensorrt_llm/_torch/modules/mamba/selective_state_update.py
- tensorrt_llm/_torch/pyexecutor/_util.py
- tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
- tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py
- tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py
- tests/unittest/disaggregated/test_mamba_transfer.py
artifact_dir: evidence/pull-bundles/tensorrt-llm/gh-13453
head_ref_oid: 5b8d147d53fa96f20243c1ebb5a867482dfb0e0d
---
# [None][feat] Use a replay method for state rollback in Mamba-2 speculative decoding

- Repository: `NVIDIA/TensorRT-LLM`
- PR: [13453](https://github.com/NVIDIA/TensorRT-LLM/pull/13453)
- Decision: `Search-matched CUDA kernel optimization PR; materialize diff/source before synthesis.`
- Reason: Search-matched CUDA kernel optimization PR; materialize diff/source before synthesis.

## Changed Paths

- `tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py`
- `tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py`
- `tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py`
- `tensorrt_llm/_torch/modules/mamba/selective_state_update.py`
- `tensorrt_llm/_torch/pyexecutor/_util.py`
- `tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py`
- `tests/unittest/_torch/modules/mamba/benchmark_replay_selective_state_update.py`
- `tests/unittest/_torch/modules/mamba/test_replay_selective_state_update.py`
- `tests/unittest/disaggregated/test_mamba_transfer.py`

## Evidence Bundle

`evidence/pull-bundles/tensorrt-llm/gh-13453`
