---
id: pr-tensorrt-llm-12731
repo: NVIDIA/TensorRT-LLM
pr: 12731
title: '[None][feat] Optimize mamba SSD prefill and extend flashinfer dispatch'
author: Wanli-Jiang
date: '2026-05-12'
url: https://github.com/NVIDIA/TensorRT-LLM/pull/12731
source_category: upstream-code
architectures:
- sm100
tags:
- attention_kv
- compiler_runtime
- cute dsl
- moe
- moe_routing
techniques: []
hardware_features: []
kernel_types:
- attention
- compiler
- moe
languages:
- cuda-cpp
- cute-dsl
captured_at: '2026-05-16'
status: merged
merge_sha: da321a21f28f16efa11776a0837f75e5876026d9
inclusion_reason: Preserve page/KV/layout and plan-run contracts; profile memory traffic separately from math utilization.
changed_paths:
- cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu
- tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py
- tensorrt_llm/_torch/models/modeling_nemotron_h.py
- tensorrt_llm/_torch/models/modeling_speculative.py
- tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py
- tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py
- tensorrt_llm/_torch/modules/mamba/ssd_combined.py
- tensorrt_llm/_torch/visual_gen/jit_kernels/flash_attention/cute/cute_dsl_utils.py
- tensorrt_llm/llmapi/reasoning_parser.py
- tests/integration/test_lists/waives.txt
- tests/unittest/_torch/models/test_nemotron_h_puzzle.py
- tests/unittest/_torch/modules/mamba/test_fuse_elementwise_ops.py
- tests/unittest/_torch/modules/mamba/test_ssd_combined_flashinfer.py
artifact_dir: evidence/pull-bundles/tensorrt-llm/gh-12731
head_ref_oid: 3e9933e023fd738612eca54739e5bdb403f30527
---
# [None][feat] Optimize mamba SSD prefill and extend flashinfer dispatch

- Repository: `NVIDIA/TensorRT-LLM`
- PR: [12731](https://github.com/NVIDIA/TensorRT-LLM/pull/12731)
- Decision: `Preserve page/KV/layout and plan-run contracts; profile memory traffic separately from math utilization.`
- Reason: Preserve page/KV/layout and plan-run contracts; profile memory traffic separately from math utilization.

## Changed Paths

- `cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu`
- `tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py`
- `tensorrt_llm/_torch/models/modeling_nemotron_h.py`
- `tensorrt_llm/_torch/models/modeling_speculative.py`
- `tensorrt_llm/_torch/modules/mamba/fuse_elementwise_ops.py`
- `tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py`
- `tensorrt_llm/_torch/modules/mamba/ssd_combined.py`
- `tensorrt_llm/_torch/visual_gen/jit_kernels/flash_attention/cute/cute_dsl_utils.py`
- `tensorrt_llm/llmapi/reasoning_parser.py`
- `tests/integration/test_lists/waives.txt`
- `tests/unittest/_torch/models/test_nemotron_h_puzzle.py`
- `tests/unittest/_torch/modules/mamba/test_fuse_elementwise_ops.py`
- `tests/unittest/_torch/modules/mamba/test_ssd_combined_flashinfer.py`

## Evidence Bundle

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