[BUG] Blackwell MLA perf for split-k
Nobody has claimed this yet.
- Dominant language
- C++
- Stars
- 10.5k
- Forks
- 2.1k
- Avg merge
- 3d 11h
- Merged PRs (30d)
- 7
Description
Describe the bug
I want to pass variable split_k via Arguments.ptr_split_kv on-device. I need to also set max value of split k Arguments.split_kv. The issue is that I need to keep max value of split k static (like 16 or so), but as soon as I set it higher than 1 perf degrades a lot (I pass only 1 to ptr_split_kv). This makes ptr_split_kv on-device not practically usable.
Another data point. Perf degrades 2x if I just split by 4x. This is really not expected as it's not doing 2x work (yes, there is some overhead, but it should be much smaller).
$ /77_blackwell_mla_2sm_fp16 --b=512 --k=2048 --page=128 --verbose --split_kv=1
###### B 512 MLA H 128 D_rope 64 D_latent 512 Q 1 K 2048 Gen None Split 1 Gen None #SM 148
[--] 128x128 fp16 persistent : 1073.23 TFLOPS/s 4.13682 TB/s
t=326.475 us, smem=210944b
[--] 128x128 fp16 individual : 1071.13 TFLOPS/s 4.12872 TB/s
t=327.115 us, smem=210944b
$ /77_blackwell_mla_2sm_fp16 --b=512 --k=2048 --page=128 --verbose --split_kv=4
###### B 512 MLA H 128 D_rope 64 D_latent 512 Q 1 K 2048 Gen None Split 4 Gen None #SM 148
[--] 128x128 fp16 persistent : 547.292 TFLOPS/s 2.10956 TB/s
t=640.213 us, smem=210944b
[--] 128x128 fp16 individual : 508.058 TFLOPS/s 1.95833 TB/s
t=689.653 us, smem=210944b
Steps/Code to reproduce bug
See description for repro.
Expected behavior
I expect perf not to degrade significantly when max split k is increased.
Environment details (please complete the following information):
B200
Additional context
Add any other context about the problem here.
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the 77_blackwell_mla_2sm_fp16 entry point and reproduce the two commands using --split_kv=1 and --split_kv=4 on B200. Trace how Arguments.ptr_split_kv and Arguments.split_kv are handled, then measure whether raising the maximum split-k value leaves the split_kv=1 case near its baseline without the reported performance regression.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100