Skip to content

Commit 10a0e08

Browse files
LucasWilkinsonsumitd2
authored andcommitted
[Bugfix] Fix Machete unittests failing with NotImplementedError (vllm-project#9218)
Signed-off-by: Sumit Dubey <[email protected]>
1 parent 5af3b76 commit 10a0e08

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

csrc/quantization/machete/machete_pytorch.cu

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,10 @@ torch::Tensor prepack_B(torch::Tensor const& B,
8989
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
9090
m.impl("machete_prepack_B", &prepack_B);
9191
m.impl("machete_gemm", &gemm);
92+
}
93+
94+
// use CatchAll since supported_schedules has no tensor arguments
95+
TORCH_LIBRARY_IMPL(TORCH_EXTENSION_NAME, CatchAll, m) {
9296
m.impl("machete_supported_schedules", &supported_schedules);
9397
}
9498

0 commit comments

Comments
 (0)