| | #include "scaled_mm_kernels.hpp" |
| | #include "scaled_mm_blockwise_sm100_fp8_dispatch.cuh" |
| | #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" |
| |
|
| | namespace vllm { |
| |
|
| | void cutlass_scaled_mm_blockwise_sm100_fp8(torch::Tensor& out, |
| | torch::Tensor const& a, |
| | torch::Tensor const& b, |
| | torch::Tensor const& a_scales, |
| | torch::Tensor const& b_scales) { |
| | if (out.dtype() == torch::kBFloat16) { |
| | cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::bfloat16_t>( |
| | out, a, b, a_scales, b_scales); |
| |
|
| | } else { |
| | TORCH_CHECK(out.dtype() == torch::kFloat16); |
| | cutlass_gemm_blockwise_sm100_fp8_dispatch<cutlass::half_t>( |
| | out, a, b, a_scales, b_scales); |
| | } |
| | } |
| |
|
| | } |
| |
|
