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
[Bugfix][B200] Fix `cutlass_mla` hang (#24966)
Signed-off-by: Alexander Matveev <amatveev@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Alexander Matveev committed
Sep 17, 2025 at 18:06 UTC
fedb75fa2790403b90ec6dc926fef9c6c5ccb7a6
1 file changed
+8
csrc/attention/mla/cutlass_sm100_mla/device/sm100_mla.hpp
+8
@@ -133,6 +133,14 @@ public:
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 4 kv_splits.
139
+ // Discuss with NVIDIA how this can be fixed.
140
+ if (B > 1) {
141
+ max_splits = min(2, max_splits);
142
+ }
143
+
144
// printf(" max_splits = %d\n", max_splits);
145
int sms_per_batch = max(1, sm_count / B);
146
// printf(" sms_per_batch = %d\n", sms_per_batch);