mamba2 : Flatten in/out projections to dispatch GEMM instead of GEMV

[pskrunner14]

Overview

With CUDA backend, I noticed that Nemotron 3 Nano (30B A3B NVFP4) model was decoding slower than alternate baseline inference engine with concurrency npl>=8. Profiling showed that it was launching mul_mat_vec_* instead of GEMM, and was dominating GPU time ~40% at npl=32.

This was due to mamba-base.cpp reshaping activations to {n_embd, n_seq_tokens, n_seqs} 3d shape before ssm_in/ssm_out. So with decode n_seq_tokens==1 the dispatcher sees ne11==1 and treats it as column-batched GEMV for what is a dense GEMM.

This PR adds a minimal fix to keep the projections flat so we optimally dispatch GEMM instead of GEMV.

Additional information

Closes #27464

Perf Impact

Measured the performance impact of this change with relevant Mamba2 models with CUDA backend on GB10 (sm121) and RTX PRO 6000 (sm120). The results are shown below as text-generation tokens/s. The numbers are on master 78ec4c3 with llama-batched-bench -npp 64 -ntg 128 -npl 1,8,32,256.

GB10 (sm121)

Nemotron 3 Nano 30B-A3B (NVFP4):

npl	current TG t/s	patched TG t/s	gain
1	67.5	67.7	0%
8	152.4	202.6	+33%
32	214.8	350.8	+63%
256	292.5	660.8	+126%

Other Mamba2-based models (npl=32):

model	current TG t/s	patched TG t/s	gain
Mamba2 2.7B Q8_0	98.4	246.7	+151%
Falcon-H1 7B BF16	80.0	199.0	+149%
Falcon-H1 7B Q4_K_M	187.7	273.1	+45%
Granite 4.0 H Tiny BF16	271.3	390.4	+44%
Granite 4.0 H Tiny Q4_K_M	456.4	551.2	+21%

RTX PRO6000 (sm120)

Nemotron 3 Nano 30B-A3B (NVFP4):

npl	current TG t/s	patched TG t/s	gain
1	333.2	332.0	0%
8	1035.5	1234.1	+19%
32	1542.0	2132.4	+38%
256	1959.8	3159.1	+61%

Other mamba2-based models (npl=32):

model	current TG t/s	patched TG t/s	gain
mamba2-2.7b BF16	630.9	1508.0	+139%
mamba2-2.7b Q4_K_M	953.0	1638.9	+72%
Falcon-H1-7B BF16	771.5	1188.2	+54%
Falcon-H1-7B Q8_0	963.3	1318.3	+37%
Falcon-H1-7B Q4_K_M	1067.9	1417.1	+33%
granite-4.0-h-tiny BF16	1427.0	2287.6	+60%
granite-4.0-h-tiny Q4_K_M	2192.4	2541.9	+16%

npl=1 is unchanged for all models (same GEMV either way).

Tests

Correctness: greedy decode output is numerically identical with master 07822bddf, perplexity identical, test-backend-ops and ctest pass. Also tested with CPU backend, no regressions in perf and passing all tests.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: Yes, AI was used for assisting with profiling and perf study with other models.

[pskrunner14]

@gaugarg-nv can you please review the PR? thanks!

[pskrunner14]

@CISC can you please review PR? Thanks

[CISC]

I think there was a reason for this, I can't recall though, @ggerganov?