squall/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix] Fix Machete unittests failing with NotImplementedError (#9218)

Browse Source
This commit is contained in:
Lucas Wilkinson 2024-10-10 13:39:56 -04:00 committed by GitHub
parent 83ea5c72b9
commit 18511aeda6
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
csrc/quantization/machete/machete_pytorch.cu
View File

@ -89,6 +89,10 @@ torch::Tensor prepack_B(torch::Tensor const& B,
TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) { TORCH_LIBRARY_IMPL_EXPAND(TORCH_EXTENSION_NAME, CUDA, m) {
m.impl("machete_prepack_B", &prepack_B); m.impl("machete_prepack_B", &prepack_B);
m.impl("machete_gemm", &gemm); m.impl("machete_gemm", &gemm);
}
// use CatchAll since supported_schedules has no tensor arguments
TORCH_LIBRARY_IMPL(TORCH_EXTENSION_NAME, CatchAll, m) {
m.impl("machete_supported_schedules", &supported_schedules); m.impl("machete_supported_schedules", &supported_schedules);
} }