feat: support gfx1250 on turbo (Part1)

Author: RuibinCheung

.gitignore

@@ -15,3 +15,4 @@ agent/workspace/
 logs/
 primus_turbo/_build_info.py
 primus_turbo/_version.py
+primus_turbo/common/build_config.py

csrc/include/primus_turbo/arch.h

@@ -10,7 +10,7 @@
 
 namespace primus_turbo {
 
-enum class GPUArch { GFX942, GFX950, UNKNOWN };
+enum class GPUArch { GFX942, GFX950, GFX1250, UNKNOWN };
 
 inline GPUArch get_current_arch() {
     static GPUArch cached_arch = []() -> GPUArch {
@@ -23,6 +23,8 @@ inline GPUArch get_current_arch() {
             return GPUArch::GFX942;
         if (prop.major == 9 && prop.minor == 5)
             return GPUArch::GFX950;
+        if (prop.major == 12 && prop.minor == 5)
+            return GPUArch::GFX1250;
         return GPUArch::UNKNOWN;
     }();
     return cached_arch;
@@ -36,6 +38,22 @@ inline bool is_gfx942() {
     return get_current_arch() == GPUArch::GFX942;
 }
 
+inline bool is_gfx1250() {
+    return get_current_arch() == GPUArch::GFX1250;
+}
+
+// gfx1250 = 32, gfx942 / gfx950 (and others) = 64.
+inline int warp_size() {
+    if (is_gfx1250()) {
+        return 32;
+    } else if (is_gfx950() || is_gfx942()) {
+        return 64;
+    } else {
+        PRIMUS_TURBO_ERROR("Unknown architecture");
+        return -1;
+    }
+}
+
 inline int32_t get_multi_processor_count(const int32_t device_id) {
     int32_t num_cu = 0;
     PRIMUS_TURBO_CHECK_HIP(

csrc/include/primus_turbo/device/memory.cuh

@@ -16,11 +16,19 @@ namespace primus_turbo::device {
 // ════════════════════════════════════════════════════════════════
 
 template <int CNT> __device__ __forceinline__ void wait_lgkmcnt() {
+#if defined(__gfx1250__)
+    asm volatile("s_wait_dscnt %0" : : "n"(CNT) : "memory");
+#else
     asm volatile("s_waitcnt lgkmcnt(%0)" : : "n"(CNT) : "memory");
+#endif
 }
 
 template <int CNT> __device__ __forceinline__ void wait_vmcnt() {
+#if defined(__gfx1250__)
+    asm volatile("s_wait_loadcnt %0" : : "n"(CNT) : "memory");
+#else
     asm volatile("s_waitcnt vmcnt(%0)" : : "n"(CNT) : "memory");
+#endif
 }
 
 // ════════════════════════════════════════════════════════════════

csrc/include/primus_turbo/device/reduce.cuh

@@ -7,6 +7,14 @@
 #include "primus_turbo/common.h"
 #include "primus_turbo/device/utils.cuh"
 
+// gfx12 (gfx1250 / gfx1251, CDNA5) removed the unified `s_waitcnt`; LDS waits use
+// the split `s_wait_dscnt` counter instead.
+#if defined(__gfx1250__)
+#define PRIMUS_TURBO_WAIT_DS_STR "s_wait_dscnt 0"
+#else
+#define PRIMUS_TURBO_WAIT_DS_STR "s_waitcnt lgkmcnt(0)"
+#endif
+
 #define FINAL_MASK 0xffffffffffffffffULL
 
 namespace primus_turbo {
@@ -187,35 +195,31 @@ template <template <class> class Func, typename T> PRIMUS_TURBO_DEVICE T BlockRe
 //   block memory-op scheduling around this VALU-bound reduction.
 PRIMUS_TURBO_DEVICE float ds_swizzle_xor1(float val) {
     float result;
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x041F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x041F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(result)
                  : "v"(val));
     return result;
 }
 
 PRIMUS_TURBO_DEVICE float ds_swizzle_xor2(float val) {
     float result;
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x081F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x081F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(result)
                  : "v"(val));
     return result;
 }
 
 PRIMUS_TURBO_DEVICE float ds_swizzle_xor8(float val) {
     float result;
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x201F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x201F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(result)
                  : "v"(val));
     return result;
 }
 
 PRIMUS_TURBO_DEVICE float ds_swizzle_xor16(float val) {
     float result;
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x401F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x401F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(result)
                  : "v"(val));
     return result;
@@ -247,36 +251,38 @@ PRIMUS_TURBO_DEVICE float warp_reduce_max_8_dpp(float val) {
     // clobber) is what makes the per-row reduction deterministic.
 
     // Step 1: Exchange with thread 4 positions away
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x101F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x101F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(tmp)
                  : "v"(v));
     val = fmaxf(val, uint_as_float(tmp));
     v   = float_as_uint(val);
 
     // Step 2: Exchange with thread 2 positions away
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x081F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x081F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(tmp)
                  : "v"(v));
     val = fmaxf(val, uint_as_float(tmp));
     v   = float_as_uint(val);
 
     // Step 3: Exchange with adjacent thread
-    asm volatile("ds_swizzle_b32 %0, %1 offset:0x041F\n\t"
-                 "s_waitcnt lgkmcnt(0)"
+    asm volatile("ds_swizzle_b32 %0, %1 offset:0x041F\n\t" PRIMUS_TURBO_WAIT_DS_STR
                  : "=v"(tmp)
                  : "v"(v));
     val = fmaxf(val, uint_as_float(tmp));
 
     return val;
 }
 
+// Full-wavefront max reduction. ``ds_swizzle`` only exchanges within 32-lane
+// groups, so a 64-lane wavefront needs an extra cross-group (lane ^ 32) step
+// while a 32-lane wavefront (gfx1250) is already fully reduced after xor16.
 PRIMUS_TURBO_DEVICE float warp_reduce_max_64_dpp(float val) {
     val = warp_reduce_max_8_dpp(val);
     val = fmaxf(val, ds_swizzle_xor8(val));
     val = fmaxf(val, ds_swizzle_xor16(val));
-    val = fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 32, THREADS_PER_WARP));
+    if constexpr (THREADS_PER_WARP == 64) {
+        val = fmaxf(val, __shfl_xor_sync(FINAL_MASK, val, 32, THREADS_PER_WARP));
+    }
     return val;
 }
 

csrc/include/primus_turbo/float8.h

@@ -25,17 +25,17 @@ inline Float8Format current_fp8_format() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
     return Float8Format::FNUZ; // dummy
 #else
-    static Float8Format fmt = [] { return is_gfx950() ? Float8Format::OCP : Float8Format::FNUZ; }();
+    static Float8Format fmt = [] { return is_gfx942() ? Float8Format::FNUZ : Float8Format::OCP; }();
     return fmt;
 #endif
 }
 
 PRIMUS_TURBO_HOST_DEVICE bool is_fp8_fnuz() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-    return false; // gfx950 OCP
-#else
+#if defined(__gfx942__)
     return true;
+#else
+    return false; // gfx950 OCP
 #endif
 #else
     return current_fp8_format() == Float8Format::FNUZ;
@@ -45,10 +45,10 @@ PRIMUS_TURBO_HOST_DEVICE bool is_fp8_fnuz() {
 struct float8_e4m3_t {
 
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-    using storage_t = __hip_fp8_e4m3; // OCP on gfx950
+#if defined(__gfx942__)
+    using storage_t = __hip_fp8_e4m3_fnuz; // FNUZ on gfx942
 #else
-    using storage_t = __hip_fp8_e4m3_fnuz; // FNUZ on others
+    using storage_t = __hip_fp8_e4m3; // OCP on others
 #endif
     storage_t val;
 #else // host side – keep both encodings
@@ -180,10 +180,10 @@ static_assert(std::is_trivially_copyable_v<float8_e4m3_t>);
 struct float8_e5m2_t {
 
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-    using storage_t = __hip_fp8_e5m2; // OCP on gfx950
+#if defined(__gfx942__)
+    using storage_t = __hip_fp8_e5m2_fnuz; // FNUZ on gfx942
 #else
-    using storage_t = __hip_fp8_e5m2_fnuz; // FNUZ on others
+    using storage_t = __hip_fp8_e5m2; // OCP on others
 #endif
     storage_t val;
 #else // host side – keep both encodings
@@ -392,10 +392,10 @@ template <> class numeric_limits<float8_e4m3_t> {
 
     PRIMUS_TURBO_HOST_DEVICE static float8_e4m3_t max() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-        return float8_e4m3_t::from_bits(0x7E);
-#else
+#if defined(__gfx942__)
         return float8_e4m3_t::from_bits(0x7F);
+#else
+        return float8_e4m3_t::from_bits(0x7E);
 #endif
 #else
         return primus_turbo::is_fp8_fnuz() ? float8_e4m3_t::from_bits(0x7F)
@@ -414,10 +414,10 @@ template <> class numeric_limits<float8_e4m3_t> {
     // FNUZ: 0x80
     PRIMUS_TURBO_HOST_DEVICE static float8_e4m3_t quiet_NaN() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-        return float8_e4m3_t::from_bits(0x7F);
-#else
+#if defined(__gfx942__)
         return float8_e4m3_t::from_bits(0x80);
+#else
+        return float8_e4m3_t::from_bits(0x7F);
 #endif
 #else
         return primus_turbo::is_fp8_fnuz() ? float8_e4m3_t::from_bits(0x80)
@@ -430,10 +430,10 @@ template <> class numeric_limits<float8_e5m2_t> {
 public:
     static constexpr bool is_specialized = true;
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-    static constexpr bool has_infinity = true;
-#else
+#if defined(__gfx942__)
     static constexpr bool has_infinity = false;
+#else
+    static constexpr bool has_infinity = true;
 #endif
 #else
     // Host: cannot determine at compile time whether this is OCP or FNUZ.
@@ -448,10 +448,10 @@ template <> class numeric_limits<float8_e5m2_t> {
 
     PRIMUS_TURBO_HOST_DEVICE static float8_e5m2_t max() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-        return float8_e5m2_t::from_bits(0x7B);
-#else
+#if defined(__gfx942__)
         return float8_e5m2_t::from_bits(0x7F);
+#else
+        return float8_e5m2_t::from_bits(0x7B);
 #endif
 #else
         return primus_turbo::is_fp8_fnuz() ? float8_e5m2_t::from_bits(0x7F)
@@ -463,10 +463,10 @@ template <> class numeric_limits<float8_e5m2_t> {
 
     PRIMUS_TURBO_HOST_DEVICE static float8_e5m2_t infinity() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-        return float8_e5m2_t::from_bits(0x7C);
-#else
+#if defined(__gfx942__)
         return max();
+#else
+        return float8_e5m2_t::from_bits(0x7C);
 #endif
 #else
         return primus_turbo::is_fp8_fnuz() ? max() : float8_e5m2_t::from_bits(0x7C);
@@ -475,10 +475,10 @@ template <> class numeric_limits<float8_e5m2_t> {
 
     PRIMUS_TURBO_HOST_DEVICE static float8_e5m2_t quiet_NaN() {
 #if PRIMUS_TURBO_DEVICE_COMPILE
-#if defined(__gfx950__)
-        return float8_e5m2_t::from_bits(0x7D);
-#else
+#if defined(__gfx942__)
         return float8_e5m2_t::from_bits(0x80);
+#else
+        return float8_e5m2_t::from_bits(0x7D);
 #endif
 #else
         return primus_turbo::is_fp8_fnuz() ? float8_e5m2_t::from_bits(0x80)

csrc/include/primus_turbo/gemm.h

@@ -4,7 +4,9 @@
 
 #pragma once
 
+#ifdef BUILD_CK_BACKEND
 #include "ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp"
+#endif
 #include "primus_turbo/dtype.h"
 #include <cstdint>
 #include <hip/hip_runtime.h>
@@ -56,6 +58,7 @@ void hipblaslt_gemm_impl(const void *A, const hipDataType A_type, const int64_t
 //==================================================================
 //  CK GEMM
 //==================================================================
+#ifdef BUILD_CK_BACKEND
 
 template <typename AType, typename BType, typename CType, typename ACCType = float>
 struct CKGemmFP8Params {
@@ -79,9 +82,12 @@ template <typename ADataType, typename BDataType, typename CDataType, typename A
           ck_tile::QuantType QuantMode>
 void ck_gemm_fp8_impl(const CKGemmFP8Params<ADataType, BDataType, CDataType, AccDataType> &params);
 
+#endif // BUILD_CK_BACKEND
+
 //==================================================================
 //  Turbo GEMM
 //==================================================================
+#ifdef BUILD_TURBO_BACKEND
 
 size_t turbo_gemm_mxfp8_workspace_size(int32_t m, int32_t n, int32_t k);
 
@@ -92,4 +98,6 @@ void turbo_gemm_mxfp8_impl(const AType *a_ptr, const BType *b_ptr,
                            int32_t n, int32_t k, void *workspace, size_t workspace_size,
                            hipStream_t stream);
 
+#endif // BUILD_TURBO_BACKEND
+
 } // namespace primus_turbo

csrc/include/primus_turbo/grouped_gemm.h

@@ -4,16 +4,20 @@
 
 #pragma once
 
+#ifdef BUILD_CK_BACKEND
 #include "ck_tile/ops/gemm_quant/pipeline/tile_gemm_quant_traits.hpp"
+#endif
 #include <cstdint>
 #include <hip/hip_runtime.h>
 
 #include "primus_turbo/common.h"
 
 namespace primus_turbo {
 
+#ifdef BUILD_CK_BACKEND
 std::int64_t get_ck_grouped_gemm_args_sizes(const int group_num);
 std::int64_t get_ck_grouped_gemm_fp8_args_sizes(const int group_num);
+#endif
 
 std::int64_t get_hipblaslt_grouped_gemm_workspace_size();
 
@@ -56,6 +60,7 @@ template <typename AType, typename BType, typename CType> struct GroupedGemmPara
     int32_t  ws_local_per_xcd = 0;
 };
 
+#ifdef BUILD_CK_BACKEND
 template <typename AType, typename BType, typename CType>
 struct CKGroupedGemmParams : public GroupedGemmParams<AType, BType, CType> {
     void *args_ptr = nullptr;
@@ -66,6 +71,7 @@ struct CKGroupedGemmFP8Params : public CKGroupedGemmParams<AType, BType, CType>
     const ACCType *aq_ptr = nullptr;
     const ACCType *bq_ptr = nullptr;
 };
+#endif // BUILD_CK_BACKEND
 
 struct HipblasltGroupedGemmParams {
     const void          *a_ptr       = nullptr;
@@ -99,6 +105,7 @@ struct HipblasltGroupedGemmParams {
 //==================================================================
 //  CK Grouped GEMM
 //==================================================================
+#ifdef BUILD_CK_BACKEND
 
 template <typename ADataType, typename BDataType, typename CDataType, typename AccDataType = float>
 void ck_grouped_gemm(const CKGroupedGemmParams<ADataType, BDataType, CDataType> &params);
@@ -116,6 +123,8 @@ template <typename ADataType, typename BDataType, typename CDataType, typename A
 void ck_grouped_gemm_fp8_variable_k(
     const CKGroupedGemmFP8Params<ADataType, BDataType, CDataType, AccDataType> &params);
 
+#endif // BUILD_CK_BACKEND
+
 //==================================================================
 //  hipBLASLt Grouped GEMM
 //==================================================================