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

[Kernel][Misc] Use TORCH_LIBRARY instead of PYBIND11_MODULE for custom ops (#5047)

Browse Source
This commit is contained in:
bnellnm
2024-06-09 16:23:30 -04:00
committed by GitHub
parent 5d7e3d0176
commit 5467ac3196
55 changed files with 833 additions and 451 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/quantization/cutlass_w8a8/scaled_mm_dq_c2x.cu
View File

@@ -1,5 +1,5 @@
#include <stddef.h>
#include <torch/extension.h>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>

2
csrc/quantization/cutlass_w8a8/scaled_mm_dq_c3x.cu
View File

@@ -4,7 +4,7 @@
#if defined CUDA_VERSION && CUDA_VERSION >= 12000
#include <torch/extension.h>
#include <torch/all.h>
#include <ATen/cuda/CUDAContext.h>

2
csrc/quantization/cutlass_w8a8/scaled_mm_dq_entry.cu
View File

@@ -1,7 +1,7 @@
#include <cudaTypedefs.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/extension.h>
#include <torch/all.h>
void cutlass_scaled_mm_dq_sm75(torch::Tensor& c, torch::Tensor const& a,
torch::Tensor const& b,