[Benchmark] Benchmark probe issue, memory probe and estimation may be inaccurate for some kernels

[lowdy1]

Issue with Current Probe Design

In the current design, estimate_kernel_peak_memory probes a kernel’s peak memory usage and returns peak_bytes:

peak_bytes = estimate_kernel_peak_memory(probe_fn=_probe)
kernel_bpt = peak_bytes // num_tokens  # if needed

This derives a bytes-per-token (BPT) estimate by assuming memory scales linearly with the number of tokens.

---

Limitation

This assumption only holds for kernels whose memory footprint scales linearly with num_tokens.

However, for some kernels Scaling may be non-linear (e.g., quadratic)

As a result, the computed kernel_bpt can be inaccurate or misleading for such kernels.

---

Summary

The current BPT estimation is:

• ✅ Reasonable for linear-scaling kernels
• ❌ Unreliable for kernels with non-linear or structured memory behavior

[Tcc0403]

What about a new argument scaling_method for the estimation function. Default to "linear".

[shivam2199]

@Tcc0403 Took a look at this. The linear assumption actually appears in two places, not just estimate_kernel_peak_memory:

1. estimate_kernel_peak_memory returns raw peak_bytes; callers divide by num_tokens themselves.
2. compute_seq_len_sweep_config inverts via max_tokens = usable_bytes / kernel_bytes_per_token — only valid for linear scaling. benchmark_sparse_multi_token_attention.py:334 already works around this by dividing by probe_L * probe_L, but the downstream sweep math still treats the result as linear bytes-per-token.

One direction worth considering: a new helper that internalizes both the probe and the inversion, leaving the existing functions alone.

def compute_seq_len_sweep_config_with_probe(
    model_cfg: ModelConfig,
    probe_fn: Callable[[], torch.Tensor],
    probe_seq_len: int,
    probe_batch_size: int = 1,
    scaling_method: Literal["linear", "quadratic_in_L"] = "linear",
    memory_utilization: float = 0.4,
    max_seq_len: Optional[int] = None,
    max_batch_size: int = 32,
) -> SeqLenSweepConfig:
    ...

For quadratic_in_L, the inversion would be L_max = sqrt(usable_bytes / c_per_L²) instead of max_tokens / batch_size.

The reason for a new helper rather than threading scaling_method through the existing function: 16+ benchmark scripts already call estimate_kernel_peak_memory, and changing its signature would mean a wide diff that could conflict with the in-flight benchmark refactors (#1199, #1180). Linear callers wouldn't need to change; only quadratic-scaling benchmarks would opt in. As a side benefit, the sparse-attention benchmark could switch to this and the manual // (probe_L * probe_L) step would go away.

A couple of things I'm unsure about:

• Whether "linear" and "quadratic_in_L" cover the cases you have in mind, or whether "quadratic_in_BL" (full attention) should also be a separate mode.
• For validation, I was thinking a synthetic O(L²) probe plus the existing sparse-attention benchmark on an A10G to check the predicted max-L doesn't OOM at the boundary — does that match what you'd want to see?

Happy to adjust the shape based on whichever direction you'd prefer, or open a draft PR if this looks roughly right.

[Tcc0403]

Thanks for taking a look. Yes I forgot this issue when merging memory estimation helper function into seqlen sweep config function. Your approach seems right, we should be good with "linear" and "quadratic" as compute_seqlen_sweep_config already states it is scaled with seqlen. I think our formula still holds in the BxL case since we batch size should be fixed. Feel free to open a PR and tag me.