Question about mamba chunk_scan kernel performance

[monellz]

Hello, I tried to reproduce the performance of chunk_scan kernel reported in the PipeThreader paper. But the actual result does not match the reported performance.

tilelang kernel: examples/linear_attention/example_mamba_chunk_scan.py
commit: 569b012
nvcc version: 12.4
torch version: 2.6.0
cmd: python example_mamba_chunk_scan.py --batch 64 --seq_len 8192 --tune
hardware: H100 PCIe

The output is

Best latency: 48.831199645996094
Best TFlops: 28.145725369921212
Best config: {'block_M': 64, 'block_N': 32, 'block_K': 64, 'block_Dstate': 128, 'num_stages': 4}

But as the paper said (Table 3) , ChunkScan for bs=64 seq=8k should cost 6.981ms.

Is the example file the one used in paper experiements?

[LeiWang1999]

@monellz Hi, thanks for your reports.

1. From my perspective, many kernels perform sub-optimally under CUDA 12.4; our tests were conducted on CUDA 12.8.
2. We tested on the H100 SXM, which is approximately 0.5× faster than the PCIe version.

@chengyupku Is this the kernel implementation we used for benchmarking?

[chengyupku]

I think this is the kernel we used for benchmarking, and using the commit from that time in my private repo, I can reproduce the results in the paper. It’s possible that some later commits caused the performance degradation. I’ll look into it.

[monellz]

I think this is the kernel we used for benchmarking, and using the commit from that time in my private repo, I can reproduce the results in the paper. It’s possible that some later commits caused the performance degradation. I’ll look into it.

Can you provide a commit hash that can reproduce the results? (The private repo is https://github.com/chengyupku/tilelang?)

I dont know whether osdi25 has artifacts evaluation process this year. But it will help a lot if tilelang can provide some benchmark scripts/commit hash for the paper experiments version.

Thanks a lot for your quick reply and this great work.

[chengyupku]

I think this is the kernel we used for benchmarking, and using the commit from that time in my private repo, I can reproduce the results in the paper. It’s possible that some later commits caused the performance degradation. I’ll look into it.

Can you provide a commit hash that can reproduce the results? (The private repo is https://github.com/chengyupku/tilelang?)

I dont know whether osdi25 has artifacts evaluation process this year. But it will help a lot if tilelang can provide some benchmark scripts/commit hash for the paper experiments version.

Thanks a lot for your quick reply and this great work.

You can find the corresponding kernel implementation here. It is a pre-release version of TileLang before it was open-sourced.

[chengyupku]

This performance regression was caused by two factors:

• the ldmatrix instruction was not inlined, and
• the producer’s complex index computation consumed too many registers, leading to register spilling.
  This issue has been fixed in [Refactor] Use forceinline in ldmatrix and update mamba scan kernel #1104.