[TRTLLM-5835][feat] Optimized Mamba2Mixer prefill#5128

Merged

Naveassaf merged 71 commits intoNVIDIA:mainfrom

tomeras91:optimize-mamba-prefill

Jun 16, 2025

Commits on May 11, 2025

format and style
tomeras91
committed
Add oad skipping support for ssm decode kernel
tomeras91
committed
import adaptations to state-spaces/mamba based mamba2 SSM prefill kernels for better numerical stability and support for initial states + varlen batching (AKA continuous batching)
tomeras91
committed

Commits on May 14, 2025

add correctness unittests for the state-spaces/mamba based mamba2 SSM prefill and decode kernels (similar to the tests in tests/unittest/_torch/thop/test_selective_scan_op.py)
tomeras91
committed
align ref state shape with kernel output (insteadl of the other way around)
tomeras91
committed

Commits on May 15, 2025

Merge branch 'main' into mamba2-kernel-updates
tomeras91
committed

Commits on May 20, 2025

Commits on May 22, 2025

Merge branch 'main' into mamba2-kernel-updates
tomeras91
authored

Commits on May 25, 2025

Commits on May 26, 2025

Commits on May 27, 2025

Commits on May 29, 2025

move cu_seqlens and seq_ids to resource_manager.prepare_resources() so done not during forward pass
tomeras91
committed

Commits on Jun 1, 2025

Merge branch 'main' into paged-mamba-cache
tomeras91
committed

Commits on Jun 4, 2025

Merge branch 'main' into optimize-mamba-prefill
tomeras91
committed

Commits on Jun 5, 2025

Commits on Jun 8, 2025

Commits on Jun 11, 2025

Commits on Jun 12, 2025

Merge branch 'main' into optimize-mamba-prefill
tomeras91
committed