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 15, 01:00 UTC
[BugFix] Fix import error on non-blackwell machines (#21020)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jul 16, 2025 at 01:27 UTC
d31a64712489d7c079fe48515c7ddd8a60bc0e71
3 files changed
+12
-16
csrc/attention/mla/sm100_cutlass_mla_kernel.cu
+10
@@ -18,6 +18,7 @@ limitations under the License.
18
* Taken from SGLANG PR https://github.com/sgl-project/sglang/pull/6929
19
* by Alcanderian JieXin Liang
20
*/
21
+#include "core/registration.h"
22
23
#include <ATen/cuda/CUDAContext.h>
24
#include <c10/cuda/CUDAGuard.h>
@@ -270,4 +271,13 @@ int64_t sm100_cutlass_mla_get_workspace_size(int64_t max_seq_len, int64_t num_ba
271
}
272
273
#endif
274
+
275
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
276
+ m.impl("sm100_cutlass_mla_decode", &sm100_cutlass_mla_decode);
277
+}
278
+
279
+TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CatchAll, m) {
280
+ m.impl("sm100_cutlass_mla_get_workspace_size", &sm100_cutlass_mla_get_workspace_size);
281
+}
282
+
283
// clang-format on
csrc/ops.h
-13
@@ -167,19 +167,6 @@ void cutlass_mla_decode(torch::Tensor const& out, torch::Tensor const& q_nope,
167
torch::Tensor const& seq_lens,
168
torch::Tensor const& page_table, double scale);
169
170
-void sm100_cutlass_mla_decode(
171
- torch::Tensor const& out, torch::Tensor const& q_nope,
172
- torch::Tensor const& q_pe, torch::Tensor const& kv_c_and_k_pe_cache,
173
- torch::Tensor const& seq_lens, torch::Tensor const& page_table,
174
- torch::Tensor const& workspace, double sm_scale,
175
- int64_t num_kv_splits =
176
- 1 /* Set to 1 to avoid cuda_graph issue by default. */);
177
-
178
-int64_t sm100_cutlass_mla_get_workspace_size(
179
- int64_t max_seq_len, int64_t num_batches, int64_t sm_count = 0,
180
- int64_t num_kv_splits =
181
- 1 /* Set to 1 to avoid cuda_graph issue by default. */);
182
-
170
torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor);
171
172
#ifndef USE_ROCM
csrc/torch_bindings.cpp
+2
-3
@@ -521,15 +521,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
521
" Tensor page_table, Tensor workspace, float "
522
"scale,"
523
" int num_kv_splits) -> ()");
524
- ops.impl("sm100_cutlass_mla_decode", torch::kCUDA, &sm100_cutlass_mla_decode);
524
+ // conditionally compiled so impl in source file
525
526
// SM100 CUTLASS MLA workspace
527
ops.def(
528
"sm100_cutlass_mla_get_workspace_size(int max_seq_len, int num_batches,"
529
" int sm_count, int num_kv_splits) "
530
"-> int");
531
- ops.impl("sm100_cutlass_mla_get_workspace_size",
532
- &sm100_cutlass_mla_get_workspace_size);
531
+ // conditionally compiled so impl in source file
532
533
// Compute NVFP4 block quantized tensor.
534
ops.def(