[Perf] [AMD] minimax: decode indexer CP for DSpark verify on the NVFP4 index cache

python/sglang/kernels/jit/csrc/distributed/quick_all_reduce/quick_all_reduce.h

@@ -0,0 +1,183 @@
+// JIT copy of sgl-kernel's kernels/aot/csrc/allreduce/quick_all_reduce.h: the same DeviceComms and
+// buffers, launching the INT4 / INT8 TP4 two-shot of quick_all_reduce.cuh.
+#pragma once
+
+#include <hip/hip_runtime.h>
+
+#include "quick_all_reduce.cuh"
+#include <algorithm>
+#include <cstdio>
+#include <limits>
+#include <optional>
+#include <stdexcept>
+#include <vector>
+
+#define HIP_CHECK(err)                                                                               \
+  do {                                                                                               \
+    hipError_t err_ = (err);                                                                         \
+    if (err_ != hipSuccess) {                                                                        \
+      std::printf("HIP error %d at %s:%d. %s\n", err_, __FILE__, __LINE__, hipGetErrorString(err_)); \
+      throw std::runtime_error("HIP error");                                                         \
+    }                                                                                                \
+  } while (0)
+
+namespace quickreduce {
+
+template <typename AllReduceKernel, typename T>
+__global__ __quickreduce_launch_bounds_two_shot__ static void allreduce_prototype_twoshot(
+    T const* A,
+    T* B,
+    uint32_t N,
+    uint32_t num_blocks,
+    int rank,
+    uint8_t** dbuffer_list,
+    uint32_t data_offset,
+    uint32_t* d_flag_counters,
+    int64_t data_size_per_phase) {
+  int block = blockIdx.x;
+  int grid = gridDim.x;
+
+  // Kept in device memory, not a launch argument, so every CUDA-graph replay sees a fresh color.
+  uint32_t flag_color = d_flag_counters[blockIdx.x];
+
+  // Loaded once, before any remote store: on CDNA one vmcnt covers loads and stores, so a pointer
+  // load between remote stores would wait for their acknowledgment.
+  uint8_t* peer[AllReduceKernel::kWorldSize];
+  for (int r = 0; r < AllReduceKernel::kWorldSize; r++) {
+    peer[r] = dbuffer_list[r];
+  }
+  while (block < num_blocks) {
+    AllReduceKernel::run(A, B, N, block, rank, peer, data_offset, flag_color, data_size_per_phase);
+    block += grid;
+    flag_color++;
+  }
+  // The whole block ends up with the same value, so a single writer suffices.
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    d_flag_counters[blockIdx.x] = flag_color;
+  }
+}
+
+struct DeviceComms {
+  // Max problem size is 2GB (in bytes) or half of uint32_t max value.
+  int64_t kMaxProblemSize = static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1;
+
+  bool initialized = false;
+  uint32_t* d_flag_counters = nullptr;
+  int world_size;
+  int rank;
+
+  uint8_t* dbuffer;
+  uint8_t** dbuffer_list;
+  hipIpcMemHandle_t buffer_ipc_handle;
+  std::vector<hipIpcMemHandle_t> all_buffer_ipc_handles;
+  std::vector<uint8_t*> buffer_list;
+  uint32_t data_offset;
+
+  DeviceComms() : initialized(false), world_size(1), rank(0) {}
+  ~DeviceComms() {
+    destroy();
+  }
+
+  void init(int world_size, int rank, std::optional<int64_t> max_problem_size = std::nullopt) {
+    destroy();
+    this->world_size = world_size;
+    this->rank = rank;
+    if (max_problem_size.has_value() && max_problem_size.value() > 0) {
+      this->kMaxProblemSize = max_problem_size.value();
+    }
+    // Allocate buffer size for worst case: F16 2-stage buffer.
+    uint32_t flags_buffer_size = 2 * world_size * kMaxNumBlocks * sizeof(uint32_t);
+    int64_t data_buffer_size = 2 * this->kMaxProblemSize;
+    int64_t total_buffer_size = flags_buffer_size + data_buffer_size;
+    data_offset = flags_buffer_size;
+    HIP_CHECK(hipExtMallocWithFlags((void**)&dbuffer, total_buffer_size, hipDeviceMallocUncached));
+
+    HIP_CHECK(hipMemset(dbuffer, 0, flags_buffer_size));
+
+    // Seeded with 1 so the first color never matches the zeroed flags.
+    HIP_CHECK(hipMalloc(&d_flag_counters, kMaxNumBlocks * sizeof(uint32_t)));
+    {
+      std::vector<uint32_t> init_color(kMaxNumBlocks, 1u);
+      HIP_CHECK(hipMemcpy(d_flag_counters, init_color.data(), kMaxNumBlocks * sizeof(uint32_t), hipMemcpyHostToDevice));
+    }
+
+    buffer_list.assign(world_size, nullptr);
+    HIP_CHECK(hipMalloc(&dbuffer_list, world_size * sizeof(uint8_t*)));
+
+    all_buffer_ipc_handles.resize(world_size);
+    HIP_CHECK(hipIpcGetMemHandle(&buffer_ipc_handle, dbuffer));
+
+    initialized = true;
+  }
+  int get_world_size() {
+    return world_size;
+  }
+  hipIpcMemHandle_t const get_handle() {
+    return buffer_ipc_handle;
+  }
+
+  void destroy() {
+    // Allocated before initialized is set, so freed on its own check after a failed init.
+    if (d_flag_counters) {
+      HIP_CHECK(hipFree(d_flag_counters));
+      d_flag_counters = nullptr;
+    }
+    if (initialized) {
+      // The host copy: dbuffer_list is device memory, and a peer may never have been opened.
+      for (int i = 0; i < world_size; i++) {
+        if (i != rank && buffer_list[i] != nullptr) {
+          HIP_CHECK(hipIpcCloseMemHandle(buffer_list[i]));
+        }
+      }
+
+      HIP_CHECK(hipFree(dbuffer));
+      HIP_CHECK(hipFree(dbuffer_list));
+
+      initialized = false;
+    }
+  }
+
+  void open_ipc_handles(std::vector<hipIpcMemHandle_t> const& ipc_handles) {
+    assert(ipc_handles.size() == all_buffer_ipc_handles.size());
+    for (int i = 0; i < world_size; i++) {
+      all_buffer_ipc_handles[i] = ipc_handles[i];
+    }
+
+    for (int i = 0; i < world_size; i++) {
+      if (i != rank) {
+        HIP_CHECK(
+            hipIpcOpenMemHandle((void**)&buffer_list[i], all_buffer_ipc_handles[i], hipIpcMemLazyEnablePeerAccess));
+      } else {
+        buffer_list[i] = dbuffer;
+      }
+    }
+
+    HIP_CHECK(hipMemcpy(dbuffer_list, buffer_list.data(), world_size * sizeof(uint8_t*), hipMemcpyHostToDevice));
+  }
+
+  // bf16 data, reduced in the fp16 kernels at Codec's level.
+  template <class Codec>
+  void allreduce(half const* A, half* B, uint32_t N, hipStream_t stream) {
+    using AllReduceKernel = AllReduceTwoshot<Codec>;
+    uint32_t num_blocks = divceil(N * sizeof(half), Codec::kTileSize);
+    uint32_t grid = std::min<uint32_t>(kMaxNumBlocks, num_blocks);
+    hipLaunchKernelGGL(
+        (allreduce_prototype_twoshot<AllReduceKernel, half>),
+        dim3(grid),
+        dim3(kBlockTwoShot),
+        0,
+        stream,
+        A,
+        B,
+        N,
+        num_blocks,
+        rank,
+        dbuffer_list,
+        data_offset,
+        d_flag_counters,
+        this->kMaxProblemSize);
+    HIP_CHECK(hipGetLastError());
+  }
+};
+
+}  // namespace quickreduce

python/sglang/kernels/jit/csrc/distributed/quick_all_reduce/quick_all_reduce_base.h

@@ -0,0 +1,204 @@
+// JIT copy of sgl-kernel's kernels/aot/csrc/allreduce/quick_all_reduce_base.h, cut to what the INT4 / INT8
+// two-shot of quick_all_reduce.cuh uses: fp16 arithmetic on bf16 data cast at load and store. Keep the
+// arithmetic in step with the AOT copy: the outputs are bit-identical to it.
+#pragma once
+
+#include <hip/hip_bf16.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+
+#include <cstdint>
+
+#define __quickreduce_device_inline__ __device__ __forceinline__
+#define __quickreduce_launch_bounds_two_shot__ __launch_bounds__(256, 4)
+
+namespace quickreduce {
+
+typedef __hip_bfloat162 nv_bfloat162;
+
+using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int;
+using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int;
+
+#if !defined(QR_MAX_BLOCKS)
+#error "QR_MAX_BLOCKS is set by the JIT loader"
+#endif
+
+static constexpr int kBlockSize = 256;
+static constexpr int kAtomStride = kBlockSize;
+
+// Max number of resident blocks; also sizes the flag buffer.
+static constexpr int kMaxNumBlocks = QR_MAX_BLOCKS;
+
+// Standard CDNA wavefront size.
+static constexpr int kWavefront = 64;
+
+static dim3 constexpr kBlockTwoShot = {kWavefront, kBlockSize / kWavefront, 1};
+
+// Threads in a quantization group; each of its block scales covers 32 fp16 values.
+static constexpr int kThreadGroupSize = 8;
+
+__quickreduce_device_inline__ __host__ unsigned long divceil(unsigned long x, unsigned long y) {
+  return ((x + y - 1) / y);
+}
+
+union BufferResource {
+  __quickreduce_device_inline__ constexpr BufferResource() : config(0x00020000U) {}
+
+  __quickreduce_device_inline__ constexpr BufferResource(void* buffer_address, uint32_t buffer_size)
+      : address(buffer_address), range(buffer_size), config(0x00020000U) {}
+
+  int32x4_t descriptor;
+  struct {
+    void* address;    // 8B: a 48b address and a 16b stride (unused)
+    uint32_t range;   // Byte range for the buffer resource
+    uint32_t config;  // Constant, DFMT=32b
+  };
+};
+
+__quickreduce_device_inline__ static int32x4_t buffer_load_dwordx4(
+    int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm("llvm.amdgcn.raw.buffer.load.v4i32");
+
+__quickreduce_device_inline__ static void
+buffer_store_dwordx4(int32x4_t data, int32x4_t srsrc, int32_t voffset, int32_t soffset, int32_t aux) __asm(
+    "llvm.amdgcn.raw.buffer.store.v4i32");
+// MODE.FP16_OVFL clamps overflowing fp16 results to +/-MAX_FP16 instead of inf, the f32 -> f16
+// conversion (v_cvt_pk_f16_f32) included. The memory clobber is load-bearing: without it the
+// compiler may hoist a conversion above the s_setreg, and that conversion still produces inf.
+__quickreduce_device_inline__ static void set_fp16_ovfl(bool const value) {
+#if defined(__gfx942__) || defined(__gfx950__)
+  if (value) {
+    asm volatile("s_setreg_imm32_b32 0x5c1, 1;" ::: "memory");
+  } else {
+    asm volatile("s_setreg_imm32_b32 0x5c1, 0;" ::: "memory");
+  }
+#endif
+}
+template <typename T>
+__quickreduce_device_inline__ void packed_assign_add(int32x4_t* A, int32x4_t* B);
+
+template <>
+__quickreduce_device_inline__ void packed_assign_add<half>(int32x4_t* A, int32x4_t* B) {
+  int32x4_t& tR_fragment = A[0];
+  int32x4_t& tA_fragment = B[0];
+
+  asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(tR_fragment[0]) : "v"(tR_fragment[0]), "v"(tA_fragment[0]));
+  asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(tR_fragment[1]) : "v"(tR_fragment[1]), "v"(tA_fragment[1]));
+  asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(tR_fragment[2]) : "v"(tR_fragment[2]), "v"(tA_fragment[2]));
+  asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(tR_fragment[3]) : "v"(tR_fragment[3]), "v"(tA_fragment[3]));
+}
+
+template <typename T>
+__quickreduce_device_inline__ int packed_max(int a, int b);
+
+template <>
+__quickreduce_device_inline__ int packed_max<half>(int a, int b) {
+  int result;
+  asm volatile("v_pk_max_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
+  return result;
+}
+
+template <typename T>
+__quickreduce_device_inline__ int packed_min(int a, int b);
+
+template <>
+__quickreduce_device_inline__ int packed_min<half>(int a, int b) {
+  int result;
+  asm volatile("v_pk_min_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
+  return result;
+}
+
+template <typename T>
+__quickreduce_device_inline__ int packed_abs_max(int a, int b);
+
+template <>
+__quickreduce_device_inline__ int packed_abs_max<half>(int a, int b) {
+  half2 wmaxh2 = __builtin_bit_cast(half2, a);
+  half2 wminh2 = __builtin_bit_cast(half2, b);
+  half2 wblockmaxh2;
+
+  wblockmaxh2.x = __hgt(__habs(wmaxh2.x), __habs(wminh2.x)) ? wmaxh2.x : wminh2.x;
+  wblockmaxh2.y = __hgt(__habs(wmaxh2.y), __habs(wminh2.y)) ? wmaxh2.y : wminh2.y;
+  return __builtin_bit_cast(int, wblockmaxh2);
+}
+
+template <typename T>
+__quickreduce_device_inline__ int packed_add(int a, int b);
+
+template <>
+__quickreduce_device_inline__ int packed_add<half>(int a, int b) {
+  int result;
+  asm volatile("v_pk_add_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
+  return result;
+}
+
+template <>
+__quickreduce_device_inline__ int packed_add<int16_t>(int a, int b) {
+  int result;
+  asm volatile("v_pk_add_i16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
+  return result;
+}
+
+template <typename T>
+__quickreduce_device_inline__ int packed_mul(int a, int b);
+
+template <>
+__quickreduce_device_inline__ int packed_mul<half>(int a, int b) {
+  int result;
+  asm volatile("v_pk_mul_f16 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b));
+  return result;
+}
+
+template <typename T>
+__quickreduce_device_inline__ int packed_rcp(int a);
+
+template <>
+__quickreduce_device_inline__ int packed_rcp<half>(int a) {
+  return __builtin_bit_cast(int, h2rcp(__builtin_bit_cast(half2, a)));
+}
+
+__quickreduce_device_inline__ float T2float_cast(half a) {
+  return __half2float(a);
+}
+
+// Max and min over each 8-lane group, combined by packed_abs_max, in every lane of the group.
+// DPP lane moves (pairs, quads, then the two halves of each 8-lane half-row); max and min are
+// exact, so the result equals the AOT reduce-to-leader-and-broadcast.
+template <typename T>
+__quickreduce_device_inline__ int group_abs_max(int32x4_t atom) {
+  int a = packed_max<T>(atom[0], atom[1]);
+  int b = packed_max<T>(atom[2], atom[3]);
+  int wmax = packed_max<T>(a, b);
+  a = packed_min<T>(atom[0], atom[1]);
+  b = packed_min<T>(atom[2], atom[3]);
+  int wmin = packed_min<T>(a, b);
+  // DPP controls: quad_perm [1,0,3,2], quad_perm [2,3,0,1], row_half_mirror.
+  wmax = packed_max<T>(wmax, __builtin_amdgcn_mov_dpp(wmax, 0xB1, 0xF, 0xF, false));
+  wmin = packed_min<T>(wmin, __builtin_amdgcn_mov_dpp(wmin, 0xB1, 0xF, 0xF, false));
+  wmax = packed_max<T>(wmax, __builtin_amdgcn_mov_dpp(wmax, 0x4E, 0xF, 0xF, false));
+  wmin = packed_min<T>(wmin, __builtin_amdgcn_mov_dpp(wmin, 0x4E, 0xF, 0xF, false));
+  wmax = packed_max<T>(wmax, __builtin_amdgcn_mov_dpp(wmax, 0x141, 0xF, 0xF, false));
+  wmin = packed_min<T>(wmin, __builtin_amdgcn_mov_dpp(wmin, 0x141, 0xF, 0xF, false));
+  return packed_abs_max<T>(wmax, wmin);
+}
+
+// Global-address-space view of a communication buffer (local or peer: both plain GPU VAs). FLAT
+// accesses also count in lgkmcnt, and their waits are imprecise: with remote FLAT stores in
+// flight, an LDS wait or a wait for an earlier load also waits for those stores.
+#define QR_GLOBAL(T, p) ((__attribute__((address_space(1))) T*)(p))
+
+// s_waitcnt vmcnt(0) (gfx9 encoding; expcnt and lgkmcnt left at their maxima). The flag release on
+// one lane covers only its own wave's stores, so every wave waits for its own before the barrier.
+__quickreduce_device_inline__ void wait_own_stores() {
+  __builtin_amdgcn_s_waitcnt(0x0F70);
+}
+
+__quickreduce_device_inline__ void set_sync_flag(uint32_t* flag_ptr, uint32_t flag) {
+  __atomic_store_n(QR_GLOBAL(uint32_t, flag_ptr), flag, __ATOMIC_RELEASE);
+}
+
+__quickreduce_device_inline__ void wait_sync_flag(uint32_t* flag_ptr, uint32_t flag) {
+  while (__atomic_load_n(QR_GLOBAL(uint32_t, flag_ptr), __ATOMIC_RELAXED) != flag) {
+  }
+}
+
+}  // namespace quickreduce