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];
   }
 

test/registered/kernels/ops/kvcache/test_hisparse.py

@@ -380,6 +380,90 @@ def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
     assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6])
 
 
+@pytest.mark.skipif(
+    not is_hip(),
+    reason="CUDA transfer_item_warp assumes 16B-aligned items with no sub-8B remainder.",
+)
+@pytest.mark.parametrize(
+    "kv_dim,miss_token",
+    [
+        # Tokens 0..3 are resident, so the queried token must be >= 4 to miss.
+        # The destination is always slot 0, so the source offset
+        # (miss_token * item size) is what decides the 16B-alignment check.
+        (256, 4),  # 1024B, exactly the gate: one 16B step per lane, no remainder
+        (257, 4),  # 1028B: wide path + 4B byte tail
+        (258, 4),  # 1032B: wide path + one 64-bit word
+        (260, 4),  # 1040B: two wide iterations on lane 0
+        (257, 5),  # 1028B, source at 5140: unaligned, wide path skipped
+        (5, 4),  # 20B: below the gate, 64-bit loop + 4B byte tail
+    ],
+)
+def test_load_cache_to_device_buffer_miss_copy_is_byte_exact(
+    kv_dim: int, miss_token: int
+) -> None:
+    """A miss must copy the item byte-exactly for any item size and alignment.
+
+    Every other ROCm case in this file uses a 32B item, far below the
+    WARP_SIZE * 16 wide-copy gate, so none of them reaches the wide path at all,
+    let alone the seam between it and the remainder loops. These sizes sit on
+    both sides of the gate and cover each remainder shape.
+    """
+    item_size_bytes = kv_dim * torch.empty((), dtype=DTYPE).element_size()
+
+    host_cache = torch.empty(
+        (HOST_CACHE_SIZE, 1, kv_dim), dtype=DTYPE, device="cpu", pin_memory=True
+    )
+    host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache))
+    device_buffer = torch.full(
+        (DEVICE_CACHE_SIZE, 1, kv_dim), -1, dtype=DTYPE, device=DEVICE
+    )
+
+    # Slots 0..3 hold tokens 0..3; slot 4 is the reserved newest slot.
+    device_buffer_locs = torch.tensor(
+        [[0, 1, 2, 3, 4]], dtype=torch.int32, device=DEVICE
+    )
+    device_buffer_tokens = torch.tensor(
+        [[0, 1, 2, 3, -1]], dtype=torch.int32, device=DEVICE
+    )
+    for slot in range(HOT_BUFFER_SIZE):
+        device_buffer[slot].copy_(host_cache[slot].to(DEVICE))
+    torch.cuda.synchronize()
+
+    top_k_tokens = torch.tensor([[miss_token]], dtype=torch.int32, device=DEVICE)
+    out = torch.full_like(top_k_tokens, -1)
+
+    load_cache_to_device_buffer_mla(
+        top_k_tokens=top_k_tokens,
+        device_buffer_tokens=device_buffer_tokens,
+        host_cache_locs=torch.arange(
+            HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE
+        ).view(1, -1),
+        device_buffer_locs=device_buffer_locs,
+        host_cache=host_cache,
+        device_buffer=device_buffer,
+        top_k_device_locs=out,
+        req_pool_indices=torch.arange(1, dtype=torch.int64, device=DEVICE),
+        seq_lens=torch.full((1,), 8, dtype=torch.int32, device=DEVICE),
+        lru_slots=torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE).view(
+            1, -1
+        ),
+        item_size_bytes=item_size_bytes,
+        num_top_k=1,
+        hot_buffer_size=HOT_BUFFER_SIZE,
+        page_size=1,
+        block_size=256,
+        num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
+    )
+    torch.cuda.synchronize()
+
+    # The miss evicts the LRU head (slot 0, physical loc 0) and lands there.
+    assert torch.equal(out.cpu(), torch.tensor([[0]], dtype=torch.int32))
+    assert torch.equal(device_buffer[0].cpu(), host_cache[miss_token])
+    # Neighbouring slots must not be corrupted by an over-copy.
+    for slot in range(1, HOT_BUFFER_SIZE):
+        assert torch.equal(device_buffer[slot].cpu(), host_cache[slot])
+
+
 def test_load_cache_to_device_buffer_multiple_misses_copy_all_slots() -> None:
     state = _make_state(
         [[9, 7, 3, 5, 11]],