NVIDIA / NVIDIA/cutlass

[BUG] Blackwell MLA perf for split-k

Open
#2,222 4 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

? - Needs Triage bug inactive-30d inactive-90d
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

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.