From d22c4cc177c4ecac1285fc2bc122311f0eb38646 Mon Sep 17 00:00:00 2001 From: TianDi101 Date: Sun, 16 Aug 2026 07:05:40 +0800 Subject: [PATCH] =?UTF-8?q?[AMD]=20perf(sgl-kernel):=20default=20block=5Fq?= =?UTF-8?q?uota=3D16=20for=20MLA=20page=5Ffirst=20KV=20gather=E2=80=A6=20(?= =?UTF-8?q?#30024)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Niko Ma Co-authored-by: figo Co-authored-by: AMD-yanfeiwang --- .../kernels/aot/csrc/kvcacheio/transfer.cu | 134 +++++++++++++++++- .../aot/python/sgl_kernel/kvcacheio.py | 28 +++- 2 files changed, 153 insertions(+), 9 deletions(-) diff --git a/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu b/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu index 248571089..bca392314 100644 --- a/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu +++ b/python/sglang/kernels/aot/csrc/kvcacheio/transfer.cu @@ -17,6 +17,7 @@ #include "utils.h" // WARP_SIZE #endif +#if !defined(USE_ROCM) && !defined(USE_MUSA) __device__ __forceinline__ void transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) { const uint64_t* __restrict__ src = static_cast(src_addr); @@ -25,17 +26,66 @@ transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_ #pragma unroll for (int j = lane_id; j < total_chunks; j += WARP_SIZE) { -#if !defined(USE_ROCM) && !defined(USE_MUSA) uint64_t tmp; asm volatile("ld.global.nc.b64 %0,[%1];" : "=l"(tmp) : "l"(src + j) : "memory"); asm volatile("st.global.cg.b64 [%0],%1;" ::"l"(dst + j), "l"(tmp) : "memory"); - -#else - uint64_t tmp = __builtin_nontemporal_load(src + j); - __builtin_nontemporal_store(tmp, dst + j); -#endif } } +#elif defined(USE_ROCM) +// ROCm: use 128-bit streaming load/store when 16B-aligned, so fewer CUs are +// needed to saturate the host fabric; falls back to 64-bit otherwise. +typedef uint32_t sgl_u32x4 __attribute__((ext_vector_type(4))); +__device__ __forceinline__ void +transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) { + const uintptr_t addr_or = reinterpret_cast(src_addr) | reinterpret_cast(dst_addr); + if ((addr_or & 0xF) == 0) { + const sgl_u32x4* __restrict__ src = static_cast(src_addr); + sgl_u32x4* __restrict__ dst = static_cast(dst_addr); + const int chunks16 = item_size_bytes / 16; + for (int j = lane_id; j < chunks16; j += WARP_SIZE) { + sgl_u32x4 tmp = __builtin_nontemporal_load(src + j); + __builtin_nontemporal_store(tmp, dst + j); + } + // Trailing bytes: item_size_bytes % 8 == 0 is guaranteed by the launcher, + // so the remainder is at most one 8B word. + const int done_bytes = chunks16 * 16; + const int rem8 = static_cast(item_size_bytes - done_bytes) / 8; + if (rem8) { + const uint64_t* __restrict__ src8 = + reinterpret_cast(static_cast(src_addr) + done_bytes); + uint64_t* __restrict__ dst8 = reinterpret_cast(static_cast(dst_addr) + done_bytes); + for (int j = lane_id; j < rem8; j += WARP_SIZE) { + uint64_t tmp = __builtin_nontemporal_load(src8 + j); + __builtin_nontemporal_store(tmp, dst8 + j); + } + } + } else { + const uint64_t* __restrict__ src = static_cast(src_addr); + uint64_t* __restrict__ dst = static_cast(dst_addr); + const int total_chunks = item_size_bytes / sizeof(uint64_t); + for (int j = lane_id; j < total_chunks; j += WARP_SIZE) { + uint64_t tmp = __builtin_nontemporal_load(src + j); + __builtin_nontemporal_store(tmp, dst + j); + } + } +} +#else +// MUSA: keep the original scalar nontemporal load/store path; the 128-bit +// ROCm path above relies on a HIP/clang ext_vector_type extension that isn't +// guaranteed to be available/correct under the MUSA compiler. +__device__ __forceinline__ void +transfer_item_warp(int32_t lane_id, const void* src_addr, void* dst_addr, int64_t item_size_bytes) { + const uint64_t* __restrict__ src = static_cast(src_addr); + uint64_t* __restrict__ dst = static_cast(dst_addr); + const int total_chunks = item_size_bytes / sizeof(uint64_t); + +#pragma unroll + for (int j = lane_id; j < total_chunks; j += WARP_SIZE) { + uint64_t tmp = __builtin_nontemporal_load(src + j); + __builtin_nontemporal_store(tmp, dst + j); + } +} +#endif template __device__ __forceinline__ T* get_global_offset_lf( @@ -945,6 +995,78 @@ inline void transfer_kv_page_first_direct_impl( }; #if defined(USE_ROCM) || !defined(CUDA_VERSION) || CUDA_VERSION < 12080 +#if defined(USE_ROCM) + // Opt-in HIP batch copy path (mirrors cudaMemcpyBatchAsync); disabled by + // default, falls back to per-page copy below. + constexpr bool kEnableHipBatch = false; + if (kEnableHipBatch) { + std::vector b_srcs, b_dsts; + std::vector b_sizes; + auto batch_append = [&](const at::Tensor& s, const at::Tensor& d, int64_t si, int64_t di, int64_t ps) { + const int64_t esz = s.element_size(); + b_srcs.push_back(static_cast(s.data_ptr()) + si * s.stride(0) * esz); + b_dsts.push_back(static_cast(d.data_ptr()) + di * d.stride(0) * esz); + b_sizes.push_back(static_cast(ps) * static_cast(s.stride(0)) * static_cast(esz)); + }; + if constexpr (IsLf2Pf) { + const bool is_mla = dst_ptrs.size() == 1; + const int64_t num_layers = is_mla ? src_ptrs.size() : src_ptrs.size() / 2; + for (const auto i : c10::irange(num_pages)) { + const int64_t s_index = src_indices_ptr[i * page_size]; + const int64_t d_index = dst_indices_ptr[i * page_size] / page_size; + for (int64_t j = 0; j < num_layers; ++j) { + batch_append( + src_ptrs[j], dst_ptrs[0].select(0, d_index).select(0, start_layer_id + j), s_index, 0, page_size); + if (!is_mla) { + batch_append( + src_ptrs[j + num_layers], + dst_ptrs[1].select(0, d_index).select(0, start_layer_id + j), + s_index, + 0, + page_size); + } + } + } + } else { + const bool is_mla = src_ptrs.size() == 1; + const int64_t num_layers = is_mla ? dst_ptrs.size() : dst_ptrs.size() / 2; + for (const auto i : c10::irange(num_pages)) { + const int64_t s_index = src_indices_ptr[i * page_size] / page_size; + const int64_t d_index = dst_indices_ptr[i * page_size]; + for (int64_t j = 0; j < num_layers; ++j) { + batch_append( + src_ptrs[0].select(0, s_index).select(0, start_layer_id + j), dst_ptrs[j], 0, d_index, page_size); + if (!is_mla) { + batch_append( + src_ptrs[1].select(0, s_index).select(0, start_layer_id + j), + dst_ptrs[j + num_layers], + 0, + d_index, + page_size); + } + } + } + } + if (!b_srcs.empty()) { + size_t fail_idx = std::numeric_limits::max(); + hipError_t err = hipMemcpyBatchAsync( + b_dsts.data(), + b_srcs.data(), + b_sizes.data(), + b_srcs.size(), + nullptr, + nullptr, + 0, + &fail_idx, + at::cuda::getCurrentCUDAStream().stream()); + if (err != hipSuccess) { + TORCH_WARN_ONCE("hipMemcpyBatchAsync failed (", hipGetErrorString(err), "), falling back to per-page copy"); + fallback_to_page_copy(); + } + } + return; + } +#endif fallback_to_page_copy(); return; diff --git a/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py b/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py index 0601c8fe6..43e3cdec8 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/kvcacheio.py @@ -1,4 +1,5 @@ -from typing import List +import os +from typing import List, Optional import torch @@ -10,6 +11,23 @@ def is_hip() -> bool: _is_hip = is_hip() +def _default_mla_block_quota() -> int: + """CU (block) quota for the MLA page_first KV gather kernel. + + Defaults to 16 on ROCm / 2 on CUDA. Override with the + SGLANG_HICACHE_BLOCK_QUOTA environment variable to tune how many CUs the + kernel is launched with. + """ + default = 16 if _is_hip else 2 + override = os.environ.get("SGLANG_HICACHE_BLOCK_QUOTA") + if override is None: + return default + try: + return int(override) + except ValueError: + return default + + def transfer_kv_per_layer( src_k: torch.Tensor, dst_k: torch.Tensor, @@ -236,9 +254,11 @@ def transfer_kv_per_layer_mla( src_indices: torch.Tensor, dst_indices: torch.Tensor, item_size: int, - block_quota: int = 2, + block_quota: Optional[int] = None, num_warps_per_block: int = 16 if _is_hip else 32, ): + if block_quota is None: + block_quota = _default_mla_block_quota() torch.ops.sgl_kernel.transfer_kv_per_layer_mla.default( src, dst, @@ -258,9 +278,11 @@ def transfer_kv_per_layer_mla_pf_lf( layer_id: int, item_size: int, src_layout_dim: int, - block_quota: int = 2, + block_quota: Optional[int] = None, num_warps_per_block: int = 16 if _is_hip else 32, ): + if block_quota is None: + block_quota = _default_mla_block_quota() torch.ops.sgl_kernel.transfer_kv_per_layer_mla_pf_lf.default( src, dst,