From ca0f8a0f4ced5fc107f27a3be386d12bc7850711 Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Tue, 11 Aug 2026 05:30:15 +0800 Subject: [PATCH] perf(hisparse): fuse the DSv4 value and scale swap-in copy on ROCm (#33484) --- python/sglang/kernels/jit/csrc/hisparse.cuh | 57 ++++++++++++- .../kernels/ops/kvcache/test_hisparse.py | 83 +++++++++++++++++++ 2 files changed, 136 insertions(+), 4 deletions(-) diff --git a/python/sglang/kernels/jit/csrc/hisparse.cuh b/python/sglang/kernels/jit/csrc/hisparse.cuh index 80818fc21..815378ae7 100644 --- a/python/sglang/kernels/jit/csrc/hisparse.cuh +++ b/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(device::hisparse::kValueBytes / sizeof(int64_t)); + constexpr int32_t kTotalWords = static_cast(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(dst_loc)); const auto [src_value_ptr, src_scale_ptr] = get_pointer_paged(const_cast(host_cache_k), static_cast(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 diff --git a/test/registered/kernels/ops/kvcache/test_hisparse.py b/test/registered/kernels/ops/kvcache/test_hisparse.py index 22d0c6d1f..bf93c55a1 100644 --- a/test/registered/kernels/ops/kvcache/test_hisparse.py +++ b/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." )