diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu index e910103c4e..84040a6a22 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8.cu @@ -1,5 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm100_fp8_dispatch.cuh" +#include "core/batch_invariant.hpp" namespace vllm { @@ -13,9 +14,17 @@ void cutlass_scaled_mm_sm100_fp8( STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(), "currently bias dtype must match output dtype ", out.scalar_type()); + if (vllm_is_batch_invariant()) { + return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue( + out, a, b, a_scales, b_scales, *bias); + } return cutlass_scaled_mm_sm100_fp8_epilogue(out, a, b, a_scales, b_scales, *bias); } else { + if (vllm_is_batch_invariant()) { + return cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue( + out, a, b, a_scales, b_scales); + } return cutlass_scaled_mm_sm100_fp8_epilogue(out, a, b, a_scales, b_scales); } diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh index 5cd55f0198..f790b3653d 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm100_fp8_dispatch.cuh @@ -294,6 +294,34 @@ inline void cutlass_gemm_sm100_fp8_dispatch( } } +template +inline void cutlass_gemm_sm100_fp8_batch_invariant_dispatch( + torch::stable::Tensor& out, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, EpilogueArgs&&... args) { + static_assert(std::is_same()); + STD_TORCH_CHECK(a.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn); + STD_TORCH_CHECK(b.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn); + + using Cutlass3xGemmM64SwapAB = + typename sm100_fp8_config_M64_swap_ab::Cutlass3xGemm; + using Cutlass3xGemmM64 = + typename sm100_fp8_config_M64::Cutlass3xGemm; + + // keep the CUTLASS config independent of M for batch invariance + uint32_t const k = a.size(1); + if (k < 4096) { + return cutlass_gemm_caller_sm100_fp8( + out, a, b, a_scales, b_scales, std::forward(args)...); + } + return cutlass_gemm_caller_sm100_fp8( + out, a, b, b_scales, a_scales, std::forward(args)...); +} + template void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out, torch::stable::Tensor const& a, @@ -320,4 +348,28 @@ void cutlass_scaled_mm_sm100_fp8_epilogue(torch::stable::Tensor& out, } } +template +void cutlass_scaled_mm_sm100_fp8_batch_invariant_epilogue( + torch::stable::Tensor& out, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, torch::stable::Tensor const& a_scales, + torch::stable::Tensor const& b_scales, EpilogueArgs&&... epilogue_args) { + STD_TORCH_CHECK(a.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn); + STD_TORCH_CHECK(b.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn); + + if (out.scalar_type() == torch::headeronly::ScalarType::BFloat16) { + return cutlass_gemm_sm100_fp8_batch_invariant_dispatch< + cutlass::float_e4m3_t, cutlass::bfloat16_t, EnableBias>( + out, a, b, a_scales, b_scales, + std::forward(epilogue_args)...); + } else { + STD_TORCH_CHECK(out.scalar_type() == torch::headeronly::ScalarType::Half); + return cutlass_gemm_sm100_fp8_batch_invariant_dispatch< + cutlass::float_e4m3_t, cutlass::half_t, EnableBias>( + out, a, b, a_scales, b_scales, + std::forward(epilogue_args)...); + } +} + } // namespace vllm diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu index fb84faa2a4..972d6c6260 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8.cu @@ -1,5 +1,6 @@ #include "scaled_mm_kernels.hpp" #include "scaled_mm_sm120_fp8_dispatch.cuh" +#include "core/batch_invariant.hpp" #include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp" namespace vllm { @@ -14,9 +15,17 @@ void cutlass_scaled_mm_sm120_fp8( STD_TORCH_CHECK(bias->scalar_type() == out.scalar_type(), "currently bias dtype must match output dtype ", out.scalar_type()); + if (vllm_is_batch_invariant()) { + return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue< + c3x::ScaledEpilogueBias>(out, a, b, a_scales, b_scales, *bias); + } return cutlass_scaled_mm_sm120_fp8_epilogue( out, a, b, a_scales, b_scales, *bias); } else { + if (vllm_is_batch_invariant()) { + return cutlass_scaled_mm_sm120_fp8_batch_invariant_epilogue< + c3x::ScaledEpilogue>(out, a, b, a_scales, b_scales); + } return cutlass_scaled_mm_sm120_fp8_epilogue( out, a, b, a_scales, b_scales); } diff --git a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh index 245f5c10fc..226e4f7a6b 100644 --- a/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh +++ b/csrc/libtorch_stable/quantization/w8a8/cutlass/c3x/scaled_mm_sm120_fp8_dispatch.cuh @@ -179,6 +179,26 @@ inline void cutlass_gemm_sm120_fp8_dispatch(torch::stable::Tensor& out, out, a, b, std::forward(args)...); } +template typename Epilogue, + typename... EpilogueArgs> +inline void cutlass_gemm_sm120_fp8_batch_invariant_dispatch( + torch::stable::Tensor& out, torch::stable::Tensor const& a, + torch::stable::Tensor const& b, EpilogueArgs&&... args) { + static_assert(std::is_same()); + STD_TORCH_CHECK(a.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn); + STD_TORCH_CHECK(b.scalar_type() == + torch::headeronly::ScalarType::Float8_e4m3fn); + + using Cutlass3xGemmM64 = + typename sm120_fp8_config_M64::Cutlass3xGemm; + + // keep the CUTLASS config independent of M for batch invariance + return cutlass_gemm_caller( + out, a, b, std::forward(args)...); +} + template