perf(hisparse): fuse the DSv4 value and scale swap-in copy on ROCm (#33484)

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

@@ -86,6 +86,56 @@ __device__ __forceinline__ void transfer_item_warp(
     dst[i] = src[i];
   }
 }
+
+// Copies one DSv4 C4 token as a single 73-word space instead of two separate
+// transfer_item_warp calls.
+//
+// A token is a 576B value and an 8B scale that sit in different runs of the
+// page row, so it cannot be moved as one contiguous range. Copying the two
+// pieces separately costs three wavefront passes on wave64 -- 64 value words,
+// 8 value words, then 1 scale word -- and because item_size_bytes reaches
+// transfer_item_warp as a runtime argument the compiler keeps the value copy
+// as a rolled loop with s_waitcnt vmcnt(0) inside it, so all three passes
+// serialize on host memory latency. Walking one 73-word space puts the value
+// tail and the scale in the same pass, and the pass count being a compile-time
+// constant lets both loads issue before either store, leaving one exposed
+// round trip instead of three. This is the same shape as the CUDA
+// device::hisparse::transfer_item, adapted to wave64.
+__device__ __forceinline__ void transfer_dsv4_item_warp(
+    int32_t lane_id,
+    const int64_t* __restrict__ src_value,
+    const int64_t* __restrict__ src_scale,
+    int64_t* __restrict__ dst_value,
+    int64_t* __restrict__ dst_scale) {
+  constexpr int32_t kValueWords = static_cast<int32_t>(device::hisparse::kValueBytes / sizeof(int64_t));
+  constexpr int32_t kTotalWords = static_cast<int32_t>(device::hisparse::kItemBytes / sizeof(int64_t));
+  constexpr int32_t kPasses = (kTotalWords + WARP_SIZE - 1) / WARP_SIZE;
+  static_assert(device::hisparse::kValueBytes % sizeof(int64_t) == 0, "value must be whole 64-bit words");
+  static_assert(device::hisparse::kScaleBytes % sizeof(int64_t) == 0, "scale must be whole 64-bit words");
+
+  const int64_t* src_slot[kPasses];
+  int64_t* dst_slot[kPasses];
+  int64_t staged[kPasses];
+
+#pragma unroll
+  for (int32_t p = 0; p < kPasses; ++p) {
+    const int32_t i = p * WARP_SIZE + lane_id;
+    const bool is_value = i < kValueWords;
+    const bool is_scale = !is_value && i < kTotalWords;
+    src_slot[p] = is_value ? src_value + i : (is_scale ? src_scale + (i - kValueWords) : nullptr);
+    dst_slot[p] = is_value ? dst_value + i : (is_scale ? dst_scale + (i - kValueWords) : nullptr);
+    if (src_slot[p] != nullptr) {
+      staged[p] = *src_slot[p];
+    }
+  }
+
+#pragma unroll
+  for (int32_t p = 0; p < kPasses; ++p) {
+    if (dst_slot[p] != nullptr) {
+      *dst_slot[p] = staged[p];
+    }
+  }
+}
 #else
 __device__ __forceinline__ void
 transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) {
@@ -558,14 +608,13 @@ __global__ void load_cache_to_device_buffer_kernel(
       // ROCm path: host cache and device buffer both use the page-padded C4
       // layout (same as the write path and the CUDA branch). We can't reuse
       // device::hisparse::transfer_item here because its warp logic is hardcoded
-      // to a 32-lane warp; on wavefront64 we use the warp-width-agnostic
-      // transfer_item_warp with paged source and destination addressing.
+      // to a 32-lane warp; on wavefront64 we use transfer_dsv4_item_warp, which
+      // moves the value and the scale in one warp-width-agnostic copy.
       using namespace device::hisparse;
       const auto [dst_value_ptr, dst_scale_ptr] = get_pointer_paged(device_buffer_k, static_cast<int32_t>(dst_loc));
       const auto [src_value_ptr, src_scale_ptr] =
           get_pointer_paged(const_cast<void*>(host_cache_k), static_cast<int32_t>(src_loc));
-      transfer_item_warp(lane_id, src_value_ptr, dst_value_ptr, kValueBytes);
-      transfer_item_warp(lane_id, src_scale_ptr, dst_scale_ptr, kScaleBytes);
+      transfer_dsv4_item_warp(lane_id, src_value_ptr, src_scale_ptr, dst_value_ptr, dst_scale_ptr);
 #else
       // CUDA path: page-padded device layout + page-padded host layout, K-only.
       // The host cache is pinned DRAM but uses the same row layout as the GPU C4

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

@@ -604,6 +604,89 @@ def test_load_cache_to_device_buffer_dsv4_mla_miss_copy_layout() -> None:
     )
 
 
+@pytest.mark.skipif(
+    not is_hip(), reason="Covers the ROCm wavefront64 fused DSv4 token copy."
+)
+def test_load_cache_to_device_buffer_dsv4_fused_copy_multi_miss() -> None:
+    """Several DSv4 misses in one launch must each land byte-exact.
+
+    The fused copy walks the 576B value and the 8B scale as one 73-word space,
+    so the seam between them falls on a lane index rather than a call boundary.
+    Vary both the source and the destination page offset, including tokens on
+    the second page, so the seam is not always at the same address.
+    """
+    hot_buffer_size = 4
+    num_pages = 2
+    # seq_len stays above the queried tokens so none of them is the newest
+    # token, which the kernel places without a host copy.
+    seq_len = 16
+    host_locs = list(range(seq_len))
+    miss_tokens = [4, 5, 6, 7]
+    # Source offsets: mid-page, last slot of page 0, first slot of page 1,
+    # last slot of page 1.
+    for token, loc in zip(miss_tokens, [10, 63, 64, 127]):
+        host_locs[token] = loc
+    # Destination offsets: first, second, last of page 0, then page 1.
+    device_locs = [0, 1, 63, 64, 65]
+
+    host_cache = torch.zeros(
+        (num_pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device="cpu", pin_memory=True
+    )
+    for loc in host_locs:
+        _write_dsv4_token(host_cache, loc, seed=loc + 1)
+
+    device_buffer = torch.full(
+        (num_pages, DSV4_PAGE_BYTES), 0xFF, dtype=torch.uint8, device=DEVICE
+    )
+
+    top_k_tokens = torch.tensor([miss_tokens], dtype=torch.int32, device=DEVICE)
+    out = torch.full_like(top_k_tokens, -1)
+
+    load_cache_to_device_buffer_dsv4_mla(
+        top_k_tokens=top_k_tokens,
+        device_buffer_tokens=torch.tensor(
+            [[0, 1, 2, 3, -1]], dtype=torch.int32, device=DEVICE
+        ),
+        host_cache_locs=torch.tensor([host_locs], dtype=torch.int64, device=DEVICE),
+        device_buffer_locs=torch.tensor(
+            [device_locs], dtype=torch.int32, device=DEVICE
+        ),
+        host_cache=host_cache,
+        device_buffer=device_buffer,
+        top_k_device_locs=out,
+        req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
+        seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE),
+        lru_slots=torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE).view(
+            1, -1
+        ),
+        item_size_bytes=DSV4_ITEM_BYTES,
+        num_top_k=len(miss_tokens),
+        hot_buffer_size=hot_buffer_size,
+        page_size=DSV4_PAGE_SIZE,
+        block_size=256,
+        num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
+    )
+    torch.cuda.synchronize()
+
+    # Which slot each miss evicts is up to the LRU, so take the destinations
+    # from the kernel; only require that they are distinct and in range.
+    landed = out.cpu().tolist()[0]
+    assert len(set(landed)) == len(landed)
+    assert set(landed).issubset(device_locs)
+
+    device_cpu = device_buffer.cpu()
+    for token, dst_loc in zip(miss_tokens, landed):
+        assert torch.equal(
+            _read_dsv4_token(device_cpu, dst_loc),
+            _read_dsv4_token(host_cache, host_locs[token]),
+        ), f"token {token} -> device loc {dst_loc}"
+
+    # Slots the kernel never wrote must keep their fill, so an over-copy that
+    # ran past the value or the scale would be caught.
+    for loc in set(device_locs) - set(landed):
+        assert torch.all(_read_dsv4_token(device_cpu, loc) == 0xFF)
+
+
 @pytest.mark.skipif(
     not is_hip(), reason="Covers a ROCm wavefront64 LRU writeback regression."
 )