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

[ Kernel ] FP8 Dynamic Per Token Quant - Add scale_ub (#6593)

Browse Source 
Co-authored-by: Varun Sundar Rabindranth <varun@neuralmagic.com>
This commit is contained in:
Varun Sundar Rabindranath
2024-07-19 21:15:26 -04:00
committed by GitHub
parent e81522e879
commit 2e26564259
6 changed files with 86 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/torch_bindings.cpp
View File

@@ -188,7 +188,7 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
// Compute dynamic-per-token FP8 quantized tensor and scaling factor.
ops.def(
"dynamic_per_token_scaled_fp8_quant(Tensor! out, Tensor input, Tensor! "
"scale) -> "
"scale, Tensor? scale_ub) -> "
"()");
ops.impl("dynamic_per_token_scaled_fp8_quant", torch::kCUDA,
&dynamic_per_token_scaled_fp8_quant);