SSM

7 ops, 36 workloads.

One table per op, one row per workload. Ratio is the fastest other implementation's device time divided by ours, so green is faster than it, plain is level with it, red is slower. Times are in ms. How these numbers are taken.

Mamba2Fwd

x: [Bsz, S, H, P]dt: [Bsz, S, H], f32A: [H], f32B, C: [Bsz, S, G, N]

• mamba2-2p7bx: [1, 2048, 80, 64]dt: [1, 2048, 80], f32A: [80], f32B, C: [1, 2048, 1, 128]dtype=bf16
• mamba2-1p3bx: [1, 8192, 64, 64]dt: [1, 8192, 64], f32A: [64], f32B, C: [1, 8192, 1, 128]dtype=f16

x: [Bsz, S, H, P]dt: [Bsz, S, H], f32A, dt_bias: [H], f32B, C: [Bsz, S, G, N]

• mamba2-2p7b-dt-biasx: [1, 2048, 80, 64]dt: [1, 2048, 80], f32A, dt_bias: [80], f32B, C: [1, 2048, 1, 128]dtype=bf16
• mamba2-1p3b-dt-biasx: [1, 8192, 64, 64]dt: [1, 8192, 64], f32A, dt_bias: [64], f32B, C: [1, 8192, 1, 128]dtype=f16

x: [Bsz, S, H, P]dt: [Bsz, S, H], f32A: [H], f32B, C: [Bsz, S, G, N]initial_states: [Bsz, H, P, N], f32

• mamba2-2p7b-initx: [1, 2048, 80, 64]dt: [1, 2048, 80], f32A: [80], f32B, C: [1, 2048, 1, 128]initial_states: [1, 80, 64, 128], f32dtype=bf16
• mamba2-1p3b-initx: [1, 8192, 64, 64]dt: [1, 8192, 64], f32A: [64], f32B, C: [1, 8192, 1, 128]initial_states: [1, 64, 64, 128], f32dtype=f16

x: [Bsz, S, H, P]dt: [Bsz, S, H], f32A, dt_bias: [H], f32B, C: [Bsz, S, G, N]initial_states: [Bsz, H, P, N], f32

• mamba2-2p7b-dt-bias-initx: [1, 2048, 80, 64]dt: [1, 2048, 80], f32A, dt_bias: [80], f32B, C: [1, 2048, 1, 128]initial_states: [1, 80, 64, 128], f32dtype=bf16
• mamba2-1p3b-dt-bias-initx: [1, 8192, 64, 64]dt: [1, 8192, 64], f32A, dt_bias: [64], f32B, C: [1, 8192, 1, 128]initial_states: [1, 64, 64, 128], f32dtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-2p7b	bf16	1.34×	0.1069	mamba torch-compile torch-ref	0.1433 0.6923 2.1725	76.2	14%	mem
mamba2-1p3b	f16	1.48×	0.2863	mamba torch-compile torch-ref	0.4245 2.0283 6.9406	91.2	16%	mem
mamba2-2p7b-dt-bias	bf16	1.35×	0.1068	mamba torch-compile torch-ref	0.1438 0.6930 2.1734	76.2	14%	mem
mamba2-1p3b-dt-bias	f16	1.48×	0.2861	mamba torch-compile torch-ref	0.4247 2.0229 6.9409	91.3	16%	mem
mamba2-2p7b-init	bf16	1.34×	0.1081	mamba torch-compile torch-ref	0.1448 0.6847 2.1694	75.3	14%	mem
mamba2-1p3b-init	f16	1.48×	0.2862	mamba torch-compile torch-ref	0.4248 2.0262 6.9403	91.2	16%	mem
mamba2-2p7b-dt-bias-init	bf16	1.35×	0.1078	mamba torch-compile torch-ref	0.1455 0.6854 2.1706	75.5	14%	mem
mamba2-1p3b-dt-bias-init	f16	1.48×	0.2867	mamba torch-compile torch-ref	0.4248 2.0227 6.9421	91.1	16%	mem

SSDChunkCumsumFwd

dt: [B, S, H], f32A: [H], f32

chunk_len=256dt_softplus=true

• mamba2-780mdt: [1, 4096, 48], f32A: [48], f32dtype=f16
• mamba2-1p3bdt: [8, 2048, 64], f32A: [64], f32dtype=bf16

dt: [B, S, H], f32A, dt_bias: [H], f32

chunk_len=256dt_softplus=true

• mamba2-780m-dt-biasdt: [1, 4096, 48], f32A, dt_bias: [48], f32dtype=f16
• mamba2-1p3b-dt-biasdt: [8, 2048, 64], f32A, dt_bias: [64], f32dtype=bf16
• mamba2-2p7bdt: [2, 32768, 80], f32A, dt_bias: [80], f32dtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-780m	f16	1.53×	0.00307	mamba torch-compile torch-ref	0.0047 0.00483 0.0727	0.448	14%	lat
mamba2-1p3b	bf16	1.48×	0.00602	mamba torch-compile torch-ref	0.00893 0.0114 0.0939	1.22	39%	lat
mamba2-780m-dt-bias	f16	1.55×	0.00307	mamba torch-compile torch-ref	0.00477 0.0048 0.0754	0.512	14%	lat
mamba2-1p3b-dt-bias	bf16	1.47×	0.00608	mamba torch-compile torch-ref	0.00893 0.0115 0.0989	1.38	38%	lat
mamba2-2p7b	f16	1.26×	0.0238	mamba torch-compile torch-ref	0.0301 0.0368 0.2465	1.76	49%	mem

SSDChunkCouplingFwd

C_mat, B_mat: [B, S, G, N]

chunk_len=256

• mamba2-780mC_mat, B_mat: [1, 4096, 1, 128]dtype=f16
• mamba2-2p7bC_mat, B_mat: [4, 2048, 1, 128]dtype=bf16
• mamba2-grouped-prefillC_mat, B_mat: [3, 3072, 8, 128]dtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-780m	f16	1.42×	0.0072	mamba torch	0.0102 0.0381	18.6	13%	lat
mamba2-2p7b	bf16	1.36×	0.0118	mamba torch	0.0161 0.0532	22.7	16%	lat
mamba2-grouped-prefill	f16	1.80×	0.0589	mamba torch	0.1060 0.4685	41	28%	mem

SSDChunkStateFwd

x: [B, NC * Q, H, P]Bmat: [B, NC * Q, G, N]dt, dA_cumsum: [B, H, NC, Q], f32

• mamba2-780mx: [1, 4096, 48, 64]Bmat: [1, 4096, 1, 128]dt, dA_cumsum: [1, 48, 16, 256], f32dtype=f16/f32, bf16/f32
• mamba2-2p7bx: [4, 2048, 80, 64]Bmat: [4, 2048, 1, 128]dt, dA_cumsum: [4, 80, 8, 256], f32dtype=bf16/f32

x: [B, NC * Q, H, P]Bmat: [B, NC * Q, G, N]dt, dA_cumsum: [B, H, NC, Q], f32seq_idx: [B, NC * Q], i32

• mamba2-780m-seq-idxx: [1, 4096, 48, 64]Bmat: [1, 4096, 1, 128]dt, dA_cumsum: [1, 48, 16, 256], f32seq_idx: [1, 4096], i32dtype=f16/f32, bf16/f32
• mamba2-1p3bx: [2, 32768, 64, 64]Bmat: [2, 32768, 1, 128]dt, dA_cumsum: [2, 64, 128, 256], f32seq_idx: [2, 32768], i32dtype=f16/f32

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-780m	f16/f32	1.04×	0.0238	mamba torch-compile torch-ref	0.0249 0.6339 8.1111	136	49%	mem
	bf16/f32	1.11×	0.0240	mamba torch-compile torch-ref	0.0266 0.6782 8.1141	135	49%	mem
mamba2-2p7b	bf16/f32	1.21×	0.0658	mamba torch-compile torch-ref	0.0798 2.4492 27.0212	164	59%	mem
mamba2-780m-seq-idx	f16/f32	1.20×	0.0289	mamba torch-compile torch-ref	0.0346 0.7518 8.1288	112	41%	mem
	bf16/f32	1.01×	0.0292	mamba torch-compile torch-ref	0.0293 0.7894 8.1290	111	40%	mem
mamba2-1p3b	f16/f32	1.39×	0.4532	mamba torch-compile torch-ref	0.6315 16.9932 172.5820	152	55%	mem

SSDStatePassingFwd

states: [B, NC, H, N]dA_chunk_cumsum: [B, H, NC], f32

• mamba2-1p3bstates: [1, 16, 64, 128]dA_chunk_cumsum: [1, 64, 16], f32dtype=f16, bf16
• mamba2-2p7bstates: [2, 128, 80, 128]dA_chunk_cumsum: [2, 80, 128], f32dtype=f16

states: [B, NC, H, N]dA_chunk_cumsum: [B, H, NC], f32initial_states: [B, H, N], f32

• mamba2-1p3b-initstates: [1, 16, 64, 128]dA_chunk_cumsum: [1, 64, 16], f32initial_states: [1, 64, 128], f32dtype=f16, bf16

states: [B, NC, H, N]dA_chunk_cumsum: [B, H, NC]initial_states: [B, H, N]

• mamba2-1p3b-flatstates: [1, 16, 64, 8192]dA_chunk_cumsum: [1, 64, 16]initial_states: [1, 64, 8192]dtype=f32

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-1p3b	f16	2.19×	0.00189	torch-compile mamba torch-ref	0.00413 0.00858 0.1214	0.139	10%	lat
	bf16	2.10×	0.00195	torch-compile mamba torch-ref	0.0041 0.00854 0.1233	0.135	9%	lat
mamba2-2p7b	f16	5.71×	0.0108	mamba torch-compile torch-ref	0.0618 0.0921 1.1655	0.487	33%	mem
mamba2-1p3b-init	f16	1.78×	0.00192	torch-compile mamba torch-ref	0.00342 0.00893 0.1216	0.137	10%	lat
	bf16	1.78×	0.00192	torch-compile mamba torch-ref	0.00342 0.00896 0.1232	0.137	10%	lat
mamba2-1p3b-flat	f32	0.95×	0.0216	torch-compile mamba torch-ref	0.0205 0.0218 0.1255	0.776	73%	mem

SSDChunkScanFwd

x: [B, NC * Q, H, P]cb: [B, NC, G, Q, Q]dA_cumsum: [B, H, NC, Q], f32C: [B, NC * Q, G, N]prev_states: [B, NC, H, P, N], f32dt: [B, H, NC, Q]

• mamba2-780mx: [1, 4096, 48, 64]cb: [1, 16, 1, 256, 256]dA_cumsum: [1, 48, 16, 256], f32C: [1, 4096, 1, 128]prev_states: [1, 16, 48, 64, 128], f32dt: [1, 48, 16, 256]dtype=f16, bf16
• mamba2-2p7bx: [4, 2048, 80, 64]cb: [4, 8, 1, 256, 256]dA_cumsum: [4, 80, 8, 256], f32C: [4, 2048, 1, 128]prev_states: [4, 8, 80, 64, 128], f32dt: [4, 80, 8, 256]dtype=bf16
• mamba2-1p3bx: [2, 32768, 64, 64]cb: [2, 128, 1, 256, 256]dA_cumsum: [2, 64, 128, 256], f32C: [2, 32768, 1, 128]prev_states: [2, 128, 64, 64, 128], f32dt: [2, 64, 128, 256]dtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-780m	f16	1.97×	0.0722	mamba torch-compile torch-ref	0.1424 0.5081 1.9714	89.3	32%	mem
	bf16	1.90×	0.0760	mamba torch-compile torch-ref	0.1444 0.5094 1.9712	84.8	31%	mem
mamba2-2p7b	bf16	1.90×	0.2369	mamba torch-compile torch-ref	0.4510 1.6433 6.5487	90.7	32%	mem
mamba2-1p3b	f16	2.01×	1.4526	mamba torch-compile torch-ref	2.9233 9.9604 40.2415	94.6	34%	mem

SSDRecurrentFwd

A: [H, P, N], f32dt: [B, H, P], f32x: [B, H, P]B_in, C_in: [B, G, N]state: [B, H, P, N], f32

• mamba2-1p3bA: [64, 64, 128], f32dt: [1, 64, 64], f32x: [1, 64, 64]B_in, C_in: [1, 1, 128]state: [1, 64, 64, 128], f32dtype=f16, bf16
• mamba2-2p7bA: [80, 64, 128], f32dt: [8, 80, 64], f32x: [8, 80, 64]B_in, C_in: [8, 1, 128]state: [8, 80, 64, 128], f32dtype=f16
• mamba2-780mA: [48, 64, 128], f32dt: [32, 48, 64], f32x: [32, 48, 64]B_in, C_in: [32, 1, 128]state: [32, 48, 64, 128], f32dtype=f16

Workload	dtype	Ratio	Device time	Alternatives		Throughput	SOL	Bound
		alt / ours	ms	name	ms	TFLOP/s	of ceiling	by
mamba2-1p3b	f16	1.90×	0.00397	mamba torch-compile flashinfer torch-ref	0.00752 0.00912 0.0136 0.0300	1.06	35%	lat
	bf16	1.93×	0.0039	mamba torch-compile flashinfer torch-ref	0.00752 0.00899 0.0136 0.0302	1.07	36%	lat
mamba2-2p7b	f16	1.37×	0.0162	mamba flashinfer torch-compile torch-ref	0.0222 0.0285 0.0305 0.1119	2.59	62%	mem
mamba2-780m	f16	1.24×	0.0361	mamba flashinfer torch-compile torch-ref	0.0449 0.0467 0.0697 0.2405	2.79	63%	mem