perf(hisparse): 128-bit non-temporal swap-in copy on ROCm (#33085)

python/sglang/kernels/jit/csrc/hisparse.cuh

@@ -32,15 +32,52 @@ __device__ __forceinline__ int hash_slot(int32_t key, int hash_size) {
 }
 
 #ifdef USE_ROCM
+// 128-bit vector type used by the wide copy path below.
+using TransferVec4 = __attribute__((__vector_size__(4 * sizeof(uint32_t)))) uint32_t;
+
 __device__ __forceinline__ void transfer_item_warp(
     int32_t lane_id, const void* __restrict__ src_addr, void* __restrict__ dst_addr, int64_t item_size_bytes) {
   const auto src = static_cast<const char*>(src_addr);
   auto dst = static_cast<char*>(dst_addr);
 
+  // Wide path: one 128-bit dwordx4 per lane instead of two 64-bit dwordx2,
+  // which halves the load/store instruction count and the number of serialized
+  // round trips per item. The source is usually pinned host DRAM, so those
+  // round trips are what the copy is actually paying for.
+  //
+  // Gated on item_size_bytes >= WARP_SIZE * 16, and this gate is load-bearing:
+  // a 16B-per-lane step only covers item_size_bytes/16 lanes, so below that
+  // threshold the wide path idles part of the wavefront while needing the same
+  // number of iterations as the 8B path. Measured on MI355X (gfx950, wave64)
+  // with 512B items, ungated widening was 14-22% SLOWER because 512/16 = 32
+  // lanes work instead of 512/8 = 64. At 1024B (= WARP_SIZE * 16) the whole
+  // wavefront stays busy and the item moves in a single instruction per lane.
+  //
+  // The load is non-temporal because these items are streamed in once; there
+  // is no reuse to preserve and they should not evict resident lines from
+  // L2/MALL. The store is deliberately left cached: its destination is the
+  // device buffer that the attention kernel reads immediately afterwards.
+  int64_t byte_pos = 0;
+  const bool aligned_16b = ((reinterpret_cast<uintptr_t>(src) | reinterpret_cast<uintptr_t>(dst)) & 0xF) == 0;
+  const bool wide_fills_wave = item_size_bytes >= static_cast<int64_t>(WARP_SIZE) * 16;
+  if (aligned_16b && wide_fills_wave) {
+    constexpr int64_t kVecBytes = static_cast<int64_t>(sizeof(TransferVec4));
+    const int64_t vec_count = item_size_bytes / kVecBytes;
+    const auto src_vec = reinterpret_cast<const TransferVec4*>(src);
+    auto dst_vec = reinterpret_cast<TransferVec4*>(dst);
+    for (int64_t i = lane_id; i < vec_count; i += WARP_SIZE) {
+      dst_vec[i] = __builtin_nontemporal_load(&src_vec[i]);
+    }
+    byte_pos = vec_count * kVecBytes;
+  }
+
+  // 64-bit path: covers the unaligned case in full, and the <= 8-byte
+  // remainder the wide path leaves behind.
   const int64_t word_count = item_size_bytes / static_cast<int64_t>(sizeof(uint64_t));
+  const int64_t word_start = byte_pos / static_cast<int64_t>(sizeof(uint64_t));
   const auto src_words = reinterpret_cast<const uint64_t*>(src);
   auto dst_words = reinterpret_cast<uint64_t*>(dst);
-  for (int64_t i = lane_id; i < word_count; i += WARP_SIZE) {
+  for (int64_t i = word_start + lane_id; i < word_count; i += WARP_SIZE) {
     dst_words[i] = src_words[i];
   }