[Misc] Add generic concat + conv + split vs N separate grouped convs benchmark

[Costa-SM]

Summary

Hi folks.

Originally this PR was for a GDN-specific harness but, following the discussion below, it covers a gap in the existing benchmarking tooling by adding a conv-level case to benchmarks/modules/benchmark_conv.py. It also fixes two existing bugs in the script.

Motivation

While trying to improve the GDN layer's performance, I found a measurement gap:

• benchmarks.ops.run --op chunk_gated_delta_rule only performs kernel-level measurements.
• benchmark_training_throughput.py dilutes smaller improvements since it does full-scale training.

So, one of the possible improvements I was working on, which consisted of a q/k/v short-conv fusion, couldn't be captured with the existing tooling.

New Benchmark

A model-agnostic comparison of two ways to run N independent depthwise short convs over [B, T, D] inputs:

• separate_* — N separate causal_conv1d calls.
• fused_* — cat the N inputs along channels, one causal_conv1d, then split.

It is parameterized purely by the number of streams N, so it applies to any module doing grouped short convs. The fused path is validated against the separate path (assert_close) before timing, so a wrong kernel can't masquerade as a faster one.

• Providers: separate_fwd / fused_fwd / separate_fwdbwd / fused_fwdbwd.
• Shapes: N ∈ {2, 3, 4}, T ∈ {2048, 4096}, D ∈ {512, 1024} (per-stream channels).
• Timing: triton.testing.do_bench with quantiles [0.5, 0.2, 0.8].

Bug fixes

The PR also fixes 2 bugs in the existing benchmark tooling:

• causal_conv1d_fwdbwd matched the causal_conv1d_fwd branch first (startswith), so the fwd+bwd providers never actually ran a backward pass. Dispatch now matches the provider exactly
• causal_conv1d returns (output, state), so the path now indexes [0] before calling .backward() on it

Test plan

Run python -m benchmarks.modules.benchmark_conv and verify both tables print; the fused_* providers pass the assert_close check against separate_* before timing.

Breaking changes

None. This only touches the benchmarking script (new benchmark + bug fixes to the existing one), so there are no model or kernel behavior changes.

[gemini-code-assist]

Code Review

This pull request introduces a new layer-level benchmark script, benchmark_gated_deltanet.py, for the GatedDeltaNet model to measure performance and memory usage across various shapes and providers. The review feedback highlights opportunities to improve memory management, specifically by clearing the CUDA cache and forcing garbage collection during the warmup phase to prevent potential Out-Of-Memory (OOM) errors, and doing the same before memory profiling to ensure accurate and consistent peak memory measurements.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[Costa-SM]

I had to update the harness because small performance improvements were not being captured very well.

Main changes:

• added repeat timing (--repeat, default 3)
• increased default timing windows to --warmup-ms 50 --rep-ms 300
• added a local untimed warmup before each measured repeat
• shuffled provider/shape measurement order per repeat with recorded --order-seed
• added per-repeat medians, min/max, and spread to JSON
• added spread and noisy markers in baseline comparisons
• added group geomeans for all, layer, layer_train, op, and conv

[yzhangcs]

Thanks a lot for putting this together, and for the thorough write-up — the methodology notes (warmup-order sensitivity, per-call working-set memory) are genuinely nice.

That said, I think most of what this adds is already covered, and we'd actually prefer to move away from per-model benchmark scripts rather than add more — we've recently been cleaning up the standalone ones. So I'd lean toward not landing a GDN-specific harness. Mapping your providers onto what's already there:

Full layer / model (layer_*) — benchmarks/benchmark_training_throughput.py times the whole model end-to-end (tokens/s + peak memory), and --profile runs torch.profiler to print per-kernel CUDA/CPU breakdowns (plus an optional chrome trace). That covers the op-vs-conv-vs-projection attribution you're after, measured on the real model:

# end-to-end throughput + per-kernel breakdown
python -m benchmarks.benchmark_training_throughput --name gated_deltanet --profile

# narrow to a single layer, and export a trace for chrome://tracing / perfetto
python -m benchmarks.benchmark_training_throughput --name gated_deltanet --num_hidden_layers 1 --profile --profile-trace trace.json

Op level (op_*) — chunk_gated_delta_rule is already registered in benchmarks/ops/registry.py, so the unified runner covers it, with built-in baseline comparison (auto git worktree), JSON output, custom shapes, etc.:

python -m benchmarks.ops.run --op chunk_gated_delta_rule --base main

Would that work for you?

[Costa-SM]

yupyup makes sense. This PR was tailored to GDN because it's what I'd been working on.

I used the tooling you suggested, but it seems to not be able to capture some of the improvements I'm working on. For example, I've been working on this (spoiler alert) and using the benchmark script in this PR, we see that it results in a ~1.17x speedup for the GDN layer. That improvement gets diluted to noise in the training throughput benchmark (~1.003x), and the op benchmark doesn't capture it since it doesn't change anything in the kernel.

In the last few commits on this branch I nuked the GDN benchmark script and pushed some changes to benchmark_conv.py that introduce a generic concat + conv + split vs N separate grouped convs comparison benchmark that covers that gap. It also fixes two bugs in the existing benchmark function:

• The causal_conv1d_fwdbwd would match with causal_conv1d_fwd first (so it never actually ran a backward pass)
• We had to index the (output, state) tuple that causal_conv1d returns before calling backward on it

Does that look interesting to you ? If so, I'll update the description too. Or I can open another PR.