---
id: pr-tensorrt-llm-12537
repo: NVIDIA/TensorRT-LLM
pr: 12537
title: '[None][feat] Add Mamba2 MTP SSM cache CUDA kernel for tree-based speculative decoding'
author: JadoTu
date: '2026-04-01'
url: https://github.com/NVIDIA/TensorRT-LLM/pull/12537
source_category: upstream-code
architectures:
- sm100
tags:
- benchmark kernel
- benchmark_test
- cuda kernel
- memory_primitives
techniques:
- cache-policy
hardware_features: []
kernel_types:
- memory
languages:
- cuda-cpp
captured_at: '2026-05-16'
status: merged
merge_sha: 1b6a2bcd2dc29914d9c06522ccedd4075fc69ed1
inclusion_reason: 'Transfer primitive policy carefully: coalescing, shared-memory staging, partial tiles, and determinism are part of the contract.'
changed_paths:
- cpp/tensorrt_llm/CMakeLists.txt
- cpp/tensorrt_llm/kernels/CMakeLists.txt
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/CMakeLists.txt
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCache.cu
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCache.h
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheKernel.cuh
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheVec16.cu
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheVec4.cu
- cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheVec8.cu
- cpp/tensorrt_llm/thop/CMakeLists.txt
- cpp/tensorrt_llm/thop/mamba2MTPSSMCacheOp.cpp
- tensorrt_llm/_torch/modules/mamba/selective_state_update.py
- tests/unittest/_torch/modules/mamba/test_mamba2_mtp_ssm_cache.py
- tests/unittest/_torch/thop/parallel/test_custom_ops.py
artifact_dir: evidence/pull-bundles/tensorrt-llm/gh-12537
head_ref_oid: 33c9959b7832e931b6eafa954ff7a49a22d14daf
---
# [None][feat] Add Mamba2 MTP SSM cache CUDA kernel for tree-based speculative decoding

- Repository: `NVIDIA/TensorRT-LLM`
- PR: [12537](https://github.com/NVIDIA/TensorRT-LLM/pull/12537)
- Decision: `Transfer primitive policy carefully: coalescing, shared-memory staging, partial tiles, and determinism are part of the contract.`
- Reason: Transfer primitive policy carefully: coalescing, shared-memory staging, partial tiles, and determinism are part of the contract.

## Changed Paths

- `cpp/tensorrt_llm/CMakeLists.txt`
- `cpp/tensorrt_llm/kernels/CMakeLists.txt`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/CMakeLists.txt`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCache.cu`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCache.h`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheKernel.cuh`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheVec16.cu`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheVec4.cu`
- `cpp/tensorrt_llm/kernels/mamba2MTPSSMCache/mamba2MTPSSMCacheVec8.cu`
- `cpp/tensorrt_llm/thop/CMakeLists.txt`
- `cpp/tensorrt_llm/thop/mamba2MTPSSMCacheOp.cpp`
- `tensorrt_llm/_torch/modules/mamba/selective_state_update.py`
- `tests/unittest/_torch/modules/mamba/test_mamba2_mtp_ssm_cache.py`
- `tests/unittest/_torch/thop/parallel/test_custom_ops.py`

## Evidence Bundle

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