[Kernel] Add support for Flashinfer Mamba SSU algorithm selection

[amitz-nv]

Purpose

Allow choosing the Flashinfer Mamba SSU algorithm.
The new command line argument is --mamba-ssu-algorithm.

On some use-cases of Nemotron 3 Nano NVFP4 (like ISL/OSL=1k/8k), using "horizontal" instead of "auto" (which eventually chooses "vertical") improves throughput.

Test Plan

On Nemotron 3 Nano NVFP4 with & without --mamba-ssu-algorithm horizontal:

• Throughput benchmark with aiperf (3 runs to account for variance), on a single H100 GPU
• Accuracy test with GSM8K, on TP=2

Test Result

Accuracy test

vLLM command:

MODEL=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4
vllm serve "$MODEL" \
  --served-model-name "$MODEL" \
  --trust-remote-code \
  --tensor-parallel-size 2 \
  --max-model-len 8192 \
  --mamba-backend flashinfer \
  --host 0.0.0.0 \
  --port 8000

To run with "horizontal", add:

--mamba-ssu-algorithm horizontal

lm_eval command:

MODEL=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4

lm_eval \
  --model local-completions \
  --model_args "base_url=http://127.0.0.1:8000/v1/completions,model=$MODEL,tokenized_requests=False,tokenizer_backend=None,num_concurrent=512,timeout=120,max_retries=5,max_length=8192" \
  --tasks gsm8k \
  --num_fewshot 5 \
  --batch_size auto \
  --seed 42 \
  --output_path "$RESULT_DIR"

Results:

Without --mamba-ssu-algorithm horizontal:

local-completions ({'base_url': 'http://127.0.0.1:37524/v1/completions', 'model': 'nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4', 'tokenized_requests': False, 'tokenizer_backend': None, 'num_concurrent': 512, 'timeout': 120, 'max_retries': 5, 'max_length': 8192}), gen_kwargs: ({}), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4390|±  |0.0137|
|     |       |strict-match    |     5|exact_match|↑  |0.8423|±  |0.0100|

With --mamba-ssu-algorithm horizontal:

local-completions ({'base_url': 'http://127.0.0.1:37525/v1/completions', 'model': 'nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4', 'tokenized_requests': False, 'tokenizer_backend': None, 'num_concurrent': 512, 'timeout': 120, 'max_retries': 5, 'max_length': 8192}), gen_kwargs: ({}), limit: None, num_fewshot: 5, batch_size: auto
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.4261|±  |0.0136|
|     |       |strict-match    |     5|exact_match|↑  |0.8438|±  |0.0100|

Throughput benchmark

vLLM command without horizontal:

MODEL=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4

VLLM_MOE_ROUTING_SIMULATION_STRATEGY=uniform_random \
vllm serve "$MODEL" \
  --served-model-name "$MODEL" \
  --max-num-seqs 512 \
  --trust-remote-code \
  --max-model-len 52000 \
  --max-num-batched-tokens 32768 \
  --enable-prefix-caching \
  --async-scheduling \
  --quantization modelopt_fp4 \
  --mamba-backend flashinfer \
  --mamba-cache-mode align \
  --mamba-ssm-cache-dtype float16 \
  --enable-mamba-cache-stochastic-rounding \
  --mamba-cache-philox-rounds 5 \
  --host 0.0.0.0 \
  --port 8000

To run with "horizontal", add:

--mamba-ssu-algorithm horizontal

AIPerf command:

MODEL=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4

aiperf profile \
  --model "$MODEL" \
  --tokenizer "$MODEL" \
  --url http://127.0.0.1:8000 \
  --endpoint-type chat \
  --ui-type None \
  --streaming \
  --concurrency 512 \
  --request-count 1024 \
  --warmup-request-count 16 \
  --isl 1024 \
  --isl-stddev 0 \
  --osl 8192 \
  --osl-stddev 0 \
  --random-seed 42 \
  --use-server-token-count \
  --tokenizer-trust-remote-code \
  --extra-inputs temperature:0 \
  --extra-inputs ignore_eos:true \
  --artifact-dir "$RESULT_DIR"

Results (on 1xH100):

Configuration	Run 1	Run 2	Run 3	Mean output TPS/GPU
FlashInfer default	13,061.091	13,078.941	13,087.962	13,075.998
FlashInfer horizontal	13,914.585	13,790.473	13,902.849	13,869.302
Improvement	+6.535%	+5.440%	+6.226%	+6.067%

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[amirkl94]

Can we have the default value as None, and if we're using FI as mamba backend, then we can resolve it as "auto"?

[amirkl94]

Can we add a log_once with the picked algorithm somwhere?

[amitz-nv]

Note that when "auto" is used, vLLM doesn't know what flashinfer eventually chose.

[amirkl94]

I think all of these tests don't really belong here, maybe in tests/kernels/mamba/test_ssu_dispatch.py ? And some of them look redundant to me.

[amirkl94]

LGTM

[mgoin]

LGTM