@@ -255,12 +346,13 @@ __global__ void __launch_bounds__(512, 1)
tmps[i] = get_tmp_buf(sg.signals[target]);
}
auto tmp_out = tmps[0];
- multi_gpu_barrier(sg, self_sg, rank);
+ start_sync(sg, self_sg, rank);
+
// stage 1: reduce scatter
for (int idx = start + tid; idx < end; idx += stride) {
tmp_out[idx - start] = packed_reduce(ptrs, idx);
}
- multi_gpu_barrier(sg, self_sg, rank);
+ end_sync(sg, self_sg, rank);
// stage 2: allgather. Note: it's important to match the tid between
// the two stages, because visibility across devices is only guaranteed
@@ -290,7 +382,7 @@ class CustomAllreduce {
bool full_nvlink_;
RankSignals sg_;
- // Stores an map from a pointer to its peer pointters from all ranks.
+ // Stores an map from a pointer to its peer pointers from all ranks.
std::unordered_map buffers_;
Signal* self_sg_;
@@ -361,8 +453,7 @@ class CustomAllreduce {
void* base_ptr;
// note: must share the base address of each allocation, or we get wrong
// address
- if (cuPointerGetAttribute(&base_ptr,
- CU_POINTER_ATTRIBUTE_RANGE_START_ADDR,
+ if (cuPointerGetAttribute(&base_ptr, rangeStartAddrAttr,
(CUdeviceptr)ptr) != CUDA_SUCCESS)
throw std::runtime_error("failed to get pointer attr");
CUDACHECK(cudaIpcGetMemHandle(
@@ -439,7 +530,7 @@ class CustomAllreduce {
*/
template
void allreduce(cudaStream_t stream, T* input, T* output, int size,
- int threads = 512, int block_limit = 36) {
+ int threads = 512, int block_limit = defaultBlockLimit) {
auto d = packed_t::P::size;
if (size % d != 0)
throw std::runtime_error(
@@ -473,8 +564,6 @@ class CustomAllreduce {
#define KL(ngpus, name) \
name<<>>(ptrs, sg_, self_sg_, output, \
rank_, size);
- // TODO(hanzhi713): Threshold is different for A100 and H100.
- // Add per device threshold.
#define REDUCE_CASE(ngpus) \
case ngpus: { \
if (world_size_ == 2) { \
@@ -497,7 +586,8 @@ class CustomAllreduce {
REDUCE_CASE(8)
default:
throw std::runtime_error(
- "custom allreduce only supports num gpus in (2,4,6,8). Actual num "
+ "custom allreduce only supports num gpus in (2,4,6,8). Actual "
+ "num "
"gpus = " +
std::to_string(world_size_));
}
diff --git a/csrc/custom_all_reduce_test.cu b/csrc/custom_all_reduce_test.cu
index b59ea40d980f4fd661fffb9769c7ba5a53a31a4e..b9651dbb73112492bf93fd9fc258add81705b976 100644
--- a/csrc/custom_all_reduce_test.cu
+++ b/csrc/custom_all_reduce_test.cu
@@ -20,9 +20,16 @@
#include
#include "cuda_profiler_api.h"
-#include "custom_all_reduce.cuh"
#include "mpi.h"
-#include "nccl.h"
+#ifdef USE_ROCM
+ #include
+typedef __hip_bfloat16 nv_bfloat16;
+ #include "rccl/rccl.h"
+ #include "custom_all_reduce_hip.cuh"
+#else
+ #include "nccl.h"
+ #include "custom_all_reduce.cuh"
+#endif
#define MPICHECK(cmd) \
do { \
@@ -44,13 +51,16 @@
} while (0)
__global__ void dummy_kernel() {
-#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
- for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms
-#else
+#ifdef USE_ROCM
for (int i = 0; i < 100; i++) {
- long long int start = clock64();
- while (clock64() - start < 150000000); // approximately 98.4ms on P40
+ uint64_t start = wall_clock64();
+ uint64_t cycles_elapsed;
+ do {
+ cycles_elapsed = wall_clock64() - start;
+ } while (cycles_elapsed < 100);
}
+#else
+ for (int i = 0; i < 100; i++) __nanosleep(1000000); // 100ms
#endif
}
@@ -121,8 +131,14 @@ void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit,
* registration, they are allocated and registered together in the test for
* convenience.
*/
+#ifdef USE_ROCM
+ CUDACHECK(hipExtMallocWithFlags(
+ (void**)&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal),
+ hipDeviceMallocUncached));
+#else
CUDACHECK(
cudaMalloc(&buffer, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
+#endif
CUDACHECK(
cudaMemset(buffer, 0, 2 * data_size * sizeof(T) + sizeof(vllm::Signal)));
CUDACHECK(cudaMalloc(&self_data_copy, data_size * sizeof(T)));
@@ -135,26 +151,24 @@ void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit,
void* rank_data;
size_t rank_data_sz = 16 * 1024 * 1024;
CUDACHECK(cudaMalloc(&rank_data, rank_data_sz));
- vllm::Signal* ipc_ptrs[8];
- for (int i = 0; i < nRanks; i++) {
- if (i == myRank)
- ipc_ptrs[i] = buffer;
- else
- CUDACHECK(cudaIpcOpenMemHandle((void**)&ipc_ptrs[i], data_handles[i],
- cudaIpcMemLazyEnablePeerAccess));
- }
- vllm::CustomAllreduce fa(ipc_ptrs, rank_data, rank_data_sz, myRank, nRanks);
+ std::vector offsets(nRanks, 0);
+ vllm::CustomAllreduce fa(buffer, rank_data, rank_data_sz, data_handles,
+ offsets, myRank);
auto* self_data =
reinterpret_cast(reinterpret_cast(buffer) +
sizeof(vllm::Signal) + data_size * sizeof(T));
// hack buffer registration
{
- void* data[8];
+ std::vector handles;
+ handles.reserve(nRanks);
for (int i = 0; i < nRanks; i++) {
- data[i] =
- ((char*)ipc_ptrs[i]) + sizeof(vllm::Signal) + data_size * sizeof(T);
+ char* begin = (char*)&data_handles[i];
+ char* end = (char*)&data_handles[i + 1];
+ handles.emplace_back(begin, end);
}
- fa.register_buffer(data);
+ std::vector offsets(nRanks,
+ sizeof(vllm::Signal) + data_size * sizeof(T));
+ fa.register_buffer(handles, offsets, self_data);
}
double* ground_truth;
@@ -266,14 +280,14 @@ void run(int myRank, int nRanks, ncclComm_t& comm, int threads, int block_limit,
if (diff >= 4e-2) {
printf(
"Rank %d: Verification mismatch at %lld: %f != (my) %f, gt=%f\n",
- myRank, j, nccl_result[j], my_result[j], ground_truth[j]);
+ myRank, j, nccl_result[j], my_result[j], ground_truth[j]);
break;
}
}
}
if (myRank == 0)
printf("Test passed: nGPUs:%d, sz (kb): %d, %d, %d\n", nRanks,
- data_size * sizeof(T) / 1024, threads, block_limit);
+ data_size * sizeof(T) / 1024, threads, block_limit);
// long double nccl_diffs = 0.0;
// long double my_diffs = 0.0;
// for (int j = 0; j < data_size; j++) {
@@ -306,24 +320,27 @@ int main(int argc, char** argv) {
ncclComm_t comm;
if (myRank == 0) ncclGetUniqueId(&id);
MPICHECK(MPI_Bcast(static_cast(&id), sizeof(id), MPI_BYTE, 0,
- MPI_COMM_WORLD));
+ MPI_COMM_WORLD));
NCCLCHECK(ncclCommInitRank(&comm, nRanks, id, myRank));
bool performance_test = true;
cudaProfilerStart();
- // Uncomment to scan through different block size configs.
- // for (int threads : {256, 512, 1024}) {
+ // for (int threads : {256, 512}) {
// for (int block_limit = 16; block_limit < 112; block_limit += 4) {
- // run(myRank, nRanks, comm, threads, block_limit, 1024 * 1024,
- // performance_test);
+ // run(myRank, nRanks, comm, threads, block_limit, 4096 * 1024);
// }
// }
- // Scan through different sizes to test performance.
+#ifdef USE_ROCM
+ for (int sz = 512; sz <= (8 << 20); sz *= 2) {
+ run(myRank, nRanks, comm, 512, 16, sz + 8 * 47, performance_test);
+ }
+#else
for (int sz = 512; sz <= (8 << 20); sz *= 2) {
run(myRank, nRanks, comm, 512, 36, sz + 8 * 47, performance_test);
}
+#endif
cudaProfilerStop();
MPICHECK(MPI_Finalize());
return EXIT_SUCCESS;
-}
+}
diff --git a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
index ef413e6dd75c5d73c515e48f2f5fa55c47f20eb0..64b7ddae3d2d780525729a314e733a0f76ce5760 100644
--- a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
+++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c2x.hpp
@@ -122,8 +122,8 @@ struct ScaledEpilogue
auto a_args = SUPER::template args_from_tensor(a_scales);
auto b_args = SUPER::template args_from_tensor(b_scales);
- typename EVTCompute0::Arguments evt0_args{b_args};
- return ArgumentType{a_args, evt0_args};
+ typename EVTCompute0::Arguments evt0_args{b_args, {}, {}};
+ return ArgumentType{a_args, evt0_args, {}};
}
};
@@ -167,8 +167,8 @@ struct ScaledEpilogueBias
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
- typename EVTCompute0::Arguments evt0_args{b_args};
- return ArgumentType{a_args, evt0_args, bias_args};
+ typename EVTCompute0::Arguments evt0_args{b_args, {}, {}};
+ return ArgumentType{a_args, evt0_args, bias_args, {}};
}
};
@@ -230,9 +230,10 @@ struct ScaledEpilogueBiasAzp
auto azp_adj_args =
SUPER::template args_from_tensor(azp_adj);
- typename EVTComputeAzp::Arguments evt_azp_args{{}, azp_adj_args};
- typename EVTComputeScaleB::Arguments evt_scale_b_args{b_args, evt_azp_args};
- return ArgumentType{a_args, evt_scale_b_args, bias_args};
+ typename EVTComputeAzp::Arguments evt_azp_args{{}, azp_adj_args, {}};
+ typename EVTComputeScaleB::Arguments evt_scale_b_args{
+ b_args, evt_azp_args, {}};
+ return ArgumentType{a_args, evt_scale_b_args, bias_args, {}};
}
};
@@ -309,11 +310,12 @@ struct ScaledEpilogueBiasAzpToken
auto azp_adj_args =
SUPER::template args_from_tensor(azp_adj);
- typename EVTComputeAzp::Arguments evt_azp_args{azp_args, azp_adj_args};
- typename EVTComputeAcc::Arguments evt_acc_args{{}, evt_azp_args};
- typename EVTComputeScaleB::Arguments evt_scale_b_args{b_args, evt_acc_args};
- return ArgumentType{a_args, evt_scale_b_args, bias_args};
+ typename EVTComputeAzp::Arguments evt_azp_args{azp_args, azp_adj_args, {}};
+ typename EVTComputeAcc::Arguments evt_acc_args{{}, evt_azp_args, {}};
+ typename EVTComputeScaleB::Arguments evt_scale_b_args{
+ b_args, evt_acc_args, {}};
+ return ArgumentType{a_args, evt_scale_b_args, bias_args, {}};
}
};
-}; // namespace vllm::c2x
\ No newline at end of file
+}; // namespace vllm::c2x
diff --git a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
index 583fa3c4551150b98f0cad97c5abd8e1a4a3d8eb..0a812dc56a994c64994dbb0c11222ddb339f2a32 100644
--- a/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+++ b/csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
@@ -22,7 +22,7 @@ struct identity {
T operator()(T lhs) const { return lhs; }
};
-template
+template
struct TrivialEpilogue {
private:
using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
@@ -44,32 +44,30 @@ struct TrivialEpilogue {
* This class provides the common load descriptors for the
* ScaledEpilogue[...] classes
*/
-template
+template
struct ScaledEpilogueBase {
protected:
using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
template
using ColOrScalarLoad = cutlass::epilogue::fusion::Sm90ColOrScalarBroadcast<
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T,
- Stride, Int<0>, Int<0>>>;
+ 0 /*Stages*/, TileShape, T, Stride, Int<0>, Int<0>>>;
template
using RowOrScalarLoad = cutlass::epilogue::fusion::Sm90RowOrScalarBroadcast<
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T,
- Stride, Int<1>, Int<0>>>;
+ 0 /*Stages*/, TileShape, T, Stride, Int<1>, Int<0>>>;
// Don't want to support nullptr by default
template
using ColLoad = cutlass::epilogue::fusion::Sm90ColBroadcast<
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T, T,
- Stride, Int<0>, Int<0>>, 128 / sizeof_bits_v, EnableNullPtr>;
+ 0 /*Stages*/, TileShape, T, T, Stride, Int<0>, Int<0>>,
+ 128 / sizeof_bits_v, EnableNullPtr>;
// Don't want to support nullptr by default
template
using RowLoad = cutlass::epilogue::fusion::Sm90RowBroadcast<
- 0 /*Stages*/, typename EpilogueDescriptor::TileShape, T, T,
- Stride, Int<1>, Int<0>>, 128 / sizeof_bits_v, EnableNullPtr>;
+ 0 /*Stages*/, TileShape, T, T, Stride, Int<1>, Int<0>>,
+ 128 / sizeof_bits_v, EnableNullPtr>;
// This utility function constructs the arguments for the load descriptors
// from a tensor. It can handle both row and column, as well as row/column or
@@ -116,11 +114,11 @@ struct ScaledEpilogueBase {
the A and B operands respectively. These scales may be either per-tensor or
per row or column.
*/
-template
+template
struct ScaledEpilogue
- : private ScaledEpilogueBase {
+ : private ScaledEpilogueBase {
private:
- using SUPER = ScaledEpilogueBase;
+ using SUPER = ScaledEpilogueBase;
using Accum = typename SUPER::Accum;
using ScaleA = typename SUPER::template ColOrScalarLoad;
using ScaleB = typename SUPER::template RowOrScalarLoad;
@@ -146,8 +144,8 @@ struct ScaledEpilogue
auto a_args = SUPER::template args_from_tensor(a_scales);
auto b_args = SUPER::template args_from_tensor(b_scales);
- typename EVTCompute0::Arguments evt0_args{b_args};
- return ArgumentType{a_args, evt0_args};
+ typename EVTCompute0::Arguments evt0_args{b_args, {}, {}};
+ return ArgumentType{a_args, evt0_args, {}};
}
};
@@ -160,11 +158,11 @@ struct ScaledEpilogue
* The bias tensor must be per-output channel.
* ScaleA and ScaleB can be per-tensor or per-token/per-channel.
*/
-template
+template
struct ScaledEpilogueBias
- : private ScaledEpilogueBase {
+ : private ScaledEpilogueBase {
private:
- using SUPER = ScaledEpilogueBase;
+ using SUPER = ScaledEpilogueBase;
using Accum = typename SUPER::Accum;
using ScaleA = typename SUPER::template ColOrScalarLoad;
using ScaleB = typename SUPER::template RowOrScalarLoad;
@@ -193,8 +191,8 @@ struct ScaledEpilogueBias
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
- typename EVTCompute0::Arguments evt0_args{b_args};
- return ArgumentType{a_args, evt0_args, bias_args};
+ typename EVTCompute0::Arguments evt0_args{b_args, {}, {}};
+ return ArgumentType{a_args, evt0_args, bias_args, {}};
}
};
@@ -203,11 +201,11 @@ struct ScaledEpilogueBias
* bias is a column vector instead of a row vector. Useful e.g. if we are
* computing a GEMM via C^T += B^T A^T. This happens in the 2:4 sparse kernels.
*/
-template
+template
struct ScaledEpilogueColumnBias
- : private ScaledEpilogueBase {
+ : private ScaledEpilogueBase {
private:
- using SUPER = ScaledEpilogueBase;
+ using SUPER = ScaledEpilogueBase;
using Accum = typename SUPER::Accum;
using ScaleA = typename SUPER::template ColOrScalarLoad;
using ScaleB = typename SUPER::template RowOrScalarLoad;
@@ -236,8 +234,8 @@ struct ScaledEpilogueColumnBias
auto b_args = SUPER::template args_from_tensor(b_scales);
auto bias_args = SUPER::template args_from_tensor(bias);
- typename EVTCompute0::Arguments evt0_args{b_args};
- return ArgumentType{a_args, evt0_args, bias_args};
+ typename EVTCompute0::Arguments evt0_args{b_args, {}, {}};
+ return ArgumentType{a_args, evt0_args, bias_args, {}};
}
};
@@ -249,11 +247,11 @@ struct ScaledEpilogueColumnBias
*
* This epilogue also supports bias, which remains per-channel.
*/
-template
+template
struct ScaledEpilogueBiasAzp
- : private ScaledEpilogueBase {
+ : private ScaledEpilogueBase {
private:
- using SUPER = ScaledEpilogueBase;
+ using SUPER = ScaledEpilogueBase;
using Accum = typename SUPER::Accum;
using ScaleA = typename SUPER::template ColOrScalarLoad;
using ScaleB = typename SUPER::template RowOrScalarLoad;
@@ -297,9 +295,10 @@ struct ScaledEpilogueBiasAzp
auto azp_adj_args =
SUPER::template args_from_tensor(azp_adj);
- typename EVTComputeAzp::Arguments evt_azp_args{{}, azp_adj_args};
- typename EVTComputeScaleB::Arguments evt_scale_b_args{b_args, evt_azp_args};
- return ArgumentType{a_args, evt_scale_b_args, bias_args};
+ typename EVTComputeAzp::Arguments evt_azp_args{{}, azp_adj_args, {}};
+ typename EVTComputeScaleB::Arguments evt_scale_b_args{
+ b_args, evt_azp_args, {}};
+ return ArgumentType{a_args, evt_scale_b_args, bias_args, {}};
}
};
@@ -313,11 +312,11 @@ struct ScaledEpilogueBiasAzp
*
* This epilogue also supports bias, which remains per-channel.
*/
-template
+template
struct ScaledEpilogueBiasAzpToken
- : private ScaledEpilogueBase {
+ : private ScaledEpilogueBase {
private:
- using SUPER = ScaledEpilogueBase;
+ using SUPER = ScaledEpilogueBase;
using Accum = typename SUPER::Accum;
using ScaleA = typename SUPER::template ColOrScalarLoad;
using ScaleB = typename SUPER::template RowOrScalarLoad;
@@ -374,10 +373,11 @@ struct ScaledEpilogueBiasAzpToken
auto azp_adj_args =
SUPER::template args_from_tensor(azp_adj);
- typename EVTComputeAzp::Arguments evt_azp_args{azp_args, azp_adj_args};
- typename EVTComputeAcc::Arguments evt_acc_args{{}, evt_azp_args};
- typename EVTComputeScaleB::Arguments evt_scale_b_args{b_args, evt_acc_args};
- return ArgumentType{a_args, evt_scale_b_args, bias_args};
+ typename EVTComputeAzp::Arguments evt_azp_args{azp_args, azp_adj_args, {}};
+ typename EVTComputeAcc::Arguments evt_acc_args{{}, evt_azp_args, {}};
+ typename EVTComputeScaleB::Arguments evt_scale_b_args{
+ b_args, evt_acc_args, {}};
+ return ArgumentType{a_args, evt_scale_b_args, bias_args, {}};
}
};
diff --git a/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp b/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
index 928a9500cbb08092176604f7903e60b7a454974c..d922a3349e1e19b73dcd2282a36b429838a71820 100644
--- a/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
+++ b/csrc/cutlass_extensions/gemm/collective/sm90_mma_tma_gmma_ss_warpspecialized_fp8_blockwise_scaling.hpp
@@ -402,7 +402,7 @@ struct CollectiveMma<
// TODO: test `scale_copy_a` with `ScaleMsPerTile` < 128
TiledCopy scale_copy_a = make_tiled_copy(SmemBlockScalingCopyAtomA{},
- Layout>{}, Layout>{}); // (1,1,1)
+ Layout>{}, Layout>{}); // (1,1,1)
TiledCopy scale_copy_b = make_tiled_copy(SmemBlockScalingCopyAtomB{},
Layout>{}, Layout>{}); // (1,1,1)
ThrCopy thr_scale_copy_a = scale_copy_a.get_slice(threadIdx.x);
diff --git a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py
index d5a5e2ef83dd808fc4b0815784b05488c3e4888b..d64f0d0a5c2a046bcf3210aab83bf1d6f9380232 100644
--- a/csrc/cutlass_extensions/vllm_cutlass_library_extension.py
+++ b/csrc/cutlass_extensions/vllm_cutlass_library_extension.py
@@ -1,7 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
import enum
-from typing import Dict, Union
+from typing import Union
from cutlass_library import *
@@ -21,7 +21,7 @@ class MixedInputKernelScheduleType(enum.Enum):
TmaWarpSpecializedCooperative = enum_auto()
-VLLMDataTypeNames: Dict[Union[VLLMDataType, DataType], str] = {
+VLLMDataTypeNames: dict[Union[VLLMDataType, DataType], str] = {
**DataTypeNames, # type: ignore
**{
VLLMDataType.u4b8: "u4b8",
@@ -29,7 +29,7 @@ VLLMDataTypeNames: Dict[Union[VLLMDataType, DataType], str] = {
}
}
-VLLMDataTypeTag: Dict[Union[VLLMDataType, DataType], str] = {
+VLLMDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {
**DataTypeTag, # type: ignore
**{
VLLMDataType.u4b8: "cutlass::vllm_uint4b8_t",
@@ -37,7 +37,7 @@ VLLMDataTypeTag: Dict[Union[VLLMDataType, DataType], str] = {
}
}
-VLLMDataTypeSize: Dict[Union[VLLMDataType, DataType], int] = {
+VLLMDataTypeSize: dict[Union[VLLMDataType, DataType], int] = {
**DataTypeSize, # type: ignore
**{
VLLMDataType.u4b8: 4,
@@ -45,7 +45,7 @@ VLLMDataTypeSize: Dict[Union[VLLMDataType, DataType], int] = {
}
}
-VLLMDataTypeVLLMScalarTypeTag: Dict[Union[VLLMDataType, DataType], str] = {
+VLLMDataTypeVLLMScalarTypeTag: dict[Union[VLLMDataType, DataType], str] = {
VLLMDataType.u4b8: "vllm::kU4B8",
VLLMDataType.u8b128: "vllm::kU8B128",
DataType.u4: "vllm::kU4",
@@ -56,7 +56,7 @@ VLLMDataTypeVLLMScalarTypeTag: Dict[Union[VLLMDataType, DataType], str] = {
DataType.bf16: "vllm::kBfloat16",
}
-VLLMDataTypeTorchDataTypeTag: Dict[Union[VLLMDataType, DataType], str] = {
+VLLMDataTypeTorchDataTypeTag: dict[Union[VLLMDataType, DataType], str] = {
DataType.u8: "at::ScalarType::Byte",
DataType.s8: "at::ScalarType::Char",
DataType.e4m3: "at::ScalarType::Float8_e4m3fn",
@@ -66,7 +66,7 @@ VLLMDataTypeTorchDataTypeTag: Dict[Union[VLLMDataType, DataType], str] = {
DataType.f32: "at::ScalarType::Float",
}
-VLLMKernelScheduleTag: Dict[Union[
+VLLMKernelScheduleTag: dict[Union[
MixedInputKernelScheduleType, KernelScheduleType], str] = {
**KernelScheduleTag, # type: ignore
**{
diff --git a/csrc/dispatch_utils.h b/csrc/dispatch_utils.h
index 03414b7e1ae9328f223ce12bc1bb7fd9c69860ad..dc6e0769b8780d091d835395d9f5b7833f561533 100644
--- a/csrc/dispatch_utils.h
+++ b/csrc/dispatch_utils.h
@@ -6,6 +6,11 @@
#include
+// Need a special dispatch case macro since we will nest the FP8 dispatch.
+// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
+#define AT_DISPATCH_FP8_CASE(enum_type, ...) \
+ AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
+
#define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
@@ -14,17 +19,32 @@
#define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
-// TODO(luka/varun): use FP8_TYPE macro after refactoring
-#ifndef USE_ROCM
- #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
- AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
- AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
-#else
+// ROCm devices might use either fn or fnuz, so set up dispatch table for both.
+// A host-based check at runtime will create a preferred FP8 type for ROCm
+// such that the correct kernel is dispatched.
+#ifdef USE_ROCM
+ #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
+ AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
+ AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
+
#define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
+ AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
+#else
+ #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
+ AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)
+
+ #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
+ AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
+ AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
#endif
+// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
+// See AT_DISPATCH_FP8_CASE above.
+#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
+ AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
+
#define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
diff --git a/csrc/layernorm_quant_kernels.cu b/csrc/layernorm_quant_kernels.cu
index c18e2a4e4abe00caf81848c4095b280961e47828..d595b9e889c84d6db632ad3af5218119d5666774 100644
--- a/csrc/layernorm_quant_kernels.cu
+++ b/csrc/layernorm_quant_kernels.cu
@@ -21,9 +21,9 @@
namespace vllm {
// TODO(woosuk): Further optimize this kernel.
-template
+template
__global__ void rms_norm_static_fp8_quant_kernel(
- FP8_TYPE* __restrict__ out, // [..., hidden_size]
+ fp8_type* __restrict__ out, // [..., hidden_size]
const scalar_t* __restrict__ input, // [..., hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
const float* __restrict__ scale, // [1]
@@ -52,7 +52,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(
float x = (float)input[blockIdx.x * hidden_size + idx];
float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
out[blockIdx.x * hidden_size + idx] =
- scaled_fp8_conversion(out_norm, scale_inv);
+ scaled_fp8_conversion(out_norm, scale_inv);
}
}
@@ -60,10 +60,10 @@ __global__ void rms_norm_static_fp8_quant_kernel(
Additional optimizations we can make in this case are
packed and vectorized operations, which help with the
memory latency bottleneck. */
-template
+template
__global__ std::enable_if_t<(width > 0) && _typeConvert::exists>
fused_add_rms_norm_static_fp8_quant_kernel(
- FP8_TYPE* __restrict__ out, // [..., hidden_size]
+ fp8_type* __restrict__ out, // [..., hidden_size]
scalar_t* __restrict__ input, // [..., hidden_size]
scalar_t* __restrict__ residual, // [..., hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
@@ -114,7 +114,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
#pragma unroll
for (int i = 0; i < width; ++i) {
out[id * width + i] =
- scaled_fp8_conversion(float(temp.data[i]), scale_inv);
+ scaled_fp8_conversion(float(temp.data[i]), scale_inv);
}
}
}
@@ -122,10 +122,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
/* Generic fused_add_rms_norm_kernel
The width field is not used here but necessary for other specializations.
*/
-template
+template
__global__ std::enable_if_t<(width == 0) || !_typeConvert::exists>
fused_add_rms_norm_static_fp8_quant_kernel(
- FP8_TYPE* __restrict__ out, // [..., hidden_size]
+ fp8_type* __restrict__ out, // [..., hidden_size]
scalar_t* __restrict__ input, // [..., hidden_size]
scalar_t* __restrict__ residual, // [..., hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
@@ -158,7 +158,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
float x = (float)residual[blockIdx.x * hidden_size + idx];
float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
out[blockIdx.x * hidden_size + idx] =
- scaled_fp8_conversion(out_norm, scale_inv);
+ scaled_fp8_conversion(out_norm, scale_inv);
}
}
@@ -176,25 +176,33 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]
dim3 block(std::min(hidden_size, 1024));
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
- vllm::rms_norm_static_fp8_quant_kernel
- <<>>(
- out.data_ptr(), input.data_ptr(),
- weight.data_ptr(), scale.data_ptr(), epsilon,
- num_tokens, hidden_size);
- });
+ VLLM_DISPATCH_FLOATING_TYPES(
+ input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
+ VLLM_DISPATCH_FP8_TYPES(
+ out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
+ vllm::rms_norm_static_fp8_quant_kernel
+ <<>>(
+ out.data_ptr(), input.data_ptr(),
+ weight.data_ptr(), scale.data_ptr(),
+ epsilon, num_tokens, hidden_size);
+ });
+ });
}
-#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
- VLLM_DISPATCH_FLOATING_TYPES( \
- input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
- vllm::fused_add_rms_norm_static_fp8_quant_kernel \
- <<>>( \
- out.data_ptr(), input.data_ptr(), \
- residual.data_ptr(), weight.data_ptr(), \
- scale.data_ptr(), epsilon, num_tokens, hidden_size); \
+#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
+ VLLM_DISPATCH_FLOATING_TYPES( \
+ input.scalar_type(), "fused_add_rms_norm_kernel_scalar_type", [&] { \
+ VLLM_DISPATCH_FP8_TYPES( \
+ out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
+ vllm::fused_add_rms_norm_static_fp8_quant_kernel \
+ <<>>( \
+ out.data_ptr(), input.data_ptr(), \
+ residual.data_ptr(), \
+ weight.data_ptr(), scale.data_ptr(), \
+ epsilon, num_tokens, hidden_size); \
+ }); \
});
-
void fused_add_rms_norm_static_fp8_quant(
torch::Tensor& out, // [..., hidden_size],
torch::Tensor& input, // [..., hidden_size]
diff --git a/csrc/moe/moe_ops.h b/csrc/moe/moe_ops.h
index 257c79a82ece1c52e76712335955d5fb4f8ca8b1..12321ddaf0d970c1da1505cb0ed24a07925699fe 100644
--- a/csrc/moe/moe_ops.h
+++ b/csrc/moe/moe_ops.h
@@ -24,3 +24,14 @@ void sgl_moe_align_block_size(torch::Tensor topk_ids, int64_t num_experts,
torch::Tensor sorted_token_ids,
torch::Tensor experts_ids,
torch::Tensor num_tokens_post_pad);
+#ifndef USE_ROCM
+torch::Tensor moe_wna16_gemm(torch::Tensor input, torch::Tensor output,
+ torch::Tensor b_qweight, torch::Tensor b_scales,
+ std::optional b_qzeros,
+ std::optional topk_weights,
+ torch::Tensor sorted_token_ids,
+ torch::Tensor expert_ids,
+ torch::Tensor num_tokens_post_pad, int64_t top_k,
+ int64_t BLOCK_SIZE_M, int64_t BLOCK_SIZE_N,
+ int64_t BLOCK_SIZE_K, int64_t bit);
+#endif
\ No newline at end of file
diff --git a/csrc/moe/moe_wna16.cu b/csrc/moe/moe_wna16.cu
new file mode 100644
index 0000000000000000000000000000000000000000..51ae76c1ec882aaac25b2fafabad5abb0a007460
--- /dev/null
+++ b/csrc/moe/moe_wna16.cu
@@ -0,0 +1,346 @@
+
+#include