[TRTLLM-5835][feat] Optimized Mamba2Mixer prefill

[tomeras91]

Description

Currently on main, the mamba2 block forward pass has some dynamic memory allocations and host-to-device copies, greatly hurting its performance. This PR improves performance by minimizing these memory operations:

1. pre-allocate the state_indices tensors on device instead of moving it to device during forward pass
2. introduce MambaMetadata class holding cu_seqlens and seq_idx needed for varlen batched prefill of the SSM op. Compute them from attn_metadata at the start of the model forward pass instead of doing it inside the Mamba2Mixer block of each mamba layer. This also means we create these tensors once and not multiple times in each layer.
3. Remove the redundant for-loop on request type (prefill/decode) in the Mamba2Mixer forward pass. Replace it with 2 if statements, checking if we have prefills / decodes in the current batch.

These changes removed many of the GPU bubbles present in the mamba forward block, as seen in these profiles. They lead to 40% latency reduction.

Benchmarks

The modifications in this PR improved Nemotron-H max throughput by ~15% and min prefill latency by ~40% without sacrificing quality:

MMLU results

Version	MMLU Score
Main	69.98
PR	69.98

Performance benchmarks

Benchmark setting:

• Used genai-perf
• Done on a single RTX6000
• ISL / OSL = 223 / 140 (similar to ShareGPT statistics)
• nax_batch_size = 64
• Max throughput setting: --request-count 500 --concurrency 500 --warmup-request-count 10
• Min latency setting: --request-count 10 --concurrency 1 --warmup-request-count 1

 
max throughput results:

Metric	TRTLLM main	PR	Change [%]
Request Latency [ms]	30,614.38	26,482.98	-13.5%
Time To First Token [ms]	25,107.02	21,563.04	-14.1%
Inter Token Latency [ms]	40.91	36.97	-9.6%
Output Token Throughput [tokens/sec]	1,235.88	1,400.71	+13.3%
Output Token Throughput Per User [tokens/sec/user]	24.84	27.68	+11.4%
Request Throughput [req/sec]	9.14	10.43	+14.1%

 
min latency results:

Metric	TRTLLM main	PR	Change [%]
Request Latency [ms]	3,283.65	2,362.60	-28.0%
Time To First Token [ms]	82.55	50.54	-38.8%
Inter Token Latency [ms]	24.58	17.77	-27.7%
Output Token Throughput [tokens/sec]	39.89	55.44	+39.0%
Output Token Throughput Per User [tokens/sec/user]	40.69	56.28	+38.3%
Request Throughput [req/sec]	0.30	0.42	+40.0%

[Copilot]

Pull Request Overview

This PR optimizes the Mamba2Mixer prefill performance by reducing dynamic memory allocations and host-to-device copies. Key changes include:

• Pre-allocating state indices on the correct device in the resource manager.
• Introducing a new Mamba2Metadata class to pre-compute and store metadata for varlen batched prefill.
• Refactoring the Mamba2Mixer forward pass to separate prefill and decode logic and eliminate a redundant loop.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
tensorrt_llm/_torch/pyexecutor/resource_manager.py	Assigns the correct device for state indices to minimize host-to-device transfers.
tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py	Refactors forward pass to split prefill and decode logic using new metadata and indices.
tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py	Introduces a metadata class to compute and hold cu_seqlens and sequence indices.
tensorrt_llm/_torch/models/modeling_nemotron_h.py	Integrates mamba metadata into layer forward passes for optimized processing.

Comments suppressed due to low confidence (2)

tensorrt_llm/_torch/pyexecutor/resource_manager.py:598

• Consider adding an inline comment explaining that the device is explicitly set using self.ssm_states.device to ensure correct GPU allocation, and verify that self.ssm_states is initialized prior to this call.

self.state_indices = torch.as_tensor(state_indices,

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py:178

• [nitpick] Consider adding a brief comment to clarify that state_indices is split into prefill and decode subsets based on the computed batch_split_size, which will aid future maintainers in understanding the code logic.

state_indices_p, state_indices_d = torch.split(state_indices, batch_split_size)

[tomeras91]

/bot run

[tomeras91]

/bot run

[tensorrt-cicd]

PR_Github #8670 [ run ] triggered by Bot

[tensorrt-cicd]

PR_Github #8670 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #6288 completed with status: 'SUCCESS'