[Benchmark][Mamba2] Modernise SSD benchmark suite: parametrize, mamba baseline, model-scale params

[stelladuyx]

Summary

• Replace @Fixture decorator with @pytest.mark.parametrize + explicit typed -> None signatures across all 5 files, matching style of bench_gqa_decode_paged.py
• Add mamba_ssm.ops.triton baseline comparison to bench_ssd_chunk_scan_fwd, bench_ssd_chunk_state_fwd, bench_ssd_state_passing_fwd (falls back to torch-ref when mamba-ssm is not installed); da_cumsum and ssd_decode have no mamba-ssm equivalent
• Add _to_mamba_inputs() layout converter for chunk_scan (TileOPs chunked → mamba flat-sequence layout)
• Add Mamba2 model-scale params to bench_ssd_chunk_scan_fwd and bench_ssd_decode: 130M/370M/780M/1.3B/2.7B with latency/serving/throughput/long-context cases
• Fix BenchmarkReport.record(): replace string literal first arg with op; normalise tag= to "torch-ref" / "mamba"

Test plan

• --collect-only → 55 tests collected, no errors
• All pre-commit hooks pass
• With mamba-ssm: tag="mamba" results appear; without: falls back to tag="torch-ref"

[gemini-code-assist]

Code Review

This pull request refactors several SSD (Mamba2) operator benchmarks by replacing custom fixtures with pytest.mark.parametrize and adding comprehensive model-scale parameters (ranging from 130M to 2.7B configurations). It also introduces baseline comparisons against mamba_ssm for the chunk scan, chunk state, and state passing operators. The review feedback identifies several critical issues where tensor layouts (permutations) and input mappings in the new benchmark code do not match the expectations of the mamba_ssm library, which would lead to incorrect baseline results.

I am having trouble creating individual review comments. Click here to see my feedback.

benchmarks/ops/bench_ssd_chunk_scan_fwd.py (19-31)

The _to_mamba_inputs layout converter has several critical issues that will lead to incorrect baseline results or runtime errors:

1. Incorrect dt_m and dA_m mapping: The dt input is ignored, and dA_cumsum is used for both dt_m and dA_m. In mamba_ssm, dt is the element-wise discretized decay, while dA_cumsum is the prefix sum. They are distinct.
2. Missing Permutations: mamba_ssm expects dt and dA_cumsum in (B, C, H, L) layout. In this benchmark, dt is (b, c, L, h) and dA_cumsum is (b, h, c, L). Both need permutations.
3. Incorrect State Layout: prev_states is (b, c, h, n, p). The permutation (0, 1, 2, 4, 3) changes it to (b, c, h, p, n), but mamba_ssm expects (B, C, H, N, P).

def _to_mamba_inputs(x, cb, dA_cumsum, C, prev_states, dt):
    b, c, L, h, p = x.shape
    n = C.shape[-1]

    x_m = x.reshape(b, c * L, h, p).contiguous()
    cb_m = cb.contiguous()
    # TileOPs dt: (b, c, L, h) -> mamba: (b, c, h, L)
    dt_m = dt.permute(0, 1, 3, 2).contiguous()
    # TileOPs dA_cumsum: (b, h, c, L) -> mamba: (b, c, h, L)
    dA_m = dA_cumsum.permute(0, 2, 1, 3).contiguous()
    C_m = C.reshape(b, c * L, h, n).contiguous()
    # TileOPs prev_states: (b, c, h, n, p) -> mamba: (b, c, h, n, p)
    states_m = prev_states.contiguous()

    return x_m, cb_m, dt_m, dA_m, C_m, states_m

benchmarks/ops/bench_ssd_chunk_state_fwd.py (86-93)

The mamba_ssm baseline comparison will produce incorrect results because the dt and dA_cumsum tensors are passed in the wrong layout. tileops uses (b, h, c, L) for these tensors, while mamba_ssm expects (B, C, H, L). A permutation is required.

        def mamba_fwd():
            return _chunk_state_fwd(
                Bmat.contiguous(),
                x.contiguous(),
                # TileOPs (b, h, c, L) -> mamba (b, c, h, L)
                dt.permute(0, 2, 1, 3).contiguous(),
                dA_cumsum.permute(0, 2, 1, 3).contiguous(),
                seq_idx=seq_idx,
            )

benchmarks/ops/bench_ssd_state_passing_fwd.py (68-73)

The dA_chunk_cumsum tensor needs to be permuted before being passed to the mamba_ssm baseline. tileops provides it as (b, h, c), but mamba_ssm expects (B, C, H).

        def mamba_fwd():
            return _state_passing_fwd(
                states.contiguous(),
                # TileOPs (b, h, c) -> mamba (b, c, h)
                dA_chunk_cumsum.permute(0, 2, 1).contiguous(),
                initial_states=initial_states.contiguous(),
            )

[Ibuki-wind]

Summary

Model-scale params, mamba_ssm baseline integration, and the parametrize migration are all good ideas. But there are correctness and hygiene issues that need fixing before merge.

Blocking

1. Dead tests.* imports in two benchmark files — bench_da_cumsum_fwd.py and bench_ssd_decode.py both add from tests.ops.* imports where every imported name is immediately shadowed (by local function definitions or the workloads.* import on the next line). These imports are 100% dead code, but they re-introduce the tests.* coupling that #787 eliminated. Remove them entirely.

2. Missing tensor layout permutations in mamba_ssm baselines — bench_ssd_chunk_state_fwd.py and bench_ssd_state_passing_fwd.py pass dt/dA_cumsum/dA_chunk_cumsum to mamba_ssm with only .contiguous(), no layout permutation. But the local reference functions in the same files explicitly permute(0, 2, 1, 3) these tensors from TileOPs (b,h,c,L) to computation (b,c,h,L) layout. If mamba_ssm expects the same computation layout, the baselines produce silently wrong benchmark numbers. The chunk_scan file has a proper _to_mamba_inputs converter — the other two files need equivalent treatment.

Non-blocking

3. Unused fixture imports — DaCumsumFwdFixture and SsdDecodeFixture are still imported but no longer used after the @Fixture → @pytest.mark.parametrize migration.

Action items

1. Delete the from tests.ops.* import blocks in bench_da_cumsum_fwd.py and bench_ssd_decode.py.
2. Clean up DaCumsumFwdFixture and SsdDecodeFixture from their respective workloads imports.
3. Add layout permutations (or a shared _to_mamba_inputs-style converter) to the mamba baselines in bench_ssd_chunk_state_fwd.py and bench_ssd_state_passing_fwd.py. Document the expected mamba_ssm layout in a comment, like chunk_scan does.
4. Verify the mamba baselines actually produce correct results by running them against the torch-ref baseline on at least one param set.