A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Attention] Tune CUTLASS MLA num_splits (#26846)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 16, 2025 at 09:36 UTC
314fa8abbf9d4f6dc89eba1d8fdf80e6cd4432ed
1 file changed
+21
-16
csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp
+21
-16
@@ -125,32 +125,37 @@ public:
125
}
126
127
static void set_split_kv (KernelArguments& args) {
128
- // printf("set_split_kv start");
128
if (args.split_kv >= 1) return;
129
auto [H, K, D, B] = args.problem_shape;
131
- // std::cout << H << " " << K << " " << D << " " << B << "\n";
130
int sm_count = args.hw_info.sm_count;
133
- // printf(" sm_count = %d\n", sm_count);
134
- int max_splits = ceil_div(K, 128);
135
- max_splits = min(16, max_splits);
136
-
137
- // TODO: This avoids a hang when the batch size larger than 1 and
138
- // there is more than 1 kv_splits.
139
- // Discuss with NVIDIA how this can be fixed.
140
- if (B > 1) {
141
- max_splits = min(1, max_splits);
131
+ float seq_length_k = static_cast<float>(K) / 1024.0f;
132
+ int max_splits = 1;
133
+
134
+ if (B <= 4 && seq_length_k >= 16) {
135
+ max_splits = 16;
136
+ }
137
+ else if (B <= 8 && seq_length_k >= 4) {
138
+ max_splits = 8;
139
+ }
140
+ else if ((B <= 16 && seq_length_k >= 8) ||
141
+ (B == 48 && seq_length_k >= 32)) {
142
+ max_splits = 4;
143
+ }
144
+ else if ((B <= 32 && seq_length_k >= 16) ||
145
+ (B == 96 && seq_length_k >= 16)) {
146
+ max_splits = 2;
147
}
143
-
144
- // printf(" max_splits = %d\n", max_splits);
148
+ else {
149
+ max_splits = 1;
150
+ }
151
+
152
+ // Wave-aware scheduling: ensure integer number of waves in K dimension
153
int sms_per_batch = max(1, sm_count / B);
146
- // printf(" sms_per_batch = %d\n", sms_per_batch);
154
int split_heur = min(max_splits, sms_per_batch);
155
int waves = ceil_div(B * split_heur, sm_count);
156
int k_waves = ceil_div(max_splits, split_heur);
157
int split_wave_aware = ceil_div(max_splits, k_waves);
158
args.split_kv = split_wave_aware;
152
- // printf(" args.split_kv = %d\n", args.split_kv);
153
-
159
}
160
161
/// Determines whether the GEMM can execute the given problem.