diff --git a/csrc/kernels/gemm/turbo_gemm.cu b/csrc/kernels/gemm/turbo_gemm.cu index caf05a336..4e9a28dd1 100644 --- a/csrc/kernels/gemm/turbo_gemm.cu +++ b/csrc/kernels/gemm/turbo_gemm.cu @@ -3,11 +3,8 @@ // See LICENSE for license information. #include "primus_turbo/gemm.h" -#include - -#ifdef BUILD_TURBO_BACKEND - #include "turbo/turbo_gemm_mxfp8_kernel.h" +#include namespace primus_turbo { @@ -29,6 +26,9 @@ void turbo_gemm_mxfp8_impl(const AType *a_ptr, const BType *b_ptr, const dtype::float8_e8m0 *b_scale_ptr, CType *c_ptr, int32_t m, int32_t n, int32_t k, void *workspace, size_t workspace_size, hipStream_t stream) { + PRIMUS_TURBO_CHECK(!primus_turbo::is_gfx1250(), + "turbo_gemm_fp8 is unavailable: the turbo backend is not built " + "on this architecture (gfx1250 is unsupported)."); constexpr int32_t MX_BLOCK_SIZE = 32; const int32_t scale_cols = (k + MX_BLOCK_SIZE - 1) / MX_BLOCK_SIZE; const size_t a_scale_bytes = (size_t) m * scale_cols * sizeof(uint32_t); @@ -73,44 +73,3 @@ INSTANTIATE_TURBO_GEMM(dtype::float8_e5m2, dtype::float8_e4m3, dtype::bfloat16) #undef INSTANTIATE_TURBO_GEMM } // namespace primus_turbo - -#else // !BUILD_TURBO_BACKEND : turbo (MFMA) GEMM not compiled on gfx1250. - -#include "primus_turbo/macros.h" - -namespace primus_turbo { - -// Stub symbols so host bindings link uniformly. The real kernels are unavailable on -// gfx1250; callers must gate on is_gfx1250() before invoking the impl. - -size_t turbo_gemm_mxfp8_workspace_size(int32_t, int32_t, int32_t) { - return 0; -} - -template -void turbo_gemm_mxfp8_impl(const AType *, const BType *, const dtype::float8_e8m0 *, - const dtype::float8_e8m0 *, CType *, int32_t, int32_t, int32_t, void *, - size_t, hipStream_t) { - PRIMUS_TURBO_ERROR("turbo_gemm_mxfp8 is unavailable: the turbo backend is not built " - "on this architecture (gfx1250 is unsupported)."); -} - -#define INSTANTIATE_TURBO_GEMM(A, B, C) \ - template void turbo_gemm_mxfp8_impl(const A *, const B *, const dtype::float8_e8m0 *, \ - const dtype::float8_e8m0 *, C *, int32_t, \ - int32_t, int32_t, void *, size_t, hipStream_t); - -INSTANTIATE_TURBO_GEMM(dtype::float8_e4m3, dtype::float8_e4m3, dtype::float16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e4m3, dtype::float8_e4m3, dtype::bfloat16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e5m2, dtype::float8_e5m2, dtype::float16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e5m2, dtype::float8_e5m2, dtype::bfloat16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e4m3, dtype::float8_e5m2, dtype::float16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e4m3, dtype::float8_e5m2, dtype::bfloat16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e5m2, dtype::float8_e4m3, dtype::float16) -INSTANTIATE_TURBO_GEMM(dtype::float8_e5m2, dtype::float8_e4m3, dtype::bfloat16) - -#undef INSTANTIATE_TURBO_GEMM - -} // namespace primus_turbo - -#endif // BUILD_TURBO_BACKEND diff --git a/csrc/kernels/quantization/dequantization_mxfp4.cu b/csrc/kernels/quantization/dequantization_mxfp4.cu index cad0c6a6d..96a506b84 100644 --- a/csrc/kernels/quantization/dequantization_mxfp4.cu +++ b/csrc/kernels/quantization/dequantization_mxfp4.cu @@ -46,6 +46,16 @@ __device__ __forceinline__ void cvt_fp4x2_scaled(uint32_t fp4, float scale, OTyp : "v"(fp4), "v"(scale)); *reinterpret_cast(out) = r; } +#elif defined(__gfx1250__) + // gfx1250 (CDNA5) only exposes the PK8 scale converter for de-quantization, and + // its E8M0 scale is selected cross-lane (lanes 16..31 read lanes 0..15). To keep + // each thread's scale independent we run the converter as a pure FP4->F32 unpack + // (scale byte = 0x7f == 2^0, opsel 0) and apply the per-thread ``scale`` in + // registers. Only the low 2 nibbles of ``fp4`` are meaningful here. + typedef float float32x8_t __attribute__((ext_vector_type(8))); + const float32x8_t r = __builtin_amdgcn_cvt_scale_pk8_f32_fp4(fp4, 0x7f7f7f7fu, 0); + out[0] = static_cast(r[0] * scale); + out[1] = static_cast(r[1] * scale); #else __builtin_trap(); #endif @@ -154,10 +164,6 @@ void dequantize_mxfp4_impl(const uint8_t *x, OType *y, const int64_t stride_x_ro const int64_t stride_scale_col, const int scale_n_rows, const int scale_n_cols, const int block_size, const bool use_rowwise, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("dequantize_mxfp4 is not implemented on gfx1250."); - } (void) stride_x_col; // packed input is contiguous along columns (stride == 1) if (n_rows == 0 || n_cols == 0) return; diff --git a/csrc/kernels/quantization/dequantization_mxfp8.cu b/csrc/kernels/quantization/dequantization_mxfp8.cu index b98bd0d56..7fc450542 100644 --- a/csrc/kernels/quantization/dequantization_mxfp8.cu +++ b/csrc/kernels/quantization/dequantization_mxfp8.cu @@ -258,10 +258,6 @@ void grouped_dequantize_mxfp8_impl(const QType *x, OType *y, const int64_t strid const int scale_n_rows, const int scale_n_cols, const int64_t *group_offs, const int64_t *group_offs_padded, int G, int block_size, bool use_rowwise, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("grouped_dequantize_mxfp8 is not implemented on gfx1250."); - } if (total_M == 0 || n_cols == 0) return; @@ -289,10 +285,6 @@ void dequantize_mxfp8_impl(const QType *x, OType *y, const int64_t stride_x_row, const int64_t stride_scale_col, const int scale_n_rows, const int scale_n_cols, const int block_size, const bool use_rowwise, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("dequantize_mxfp8 is not implemented on gfx1250."); - } (void) stride_x_col; // input is contiguous along columns (stride == 1) if (n_rows == 0 || n_cols == 0) return; diff --git a/csrc/kernels/quantization/quantization_mxfp4.cu b/csrc/kernels/quantization/quantization_mxfp4.cu index f4bc11838..1d22b5d45 100644 --- a/csrc/kernels/quantization/quantization_mxfp4.cu +++ b/csrc/kernels/quantization/quantization_mxfp4.cu @@ -39,9 +39,9 @@ using namespace primus_turbo::detail; // ============================================================================ // Hardware architecture parameters -constexpr int WARP_SIZE = 64; // AMD wavefront size -constexpr int THREADS_PER_BLOCK = 256; // 4 warps per block -constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE; +constexpr int WARP_SIZE = THREADS_PER_WARP; +constexpr int WARPS_PER_BLOCK = 4; +constexpr int THREADS_PER_BLOCK = WARP_SIZE * WARPS_PER_BLOCK; // Tile dimensions for main kernel loop constexpr int BLOCK_M = 128; // rows per thread block @@ -214,6 +214,15 @@ __device__ __forceinline__ uint16_t cvt_f32x4_to_fp4x4(float v0, float v1, float // Combine into 16-bit result (4 FP4 values) result |= (tmp << 8); return result; +#elif defined(__gfx1250__) + // gfx1250 (CDNA5) only exposes the PK8 scale converter (8 f32 -> 8 fp4 with a + // single per-lane float32 scale). Pad the upper 4 lanes with zeros and keep the + // low 16 bits (4 fp4). Element i maps to nibble i, and the scale divides the + // input -- identical semantics to the gfx950 pk_fp4 path. + typedef float float32x8_t __attribute__((ext_vector_type(8))); + const float32x8_t v = {v0, v1, v2, v3, 0.f, 0.f, 0.f, 0.f}; + const uint32_t packed = __builtin_amdgcn_cvt_scalef32_pk8_fp4_f32(v, scale); + return static_cast(packed & 0xFFFFu); #else __builtin_trap(); return 0; @@ -254,6 +263,13 @@ __device__ __forceinline__ uint16_t cvt_f32x4_to_fp4x4_sr(float v0, float v1, fl // Combine into 16-bit result (4 FP4 values) result |= (tmp << 8); return result; +#elif defined(__gfx1250__) + // gfx1250 stochastic-rounding PK8 converter: 8 f32 -> 8 fp4 with one rng seed + // and a per-lane float32 scale. Pad the upper half and keep the low 4 fp4. + typedef float float32x8_t __attribute__((ext_vector_type(8))); + const float32x8_t v = {v0, v1, v2, v3, 0.f, 0.f, 0.f, 0.f}; + const uint32_t packed = __builtin_amdgcn_cvt_scalef32_sr_pk8_fp4_f32(v, rng, scale); + return static_cast(packed & 0xFFFFu); #else __builtin_trap(); return 0; @@ -696,11 +712,17 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK, 4) void quantize_mxfp4_dual_kern // Layout: [N_chunk][column_within_chunk][m_chunk] __shared__ uint8_t s_colwise_scale[NUM_CHUNKS_N][MXFP4_BLOCK_SIZE][NUM_CHUNKS_M]; - // Zero-initialize for boundary handling (OOB entries stay 0) - static_assert(sizeof(s_colwise_fp4) == THREADS_PER_BLOCK * sizeof(uint64_t), - "s_colwise_fp4 size must match thread count for zero-init"); + // Zero-initialize for boundary handling (OOB entries stay 0). The block width + // scales with the wavefront size, so a grid-stride loop covers the buffer for + // both 256-thread (gfx950) and 128-thread (gfx1250) blocks. + constexpr int COLWISE_FP4_U64 = sizeof(s_colwise_fp4) / sizeof(uint64_t); + static_assert(sizeof(s_colwise_fp4) % sizeof(uint64_t) == 0, + "s_colwise_fp4 must be a whole number of uint64_t words"); if (!shuffle_colwise) { - reinterpret_cast(s_colwise_fp4)[tid] = 0; +#pragma unroll + for (int i = tid; i < COLWISE_FP4_U64; i += THREADS_PER_BLOCK) { + reinterpret_cast(s_colwise_fp4)[i] = 0; + } } // ======================================================================== @@ -1077,25 +1099,29 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK, 4) void quantize_mxfp4_dual_kern if (!shuffle_colwise) { constexpr int ITEMS_PER_COL = NUM_CHUNKS_M * THREADS_PER_ROW; constexpr int SEGS_PER_COL = ITEMS_PER_COL / 4; // uint64_t segments per column - static_assert(THREADS_PER_BLOCK == NUM_CHUNKS_N * MXFP4_BLOCK_SIZE * SEGS_PER_COL, - "Thread count must exactly cover all colwise FP4 segments"); + constexpr int TOTAL_SEGS = NUM_CHUNKS_N * MXFP4_BLOCK_SIZE * SEGS_PER_COL; + static_assert(TOTAL_SEGS % THREADS_PER_BLOCK == 0, + "Thread count must evenly cover all colwise FP4 segments"); - const int n_chunk = tid / (MXFP4_BLOCK_SIZE * SEGS_PER_COL); - const int local_tid = tid % (MXFP4_BLOCK_SIZE * SEGS_PER_COL); - const int col_in_chunk = local_tid / SEGS_PER_COL; - const int seg = local_tid % SEGS_PER_COL; +#pragma unroll + for (int item = tid; item < TOTAL_SEGS; item += THREADS_PER_BLOCK) { + const int n_chunk = item / (MXFP4_BLOCK_SIZE * SEGS_PER_COL); + const int local_tid = item % (MXFP4_BLOCK_SIZE * SEGS_PER_COL); + const int col_in_chunk = local_tid / SEGS_PER_COL; + const int seg = local_tid % SEGS_PER_COL; - const int global_col = base_n + n_chunk * MXFP4_BLOCK_SIZE + col_in_chunk; + const int global_col = base_n + n_chunk * MXFP4_BLOCK_SIZE + col_in_chunk; - if (global_col < N) { - const uint64_t data = *reinterpret_cast( - &s_colwise_fp4[n_chunk][col_in_chunk][seg * 4]); - const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); - if (row_start < M_pad) { - // Make sure the colwise store bypass L2 cache - __builtin_nontemporal_store( - data, reinterpret_cast(colwise_fp4 + global_col * M_packed + - base_m / 2 + seg * 8)); + if (global_col < N) { + const uint64_t data = *reinterpret_cast( + &s_colwise_fp4[n_chunk][col_in_chunk][seg * 4]); + const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); + if (row_start < M_pad) { + // Make sure the colwise store bypass L2 cache + __builtin_nontemporal_store( + data, reinterpret_cast(colwise_fp4 + global_col * M_packed + + base_m / 2 + seg * 8)); + } } } } @@ -1111,14 +1137,10 @@ void quantize_mxfp4_dual_impl(const DType *input, dtype::float4x2_e2m1 *rowwise_ int colwise_scale_M, int colwise_scale_N, int colwise_scale_M_pad, int colwise_scale_N_pad, ScalingRecipe rowwise_recipe, ScalingRecipe colwise_recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("quantize_mxfp4_dual is not implemented on gfx1250."); - } // Batched (G > 1) input is handled by replicating the per-matrix grid along // blockIdx.z; each z-slice quantizes one (M, N) group offset by its stride. dim3 grid((M_pad + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N, G); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); const uint32_t sr_seed = global_sr_counter.fetch_add(1, std::memory_order_relaxed); // Per-group strides into the contiguous (G, ...) output/scale buffers. FP4 @@ -1228,14 +1250,10 @@ void quantize_mxfp4_impl(const DType *input, dtype::float4x2_e2m1 *output, uint8 QuantizeMode mode, int G, int M, int N, int M_pad, int N_pad, int scale_stride, int scale_N, int scale_M_pad, int scale_N_pad, ScalingRecipe recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("quantize_mxfp4 is not implemented on gfx1250."); - } // Batched (G > 1) input replicates the per-matrix grid along blockIdx.z; // each z-slice quantizes one (M, N) group offset by its per-group stride. dim3 grid((M_pad + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N, G); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); const uint32_t sr_seed = global_sr_counter.fetch_add(1, std::memory_order_relaxed); // Per-group strides into the contiguous (G, ...) output/scale buffers. FP4 @@ -1375,9 +1393,13 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK, 4) void grouped_quantize_mxfp4_d s_colwise_fp4[NUM_CHUNKS_N][MXFP4_BLOCK_SIZE][NUM_CHUNKS_M * THREADS_PER_ROW]; __shared__ uint8_t s_colwise_scale[NUM_CHUNKS_N][MXFP4_BLOCK_SIZE][NUM_CHUNKS_M]; - static_assert(sizeof(s_colwise_fp4) == THREADS_PER_BLOCK * sizeof(uint64_t), - "s_colwise_fp4 size must match thread count for zero-init"); - reinterpret_cast(s_colwise_fp4)[tid] = 0; + constexpr int COLWISE_FP4_U64 = sizeof(s_colwise_fp4) / sizeof(uint64_t); + static_assert(sizeof(s_colwise_fp4) % sizeof(uint64_t) == 0, + "s_colwise_fp4 must be a whole number of uint64_t words"); +#pragma unroll + for (int i = tid; i < COLWISE_FP4_U64; i += THREADS_PER_BLOCK) { + reinterpret_cast(s_colwise_fp4)[i] = 0; + } for (int round = 0; round < TOTAL_CHUNKS; round += WARPS_PER_BLOCK) { const int chunk_idx = round + warp_id; @@ -1655,24 +1677,28 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK, 4) void grouped_quantize_mxfp4_d constexpr int ITEMS_PER_COL = NUM_CHUNKS_M * THREADS_PER_ROW; constexpr int SEGS_PER_COL = ITEMS_PER_COL / 4; // uint64_t segments per column - static_assert(THREADS_PER_BLOCK == NUM_CHUNKS_N * MXFP4_BLOCK_SIZE * SEGS_PER_COL, - "Thread count must exactly cover all colwise FP4 segments"); - - const int n_chunk = tid / (MXFP4_BLOCK_SIZE * SEGS_PER_COL); - const int local_tid = tid % (MXFP4_BLOCK_SIZE * SEGS_PER_COL); - const int col_in_chunk = local_tid / SEGS_PER_COL; - const int seg = local_tid % SEGS_PER_COL; - - const int global_col = base_n + n_chunk * MXFP4_BLOCK_SIZE + col_in_chunk; - if (global_col < N) { - const uint64_t data = - *reinterpret_cast(&s_colwise_fp4[n_chunk][col_in_chunk][seg * 4]); - const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); - if (row_start < M_pad_col) { - __builtin_nontemporal_store( - data, reinterpret_cast(colwise_fp4 + - static_cast(global_col) * M_packed + - base_m / 2 + seg * 8)); + constexpr int TOTAL_SEGS = NUM_CHUNKS_N * MXFP4_BLOCK_SIZE * SEGS_PER_COL; + static_assert(TOTAL_SEGS % THREADS_PER_BLOCK == 0, + "Thread count must evenly cover all colwise FP4 segments"); + +#pragma unroll + for (int item = tid; item < TOTAL_SEGS; item += THREADS_PER_BLOCK) { + const int n_chunk = item / (MXFP4_BLOCK_SIZE * SEGS_PER_COL); + const int local_tid = item % (MXFP4_BLOCK_SIZE * SEGS_PER_COL); + const int col_in_chunk = local_tid / SEGS_PER_COL; + const int seg = local_tid % SEGS_PER_COL; + + const int global_col = base_n + n_chunk * MXFP4_BLOCK_SIZE + col_in_chunk; + if (global_col < N) { + const uint64_t data = *reinterpret_cast( + &s_colwise_fp4[n_chunk][col_in_chunk][seg * 4]); + const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); + if (row_start < M_pad_col) { + __builtin_nontemporal_store( + data, reinterpret_cast( + colwise_fp4 + static_cast(global_col) * M_packed + + base_m / 2 + seg * 8)); + } } } } @@ -1687,17 +1713,13 @@ void grouped_quantize_mxfp4_dual_impl(const DType *input, dtype::float4x2_e2m1 * int colwise_scale_stride, int rowwise_scale_N, int colwise_scale_N, ScalingRecipe rowwise_recipe, ScalingRecipe colwise_recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("grouped_quantize_mxfp4_dual is not implemented on gfx1250."); - } PRIMUS_TURBO_CHECK(rowwise_recipe.shuffle_out == false && rowwise_recipe.shuffle_scale == false, "grouped MXFP4 dual does not support shuffle"); PRIMUS_TURBO_CHECK(colwise_recipe.shuffle_out == false && colwise_recipe.shuffle_scale == false, "grouped MXFP4 dual does not support shuffle"); dim3 grid((M_pad_col + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); const uint32_t sr_seed = global_sr_counter.fetch_add(1, std::memory_order_relaxed); #define GROUPED_QUANTIZE_MXFP4_DUAL_ARGS \ @@ -2031,15 +2053,11 @@ void grouped_quantize_mxfp4_impl(const DType *input, dtype::float4x2_e2m1 *outpu const int64_t *group_offs_padded_colwise, QuantizeMode mode, int G, int total_M, int N, int M_pad_col, int N_pad, int scale_stride, int scale_N, ScalingRecipe recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("grouped_quantize_mxfp4 is not implemented on gfx1250."); - } PRIMUS_TURBO_CHECK(recipe.shuffle_out == false && recipe.shuffle_scale == false, "grouped MXFP4 single does not support shuffle"); dim3 grid((M_pad_col + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); const uint32_t sr_seed = global_sr_counter.fetch_add(1, std::memory_order_relaxed); #define GROUPED_QUANTIZE_MXFP4_ARGS \ diff --git a/csrc/kernels/quantization/quantization_mxfp8.cu b/csrc/kernels/quantization/quantization_mxfp8.cu index b8fa6e0a9..e63ff2699 100644 --- a/csrc/kernels/quantization/quantization_mxfp8.cu +++ b/csrc/kernels/quantization/quantization_mxfp8.cu @@ -37,9 +37,9 @@ using namespace primus_turbo::detail; // ============================================================================ // Hardware architecture parameters -constexpr int WARP_SIZE = 64; // AMD wavefront size -constexpr int THREADS_PER_BLOCK = 256; // 4 warps per block -constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE; +constexpr int WARP_SIZE = THREADS_PER_WARP; +constexpr int WARPS_PER_BLOCK = 4; +constexpr int THREADS_PER_BLOCK = WARP_SIZE * WARPS_PER_BLOCK; // Tile dimensions for main kernel loop constexpr int BLOCK_M = 128; // rows per thread block @@ -174,6 +174,42 @@ __device__ __forceinline__ uint32_t cvt_f32x4_to_fp8x4(float v0, float v1, float } return result; +#elif defined(__gfx1250__) + // gfx1250 (CDNA5) only exposes the PK8 scale converters (8 f32 -> 8 fp8/bf8 with + // a single per-lane float32 scale). Pad the upper 4 lanes with zeros and keep the + // low 32 bits (4 fp8). Element i maps to byte i and the scale divides the input -- + // identical semantics to the gfx950 pk_fp8 / pk_bf8 path. Soft-clamp first to + // preserve the gfx950 behaviour of avoiding NaN on overflow. + // The PK8 fp8/bf8 converters return a uint2 (8 packed fp8 bytes); the low dword + // holds our 4 useful values (elements 0..3). + typedef float float32x8_t __attribute__((ext_vector_type(8))); + typedef uint32_t uint32x2_t __attribute__((ext_vector_type(2))); + if constexpr (std::is_same_v) { + const float lim = FP8E4M3_MAX * scale; + const float32x8_t v = {fminf(fmaxf(v0, -lim), lim), + fminf(fmaxf(v1, -lim), lim), + fminf(fmaxf(v2, -lim), lim), + fminf(fmaxf(v3, -lim), lim), + 0.f, + 0.f, + 0.f, + 0.f}; + const uint32x2_t packed = __builtin_amdgcn_cvt_scalef32_pk8_fp8_f32(v, scale); + return packed[0]; + } else if constexpr (std::is_same_v) { + const float lim = FP8E5M2_MAX * scale; + const float32x8_t v = {fminf(fmaxf(v0, -lim), lim), + fminf(fmaxf(v1, -lim), lim), + fminf(fmaxf(v2, -lim), lim), + fminf(fmaxf(v3, -lim), lim), + 0.f, + 0.f, + 0.f, + 0.f}; + const uint32x2_t packed = __builtin_amdgcn_cvt_scalef32_pk8_bf8_f32(v, scale); + return packed[0]; + } + return 0; #else __builtin_trap(); return 0; @@ -898,24 +934,28 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK, 4) void quantize_mxfp8_dual_kern if (!shuffle_colwise) { constexpr int ITEMS_PER_COL = NUM_CHUNKS_M * THREADS_PER_ROW; constexpr int SEGS_PER_COL = ITEMS_PER_COL / 4; // uint4 segments per column + constexpr int TOTAL_SEGS = NUM_CHUNKS_N * MXFP8_BLOCK_SIZE * SEGS_PER_COL; static_assert(ITEMS_PER_COL % 4 == 0, "ITEMS_PER_COL must be divisible by 4"); - static_assert(THREADS_PER_BLOCK == NUM_CHUNKS_N * MXFP8_BLOCK_SIZE * SEGS_PER_COL, - "Thread count must exactly cover all colwise FP8 segments"); - - const int n_chunk = tid / (MXFP8_BLOCK_SIZE * SEGS_PER_COL); - const int local_tid = tid % (MXFP8_BLOCK_SIZE * SEGS_PER_COL); - const int col_in_chunk = local_tid / SEGS_PER_COL; - const int seg = local_tid % SEGS_PER_COL; - - const int global_col = base_n + n_chunk * MXFP8_BLOCK_SIZE + col_in_chunk; - if (global_col < N) { - const uint4 data = *reinterpret_cast( - &s_colwise_fp8[n_chunk][col_in_chunk][seg * 4]); - const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); - if (row_start < M_pad) { - *reinterpret_cast( - colwise_fp8 + static_cast(global_col) * M_packed + row_start) = - data; + static_assert(TOTAL_SEGS % THREADS_PER_BLOCK == 0, + "Thread count must evenly cover all colwise FP8 segments"); + +#pragma unroll + for (int item = tid; item < TOTAL_SEGS; item += THREADS_PER_BLOCK) { + const int n_chunk = item / (MXFP8_BLOCK_SIZE * SEGS_PER_COL); + const int local_tid = item % (MXFP8_BLOCK_SIZE * SEGS_PER_COL); + const int col_in_chunk = local_tid / SEGS_PER_COL; + const int seg = local_tid % SEGS_PER_COL; + + const int global_col = base_n + n_chunk * MXFP8_BLOCK_SIZE + col_in_chunk; + if (global_col < N) { + const uint4 data = *reinterpret_cast( + &s_colwise_fp8[n_chunk][col_in_chunk][seg * 4]); + const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); + if (row_start < M_pad) { + *reinterpret_cast( + colwise_fp8 + static_cast(global_col) * M_packed + row_start) = + data; + } } } } @@ -1643,24 +1683,28 @@ __global__ __launch_bounds__(THREADS_PER_BLOCK, 4) void grouped_quantize_mxfp8_d if (!shuffle_colwise) { constexpr int ITEMS_PER_COL = NUM_CHUNKS_M * THREADS_PER_ROW; constexpr int SEGS_PER_COL = ITEMS_PER_COL / 4; + constexpr int TOTAL_SEGS = NUM_CHUNKS_N * MXFP8_BLOCK_SIZE * SEGS_PER_COL; static_assert(ITEMS_PER_COL % 4 == 0, "ITEMS_PER_COL must be divisible by 4"); - static_assert(THREADS_PER_BLOCK == NUM_CHUNKS_N * MXFP8_BLOCK_SIZE * SEGS_PER_COL, - "Thread count must exactly cover all colwise FP8 segments"); - - const int n_chunk = tid / (MXFP8_BLOCK_SIZE * SEGS_PER_COL); - const int local_tid = tid % (MXFP8_BLOCK_SIZE * SEGS_PER_COL); - const int col_in_chunk = local_tid / SEGS_PER_COL; - const int seg = local_tid % SEGS_PER_COL; - - const int global_col = base_n + n_chunk * MXFP8_BLOCK_SIZE + col_in_chunk; - if (global_col < N) { - const uint4 data = *reinterpret_cast( - &s_colwise_fp8[n_chunk][col_in_chunk][seg * 4]); - const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); - if (row_start < M_pad) { - *reinterpret_cast( - colwise_fp8 + static_cast(global_col) * M_packed + row_start) = - data; + static_assert(TOTAL_SEGS % THREADS_PER_BLOCK == 0, + "Thread count must evenly cover all colwise FP8 segments"); + +#pragma unroll + for (int item = tid; item < TOTAL_SEGS; item += THREADS_PER_BLOCK) { + const int n_chunk = item / (MXFP8_BLOCK_SIZE * SEGS_PER_COL); + const int local_tid = item % (MXFP8_BLOCK_SIZE * SEGS_PER_COL); + const int col_in_chunk = local_tid / SEGS_PER_COL; + const int seg = local_tid % SEGS_PER_COL; + + const int global_col = base_n + n_chunk * MXFP8_BLOCK_SIZE + col_in_chunk; + if (global_col < N) { + const uint4 data = *reinterpret_cast( + &s_colwise_fp8[n_chunk][col_in_chunk][seg * 4]); + const int row_start = base_m + seg * (4 * ELEMS_PER_THREAD); + if (row_start < M_pad) { + *reinterpret_cast( + colwise_fp8 + static_cast(global_col) * M_packed + row_start) = + data; + } } } } @@ -1676,10 +1720,6 @@ void grouped_quantize_mxfp8_dual_impl( int rowwise_scale_M_pad, int rowwise_scale_N_pad, int colwise_scale_M, int colwise_scale_N, int colwise_scale_M_pad, int colwise_scale_N_pad, ScalingRecipe rowwise_recipe, ScalingRecipe colwise_recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("grouped_quantize_mxfp8_dual is not implemented on gfx1250."); - } PRIMUS_TURBO_CHECK(rowwise_recipe.use_rht == false, "MXFP8 not support RHT"); PRIMUS_TURBO_CHECK(colwise_recipe.use_rht == false, "MXFP8 not support RHT"); PRIMUS_TURBO_CHECK(rowwise_recipe.use_sr == false, "MXFP8 not support SR"); @@ -1690,7 +1730,7 @@ void grouped_quantize_mxfp8_dual_impl( const int total_M_padded = (total_M + G * BLOCK_M + BLOCK_M - 1) / BLOCK_M * BLOCK_M; dim3 grid((total_M_padded + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); #define QUANTIZE_MXFP8_DUAL_GROUPED \ input, rowwise_output, rowwise_scale, colwise_output, colwise_scale, group_offs, \ @@ -1761,17 +1801,13 @@ void grouped_quantize_mxfp8_impl(const IType *input, OType *output, uint8_t *sca QuantizeMode mode, int G, int total_M, int N, int N_pad, int scale_stride, int scale_N, int scale_M_pad, int scale_N_pad, ScalingRecipe recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("grouped_quantize_mxfp8 is not implemented on gfx1250."); - } PRIMUS_TURBO_CHECK(recipe.use_rht == false, "MXFP8 not support RHT"); PRIMUS_TURBO_CHECK(recipe.use_sr == false, "MXFP8 not support SR"); const int total_M_padded = (total_M + G * BLOCK_M + BLOCK_M - 1) / BLOCK_M * BLOCK_M; dim3 grid((total_M_padded + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); #define QUANTIZE_MXFP8_GROUPED_ARGS \ input, output, scale, group_offs, group_offs_padded, G, N, total_M_padded, N_pad, \ @@ -1826,14 +1862,10 @@ void quantize_mxfp8_dual_impl(const IType *input, OType *rowwise_output, uint8_t int colwise_scale_N, int colwise_scale_M_pad, int colwise_scale_N_pad, ScalingRecipe rowwise_recipe, ScalingRecipe colwise_recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("quantize_mxfp8_dual is not implemented on gfx1250."); - } // Batched (G > 1) input is handled by replicating the per-matrix grid along // blockIdx.z; each z-slice quantizes one (M, N) group offset by its stride. dim3 grid((M_pad + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N, G); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); PRIMUS_TURBO_CHECK(rowwise_recipe.use_rht == false, "MXFP8 not support RHT"); PRIMUS_TURBO_CHECK(colwise_recipe.use_rht == false, "MXFP8 not support RHT"); @@ -1921,14 +1953,10 @@ void quantize_mxfp8_impl(const IType *input, OType *output, uint8_t *scale, Quan int G, int M, int N, int M_pad, int N_pad, int scale_stride, int scale_N, int scale_M_pad, int scale_N_pad, ScalingRecipe recipe, hipStream_t stream) { - // TODO(ruibin): add kernel support for gfx1250 - if (is_gfx1250()) { - PRIMUS_TURBO_ERROR("quantize_mxfp8 is not implemented on gfx1250."); - } // Batched (G > 1) input replicates the per-matrix grid along blockIdx.z; // each z-slice quantizes one (M, N) group offset by its per-group stride. dim3 grid((M_pad + BLOCK_M - 1) / BLOCK_M, (N_pad + BLOCK_N - 1) / BLOCK_N, G); - dim3 block(THREADS_PER_BLOCK); + dim3 block(warp_size() * WARPS_PER_BLOCK); PRIMUS_TURBO_CHECK(recipe.use_rht == false, "MXFP8 not support RHT"); PRIMUS_TURBO_CHECK(recipe.use_sr == false, "MXFP8 not support SR"); diff --git a/csrc/pytorch/gemm/turbo_gemm.cpp b/csrc/pytorch/gemm/turbo_gemm.cpp index 8c2708eac..b91c8b759 100644 --- a/csrc/pytorch/gemm/turbo_gemm.cpp +++ b/csrc/pytorch/gemm/turbo_gemm.cpp @@ -67,9 +67,6 @@ at::Tensor turbo_gemm_fp8(at::Tensor A, at::Tensor scaleA_inv, at::Tensor B, at: const at::ScalarType out_dtype, bool transA, bool transB, bool transC, const std::string &granularity) { // Runtime guard: the turbo (MFMA) backend is not compiled on gfx1250. - PRIMUS_TURBO_CHECK(!primus_turbo::is_gfx1250(), - "turbo_gemm_fp8 is unavailable: the turbo backend is not built " - "on this architecture (gfx1250 is unsupported)."); PRIMUS_TURBO_CHECK(is_8bit_floating_point_dtype(A.scalar_type()), "A must be FP8."); PRIMUS_TURBO_CHECK(is_8bit_floating_point_dtype(B.scalar_type()), "B must be FP8."); PRIMUS_TURBO_CHECK(is_16bit_floating_point_dtype(out_dtype) || out_dtype == at::kFloat,