From 1a8e4876b6ddccc5f2712efdca7f87fdc5514c7b Mon Sep 17 00:00:00 2001 From: AMD-yanfeiwang Date: Tue, 11 Aug 2026 04:27:15 +0800 Subject: [PATCH] perf(hisparse): 128-bit non-temporal swap-in copy on ROCm (#33085) --- python/sglang/kernels/jit/csrc/hisparse.cuh | 39 ++++++++- .../kernels/ops/kvcache/test_hisparse.py | 84 +++++++++++++++++++ 2 files changed, 122 insertions(+), 1 deletion(-) diff --git a/python/sglang/kernels/jit/csrc/hisparse.cuh b/python/sglang/kernels/jit/csrc/hisparse.cuh index dffdffb32..80818fc21 100644 --- a/python/sglang/kernels/jit/csrc/hisparse.cuh +++ b/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(src_addr); auto dst = static_cast(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(src) | reinterpret_cast(dst)) & 0xF) == 0; + const bool wide_fills_wave = item_size_bytes >= static_cast(WARP_SIZE) * 16; + if (aligned_16b && wide_fills_wave) { + constexpr int64_t kVecBytes = static_cast(sizeof(TransferVec4)); + const int64_t vec_count = item_size_bytes / kVecBytes; + const auto src_vec = reinterpret_cast(src); + auto dst_vec = reinterpret_cast(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(sizeof(uint64_t)); + const int64_t word_start = byte_pos / static_cast(sizeof(uint64_t)); const auto src_words = reinterpret_cast(src); auto dst_words = reinterpret_cast(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]; } diff --git a/test/registered/kernels/ops/kvcache/test_hisparse.py b/test/registered/kernels/ops/kvcache/test_hisparse.py index 8824099e4..22d0c6d1f 100644 --- a/test/registered/kernels/ops/kvcache/test_hisparse.py +++ b/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]],