| |
| #pragma once |
|
|
| #include <torch/all.h> |
|
|
| void w4a16_decode_gemv_bf16(torch::Tensor const& x_bf16, |
| torch::Tensor const& weight_packed, |
| torch::Tensor const& sfb, |
| double alpha, |
| torch::Tensor& out); |
|
|
| void grouped_w4a16_gemv_bf16(torch::Tensor const& activations, |
| torch::Tensor const& weight_stack, |
| torch::Tensor const& sfb_stack, |
| torch::Tensor const& alpha_stack, |
| torch::Tensor const& expert_idx, |
| int64_t w_stride, |
| int64_t sfb_stride, |
| torch::Tensor& out); |
|
|
| void quantize_activations_nvfp4_bf16(torch::Tensor const& activations, |
| torch::Tensor& packed, |
| torch::Tensor& sfa); |
|
|
| void quantize_weights_nvfp4_bf16(torch::Tensor const& weights, |
| torch::Tensor& packed, |
| torch::Tensor& sfb); |
|
|
| void grouped_w4a4_gemv_bf16(torch::Tensor const& activations_packed, |
| torch::Tensor const& weight_stack, |
| torch::Tensor const& sfa, |
| torch::Tensor const& sfb_stack, |
| torch::Tensor const& alpha_stack, |
| torch::Tensor const& expert_idx, |
| torch::Tensor& out); |
|
|
