From 7ad55e43865c3555bad01360743f0a8d92c09677 Mon Sep 17 00:00:00 2001 From: cctry Date: Mon, 21 Sep 2026 10:38:23 -0700 Subject: [PATCH] [HiCache] TMA-staged host<->device KV transfer kernel (sm_90+) (#40278) --- .../jit/csrc/kvcacheio/hicache_tma.cuh | 646 ++++++++++++++++++ python/sglang/kernels/ops/kvcache/hicache.py | 146 +++- python/sglang/srt/environ.py | 3 + python/sglang/srt/mem_cache/pool_host/mha.py | 12 +- python/sglang/srt/mem_cache/pool_host/mla.py | 8 +- .../benchmark/kvcache/bench_hicache.py | 104 ++- .../kernels/ops/kvcache/test_hicache_tma.py | 173 +++++ 7 files changed, 1060 insertions(+), 32 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/kvcacheio/hicache_tma.cuh create mode 100644 test/registered/kernels/ops/kvcache/test_hicache_tma.py diff --git a/python/sglang/kernels/jit/csrc/kvcacheio/hicache_tma.cuh b/python/sglang/kernels/jit/csrc/kvcacheio/hicache_tma.cuh new file mode 100644 index 000000000..1dda34cda --- /dev/null +++ b/python/sglang/kernels/jit/csrc/kvcacheio/hicache_tma.cuh @@ -0,0 +1,646 @@ +// HiCache host<->device KV transfer staged through shared memory by the TMA +// bulk-copy engine (sm_90+). +// +// One CTA owns a ring of shared-memory stages. A single loader warp fills +// stages with `cp.async.bulk` (global -> shared, completion counted on an +// mbarrier), which keeps the whole ring in flight with no registers or issue +// slots; the register-staging kernel in hicache.cuh cannot hold enough host +// loads in flight per SM for that. Store warps drain filled stages. Row size is +// a runtime parameter (multiple of 16 B), so one compiled module serves every +// KV shape. +// +// Two hardware facts fix the shape of the kernel (numbers in the PR): every SM +// has a fixed write port to L2, so one CTA cannot exceed it and the block +// quota decides how much of the host link is used; and the TMA unit processes +// bulk ops at a fixed per-op rate, so a source run must move as one op +// (contiguous run -> one 1D bulk copy; strided page run -> one 2D tensor-map +// box), never one op per row. Revisit both if a future part widens the SM +// write port or the TMA op rate. +// +// Work unit ("chunk") = (buffer K|V, layer, run of consecutive positions of the +// index arrays). The loader prefetches the next chunk's indices before blocking +// on the ring so their latency overlaps the wait, and stashes the destination +// indices in smem for the store warps. + +#pragma once + +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace sglang { + +namespace device::ptx { + +// global -> shared::cta, completion counted on `bar` (arm with mbar_arrive_expect_tx). +SGL_DEVICE void bulk_g2s(void* dst_smem, const void* src_gmem, uint32_t bytes, uint64_t* bar) { + asm volatile("cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];" ::"r"( + to_shared(dst_smem)), + "l"(src_gmem), + "r"(bytes), + "r"(to_shared(bar)) + : "memory"); +} + +// 2D tiled tensor-map box at element coords (x, y) -> shared::cta. +SGL_DEVICE void bulk_tensor_2d_g2s(void* dst_smem, const CUtensorMap* map, int32_t x, int32_t y, uint64_t* bar) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile.mbarrier::complete_tx::bytes [%0], [%1, {%2, %3}], [%4];" :: + "r"(to_shared(dst_smem)), + "l"(map), + "r"(x), + "r"(y), + "r"(to_shared(bar)) + : "memory"); +} + +// shared::cta -> global, tracked by the issuing thread's bulk group. +SGL_DEVICE void bulk_s2g(void* dst_gmem, const void* src_smem, uint32_t bytes) { + asm volatile("cp.async.bulk.global.shared::cta.bulk_group [%0], [%1], %2;" ::"l"(dst_gmem), + "r"(to_shared(src_smem)), + "r"(bytes) + : "memory"); +} + +SGL_DEVICE void bulk_commit_group() { + asm volatile("cp.async.bulk.commit_group;" ::: "memory"); +} + +// Block until every committed bulk group has finished reading its smem source. +SGL_DEVICE void bulk_wait_group_read_all() { + asm volatile("cp.async.bulk.wait_group.read 0;" ::: "memory"); +} + +// Same, but the most recent group may still be reading. +SGL_DEVICE void bulk_wait_group_read_one() { + asm volatile("cp.async.bulk.wait_group.read 1;" ::: "memory"); +} + +// Block until every committed bulk group has fully landed in global memory. +SGL_DEVICE void bulk_wait_group_all() { + asm volatile("cp.async.bulk.wait_group 0;" ::: "memory"); +} + +SGL_DEVICE void fence_mbarrier_init() { + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); +} + +} // namespace device::ptx + +struct HicacheTmaParams { + // Either a direct base pointer (`*_is_table == false`) or a device array of + // `num_layers` uint64 base pointers. `v_*` is unused when `has_v == false` (MLA). + const void* __restrict__ k_src; + const void* __restrict__ v_src; + void* __restrict__ k_dst; + void* __restrict__ v_dst; + const void* __restrict__ indices_src; + const void* __restrict__ indices_dst; + int64_t src_stride; // bytes between consecutive token rows + int64_t dst_stride; + uint32_t row_bytes; // bytes copied per token row, multiple of 16 + uint32_t length; // number of token indices + uint32_t num_layers; + uint64_t units_per_row_magic; // ceil(2^32 / (row_bytes / 16)); see store loop + bool src_is_table; + bool dst_is_table; + bool has_v; + // Strided source rows ([K, V] views): one box per chunk instead of one op per row. + bool has_src_map; + CUtensorMap src_map[2]; +}; + +// Rows of one chunk are spread over the loader lanes; each lane prefetches at +// most this many row indices, which bounds rows per chunk to 32x this. +inline constexpr uint32_t kHicacheTmaRowsPerLane = 4; +inline constexpr uint32_t kHicacheTmaMaxRows = kHicacheTmaRowsPerLane * device::kWarpThreads; +// Tensor-map boxes are limited to 256 elements per dimension; rows are mapped as +// 8-byte elements so this is the widest row a 2D box can cover. +inline constexpr uint32_t kHicacheTmaMaxMapRowBytes = 256 * 8; + +// Rows per chunk: largest power of two that fits the stage, so chunks never +// straddle a (power-of-two) page and a page run stays one bulk op. +__host__ __device__ constexpr uint32_t hicache_tma_rows_per_chunk(uint32_t stage_bytes, uint32_t row_bytes) { + uint32_t rows = 1; + while (rows * 2 <= stage_bytes / row_bytes && rows * 2 <= kHicacheTmaMaxRows) + rows *= 2; + return rows; +} + +template +struct HicacheTmaSmem { + alignas(128) uint8_t stages[kNumStages][kStageBytes]; + int64_t dst_idx[kNumStages][kHicacheTmaMaxRows]; // destination row indices of the staged chunk + uint32_t dst_run[kNumStages]; // destination rows form one contiguous span + uint64_t full[kNumStages]; // loader -> storers: stage filled (count 1) + uint64_t empty[kNumStages]; // storers -> loader: stage drained (count kStoreWarps) +}; + +template +__global__ void __launch_bounds__((1 + kStoreWarps) * device::kWarpThreads, 1) + hicache_tma_transfer_kernel(const __grid_constant__ HicacheTmaParams p) { +#if SGL_ARCH_HOPPER_OR_GREATER + using namespace device; + using Smem = HicacheTmaSmem; + extern __shared__ __align__(128) uint8_t smem_raw[]; + auto& smem = *reinterpret_cast(smem_raw); + + const uint32_t warp = threadIdx.x / kWarpThreads; + const uint32_t lane = threadIdx.x % kWarpThreads; + + if (threadIdx.x == 0) { + for (uint32_t s = 0; s < kNumStages; ++s) { + ptx::mbar_init(&smem.full[s], 1); + ptx::mbar_init(&smem.empty[s], kStoreWarps); + } + ptx::fence_mbarrier_init(); + } + __syncthreads(); + + const uint32_t rows_per_chunk = hicache_tma_rows_per_chunk(kStageBytes, p.row_bytes); + const uint32_t token_chunks = div_ceil(p.length, rows_per_chunk); + const uint32_t num_chunks = (p.has_v ? 2u : 1u) * p.num_layers * token_chunks; + const T* idx_src = static_cast(p.indices_src); + const T* idx_dst = static_cast(p.indices_dst); + + struct ChunkInfo { + uint32_t t0; // first index position + uint32_t rows; // rows in this chunk + uint32_t layer; + bool is_v; + }; + auto describe = [&](uint32_t chunk) { + const uint32_t tc = chunk % token_chunks; + const uint32_t rest = chunk / token_chunks; + const uint32_t t0 = tc * rows_per_chunk; + return ChunkInfo{t0, min(rows_per_chunk, p.length - t0), rest % p.num_layers, (rest / p.num_layers) != 0}; + }; + auto base_ptr = [&](const void* direct_or_table, bool is_table, uint32_t layer) -> const void* { + return is_table ? reinterpret_cast(static_cast(direct_or_table)[layer]) + : direct_or_table; + }; + + if (warp == 0) { + // ---- loader: global -> smem ring via TMA bulk copies + struct Prefetch { + const void* src_base; + T src[kHicacheTmaRowsPerLane]; // rows lane, lane + 32, ... + T dst[kHicacheTmaRowsPerLane]; + }; + auto prefetch = [&](uint32_t chunk) { + const ChunkInfo c = describe(chunk); + Prefetch pf; + pf.src_base = base_ptr(c.is_v ? p.v_src : p.k_src, p.src_is_table, c.layer); +#pragma unroll + for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) { + const uint32_t r = lane + k * kWarpThreads; + pf.src[k] = r < c.rows ? idx_src[c.t0 + r] : T{0}; + pf.dst[k] = r < c.rows ? idx_dst[c.t0 + r] : T{0}; + } + return pf; + }; + + uint32_t chunk = blockIdx.x; + Prefetch next = chunk < num_chunks ? prefetch(chunk) : Prefetch{}; + for (uint32_t it = 0; chunk < num_chunks; chunk += gridDim.x, ++it) { + const ChunkInfo c = describe(chunk); + const Prefetch cur = next; + if (chunk + gridDim.x < num_chunks) next = prefetch(chunk + gridDim.x); + + // A run: every row sits at first + r (whole pages in order), on either side. + const T first = __shfl_sync(warp::kFullMask, cur.src[0], 0); + const T first_dst = __shfl_sync(warp::kFullMask, cur.dst[0], 0); + bool run = true, run_dst = true; +#pragma unroll + for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) { + const uint32_t r = lane + k * kWarpThreads; + run &= r >= c.rows || cur.src[k] == first + static_cast(r); + run_dst &= r >= c.rows || cur.dst[k] == first_dst + static_cast(r); + } + run = __all_sync(warp::kFullMask, run); + run_dst = __all_sync(warp::kFullMask, run_dst); + const bool contiguous = run && p.src_stride == p.row_bytes; + const bool boxed = run && !contiguous && p.has_src_map; + + const uint32_t s = it % kNumStages; + ptx::mbar_wait_parity(&smem.empty[s], ((it / kNumStages) & 1) ^ 1); // fresh ring passes +#pragma unroll + for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) { + const uint32_t r = lane + k * kWarpThreads; + if (r < c.rows) smem.dst_idx[s][r] = static_cast(cur.dst[k]); + } + if (lane == 0) smem.dst_run[s] = run_dst && p.dst_stride == p.row_bytes; + // A box always lands rows_per_chunk rows (out-of-range rows are zero-filled). + const uint32_t tx_bytes = (boxed ? rows_per_chunk : c.rows) * p.row_bytes; + if (lane == 0) ptx::mbar_arrive_expect_tx(&smem.full[s], tx_bytes); + __syncwarp(); + + uint8_t* stage = smem.stages[s]; + if (contiguous) { + if (lane == 0) { + ptx::bulk_g2s( + stage, + pointer::offset(cur.src_base, static_cast(first) * p.src_stride), + c.rows * p.row_bytes, + &smem.full[s]); + } + } else if (boxed) { + if (lane == 0) { + ptx::bulk_tensor_2d_g2s(stage, &p.src_map[c.is_v ? 1 : 0], 0, static_cast(first), &smem.full[s]); + } + } else { +#pragma unroll + for (uint32_t k = 0; k < kHicacheTmaRowsPerLane; ++k) { + const uint32_t r = lane + k * kWarpThreads; + if (r < c.rows) { + ptx::bulk_g2s( + stage + r * p.row_bytes, + pointer::offset(cur.src_base, static_cast(cur.src[k]) * p.src_stride), + p.row_bytes, + &smem.full[s]); + } + } + } + } + } else { + // ---- storers: smem ring -> global, 16-byte units interleaved across all + // store threads. Row addressing comes from the loader-staged dst_idx. + constexpr uint32_t kUnroll = 4; + constexpr uint32_t kStoreThreads = kStoreWarps * kWarpThreads; + constexpr uint32_t kUnitsPerIter = kStoreThreads * kUnroll; + const uint32_t tid = threadIdx.x - kWarpThreads; + const uint32_t units_per_row = p.row_bytes / 16; + + // unit -> (row, col) without a hardware divide: magic multiply plus a one-step fixup. + auto locate = [&](uint32_t u, uint32_t& row, uint32_t& col) { + row = static_cast((static_cast(u) * p.units_per_row_magic) >> 32); + int32_t rem = static_cast(u - row * units_per_row); + if (rem < 0) { + --row; + rem += units_per_row; + } + col = static_cast(rem); + }; + auto unit_dst = [&](void* dst_base, const int64_t* dst_idx, uint32_t u) -> uint4* { + uint32_t row, col; + locate(u, row, col); + return static_cast(pointer::offset(dst_base, dst_idx[row] * p.dst_stride, col * 16)); + }; + + constexpr uint32_t kNoStage = ~0u; + uint32_t bulk_pending = kNoStage; // stage of the last bulk store not yet released (warp 0) + uint32_t it = 0; + for (uint32_t chunk = blockIdx.x; chunk < num_chunks; chunk += gridDim.x, ++it) { + const ChunkInfo c = describe(chunk); + void* dst_base = const_cast(base_ptr(c.is_v ? p.v_dst : p.k_dst, p.dst_is_table, c.layer)); + const uint32_t s = it % kNumStages; + const uint32_t n_units = c.rows * units_per_row; + const uint32_t n_full = n_units - n_units % kUnitsPerIter; + const uint4* stage = reinterpret_cast(smem.stages[s]); + const int64_t* dst_idx = smem.dst_idx[s]; + + ptx::mbar_wait_parity(&smem.full[s], (it / kNumStages) & 1); + if (smem.dst_run[s]) { + // Contiguous span: one bulk store from store warp 0 runs at the SM's write + // port; it releases the previous bulk stage once that stage's smem read is + // done, keeping two stores in flight. The other store warps have nothing + // to read and release the stage right away. + if (tid < kWarpThreads) { + if (lane == 0) { + ptx::bulk_s2g(pointer::offset(dst_base, dst_idx[0] * p.dst_stride), stage, n_units * 16); + ptx::bulk_commit_group(); + if (bulk_pending != kNoStage) { + ptx::bulk_wait_group_read_one(); + ptx::mbar_arrive(&smem.empty[bulk_pending]); + } + } + bulk_pending = s; + } else if (lane == 0) { + ptx::mbar_arrive(&smem.empty[s]); + } + continue; + } + if (tid < kWarpThreads && bulk_pending != kNoStage) { + if (lane == 0) { + ptx::bulk_wait_group_read_all(); + ptx::mbar_arrive(&smem.empty[bulk_pending]); + } + bulk_pending = kNoStage; + } + for (uint32_t u0 = tid; u0 < n_full; u0 += kUnitsPerIter) { + uint4 v[kUnroll]; + uint4* dst[kUnroll]; +#pragma unroll + for (uint32_t k = 0; k < kUnroll; ++k) { + v[k] = stage[u0 + k * kStoreThreads]; + dst[k] = unit_dst(dst_base, dst_idx, u0 + k * kStoreThreads); + } +#pragma unroll + for (uint32_t k = 0; k < kUnroll; ++k) + __stcs(dst[k], v[k]); + } + for (uint32_t u = n_full + tid; u < n_units; u += kStoreThreads) { + __stcs(unit_dst(dst_base, dst_idx, u), stage[u]); + } + __syncwarp(); + if (lane == 0) ptx::mbar_arrive(&smem.empty[s]); + } + if (tid == 0) ptx::bulk_wait_group_all(); // bulk stores must land before the grid completes + (void)bulk_pending; // the final stage is never reused + } +#endif +} + +template +struct HiCacheTmaKernel { + using Smem = HicacheTmaSmem; + static_assert(kStageBytes % 128 == 0, "stage must stay 128-byte aligned for bulk copies"); + static_assert(kNumStages >= 3 && kStoreWarps >= 1, "two bulk stores in flight plus one loading stage"); + static constexpr uint32_t kThreads = (1 + kStoreWarps) * device::kWarpThreads; + + template + static constexpr auto kernel = hicache_tma_transfer_kernel; + + static uint32_t rows_per_chunk(uint32_t row_bytes) { + return hicache_tma_rows_per_chunk(kStageBytes, row_bytes); + } + + // Whether the device can hold the smem ring in one CTA; sm_90+ parts with + // small opt-in shared memory (consumer Blackwell) must keep the register kernel. + static bool fits_device(int64_t device_id) { + int max_smem = 0; + host::RuntimeDeviceCheck( + cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, static_cast(device_id))); + return static_cast(max_smem) >= sizeof(Smem); + } + + // ceil(2^32 / units_per_row): (u * magic) >> 32 overestimates u / units_per_row + // by at most one for the unit counts a stage can hold; the kernel fixes that up. + static uint64_t units_per_row_magic(uint32_t row_bytes) { + const uint64_t upr = row_bytes / 16; + return ((uint64_t{1} << 32) + upr - 1) / upr; + } + + static auto encode_tiled_fn() -> PFN_cuTensorMapEncodeTiled_v12000 { + static const auto fn = [] { + void* sym = nullptr; + cudaDriverEntryPointQueryResult status; + host::RuntimeDeviceCheck( + cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &sym, 12000, cudaEnableDefault, &status)); + host::RuntimeCheck(status == cudaDriverEntryPointSuccess && sym != nullptr, "cuTensorMapEncodeTiled unavailable"); + return reinterpret_cast(sym); + }(); + return fn; + } + + // 2D view of a strided row buffer: [rows][row_bytes / 8] uint64 elements with + // pitch `stride_bytes`; one box covers rows_per_chunk consecutive rows. + static void encode_src_map(CUtensorMap* map, const void* base, int64_t num_rows, uint32_t row_bytes, int64_t stride) { + const cuuint64_t gdim[2] = {row_bytes / 8, static_cast(num_rows)}; + const cuuint64_t gstride[1] = {static_cast(stride)}; + const cuuint32_t box[2] = {row_bytes / 8, rows_per_chunk(row_bytes)}; + const cuuint32_t estride[2] = {1, 1}; + const CUresult res = encode_tiled_fn()( + map, + CU_TENSOR_MAP_DATA_TYPE_UINT64, + 2, + const_cast(base), + gdim, + gstride, + box, + estride, + CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, + CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + host::RuntimeCheck(res == CUDA_SUCCESS, "cuTensorMapEncodeTiled failed: ", static_cast(res)); + } + + static void launch(const HicacheTmaParams& params, bool use_int32, DLDevice device) { + using namespace host; + RuntimeCheck(params.row_bytes > 0 && params.row_bytes % 16 == 0, "HiCache TMA: row bytes must be a multiple of 16"); + RuntimeCheck(params.row_bytes <= kStageBytes, "HiCache TMA: row bytes exceed the smem stage"); + RuntimeCheck( + params.src_stride % 16 == 0 && params.dst_stride % 16 == 0, "HiCache TMA: strides must be multiples of 16"); + if (params.length == 0 || params.num_layers == 0) return; + + const uint32_t chunks = + (params.has_v ? 2u : 1u) * params.num_layers * div_ceil(params.length, rows_per_chunk(params.row_bytes)); + constexpr std::size_t kSmemBytes = sizeof(Smem); + + static const bool attr_set = [] { + for (auto fn : {kernel, kernel}) { + RuntimeDeviceCheck( + cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast(kSmemBytes))); + } + return true; + }(); + (void)attr_set; + LaunchKernel(std::min(chunks, kBlockQuota), kThreads, device, kSmemBytes)( + use_int32 ? kernel : kernel, params); + } + + // Cache operand viewed as [-1, D] rows; binds row dim, stride and dtype. + static void verify_cache( + const tvm::ffi::TensorView& t, host::SymbolicSize& D, host::SymbolicSize& stride, host::SymbolicDType& dtype) { + using namespace host; + TensorMatcher({-1, D}) // + .with_strides({stride, 1}) + .with_dtype(dtype) + .with_device() + .verify(t); + } + + static void verify_indices( + const tvm::ffi::TensorView& a, + const tvm::ffi::TensorView& b, + host::SymbolicSize& L, + host::SymbolicDType& dtype, + host::SymbolicDevice& device) { + using namespace host; + TensorMatcher({L}) // + .with_dtype(dtype) + .with_device(device) + .verify(a) + .verify(b); + } + + // One layer, direct pointers. `v_*` are ignored when `has_v == false`. + static void run_one_impl( + const tvm::ffi::TensorView k_cache_dst, + const tvm::ffi::TensorView v_cache_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView k_cache_src, + const tvm::ffi::TensorView v_cache_src, + const tvm::ffi::TensorView indices_src, + bool has_v) { + using namespace host; + auto D = SymbolicSize{"row dim"}; + auto N = SymbolicSize{"src stride"}; + auto M = SymbolicSize{"dst stride"}; + auto L = SymbolicSize{"indices length"}; + auto cache_dtype = SymbolicDType{}; + auto indices_dtype = SymbolicDType{}; + auto indices_device = SymbolicDevice{}; + + verify_cache(k_cache_src, D, N, cache_dtype); + verify_cache(k_cache_dst, D, M, cache_dtype); + if (has_v) { + verify_cache(v_cache_src, D, N, cache_dtype); + verify_cache(v_cache_dst, D, M, cache_dtype); + } + verify_indices(indices_src, indices_dst, L, indices_dtype, indices_device); + + const auto dtype_size = dtype_bytes(cache_dtype.unwrap()); + const auto row_bytes = static_cast(D.unwrap() * dtype_size); + const auto src_stride = static_cast(N.unwrap() * dtype_size); + HicacheTmaParams params{ + .k_src = k_cache_src.data_ptr(), + .v_src = has_v ? v_cache_src.data_ptr() : nullptr, + .k_dst = k_cache_dst.data_ptr(), + .v_dst = has_v ? v_cache_dst.data_ptr() : nullptr, + .indices_src = indices_src.data_ptr(), + .indices_dst = indices_dst.data_ptr(), + .src_stride = src_stride, + .dst_stride = static_cast(M.unwrap() * dtype_size), + .row_bytes = row_bytes, + .length = static_cast(L.unwrap()), + .num_layers = 1, + .units_per_row_magic = units_per_row_magic(row_bytes), + .src_is_table = false, + .dst_is_table = false, + .has_v = has_v, + .has_src_map = false, + }; + if (src_stride != row_bytes && row_bytes <= kHicacheTmaMaxMapRowBytes) { + params.has_src_map = true; + encode_src_map(¶ms.src_map[0], params.k_src, k_cache_src.shape()[0], row_bytes, src_stride); + if (has_v) encode_src_map(¶ms.src_map[1], params.v_src, v_cache_src.shape()[0], row_bytes, src_stride); + } + launch(params, indices_dtype.unwrap().bits == 32, indices_device.unwrap()); + } + + // All layers through device-side pointer tables; strides and row bytes explicit. + static void run_all_impl( + const tvm::ffi::TensorView k_ptr_dst, + const tvm::ffi::TensorView v_ptr_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView k_ptr_src, + const tvm::ffi::TensorView v_ptr_src, + const tvm::ffi::TensorView indices_src, + int64_t src_stride_bytes, + int64_t dst_stride_bytes, + int64_t row_bytes, + bool has_v) { + using namespace host; + auto N = SymbolicSize{"num_layers"}; + auto L = SymbolicSize{"indices length"}; + auto indices_dtype = SymbolicDType{}; + auto device_ = SymbolicDevice{}; + + auto verify_table = [&](const tvm::ffi::TensorView& t) { + TensorMatcher({N}).with_dtype().with_device(device_).verify(t); + }; + verify_table(k_ptr_src); + verify_table(k_ptr_dst); + if (has_v) { + verify_table(v_ptr_src); + verify_table(v_ptr_dst); + } + verify_indices(indices_src, indices_dst, L, indices_dtype, device_); + + const HicacheTmaParams params{ + .k_src = k_ptr_src.data_ptr(), + .v_src = has_v ? v_ptr_src.data_ptr() : nullptr, + .k_dst = k_ptr_dst.data_ptr(), + .v_dst = has_v ? v_ptr_dst.data_ptr() : nullptr, + .indices_src = indices_src.data_ptr(), + .indices_dst = indices_dst.data_ptr(), + .src_stride = src_stride_bytes, + .dst_stride = dst_stride_bytes, + .row_bytes = static_cast(row_bytes), + .length = static_cast(L.unwrap()), + .num_layers = static_cast(N.unwrap()), + .units_per_row_magic = units_per_row_magic(static_cast(row_bytes)), + .src_is_table = true, + .dst_is_table = true, + .has_v = has_v, + .has_src_map = false, + }; + launch(params, indices_dtype.unwrap().bits == 32, device_.unwrap()); + } + + static void run_one( + const tvm::ffi::TensorView k_cache_dst, + const tvm::ffi::TensorView v_cache_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView k_cache_src, + const tvm::ffi::TensorView v_cache_src, + const tvm::ffi::TensorView indices_src) { + run_one_impl(k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src, true); + } + + static void run_one_mla( + const tvm::ffi::TensorView cache_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView cache_src, + const tvm::ffi::TensorView indices_src) { + run_one_impl(cache_dst, cache_dst, indices_dst, cache_src, cache_src, indices_src, false); + } + + static void run_all( + const tvm::ffi::TensorView k_ptr_dst, + const tvm::ffi::TensorView v_ptr_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView k_ptr_src, + const tvm::ffi::TensorView v_ptr_src, + const tvm::ffi::TensorView indices_src, + const int64_t src_stride_bytes, + const int64_t dst_stride_bytes, + const int64_t row_bytes) { + run_all_impl( + k_ptr_dst, + v_ptr_dst, + indices_dst, + k_ptr_src, + v_ptr_src, + indices_src, + src_stride_bytes, + dst_stride_bytes, + row_bytes, + true); + } + + static void run_all_mla( + const tvm::ffi::TensorView ptr_dst, + const tvm::ffi::TensorView indices_dst, + const tvm::ffi::TensorView ptr_src, + const tvm::ffi::TensorView indices_src, + const int64_t src_stride_bytes, + const int64_t dst_stride_bytes, + const int64_t row_bytes) { + run_all_impl( + ptr_dst, + ptr_dst, + indices_dst, + ptr_src, + ptr_src, + indices_src, + src_stride_bytes, + dst_stride_bytes, + row_bytes, + false); + } +}; + +} // namespace sglang diff --git a/python/sglang/kernels/ops/kvcache/hicache.py b/python/sglang/kernels/ops/kvcache/hicache.py index 8e454390a..8b3ff3803 100644 --- a/python/sglang/kernels/ops/kvcache/hicache.py +++ b/python/sglang/kernels/ops/kvcache/hicache.py @@ -3,6 +3,8 @@ from __future__ import annotations import logging from typing import TYPE_CHECKING +import torch + from sglang.kernels.jit.utils import ( cache_once, is_hip_runtime, @@ -10,9 +12,9 @@ from sglang.kernels.jit.utils import ( make_cpp_args, ) from sglang.kernels.kernel_api_logging import debug_kernel_api +from sglang.srt.environ import envs if TYPE_CHECKING: - import torch from tvm_ffi.module import Module _is_hip = is_hip_runtime() @@ -81,13 +83,81 @@ def _jit_hicache_staged_module( ) +# TMA staging ring per CTA: 32 KB stages x 6 keeps the host loads in flight +# under a 4-block launch; smaller stages make the loader's per-chunk cost the +# limit. 4 store warps drain a strided-destination stage faster than it fills. +TMA_STAGE_BYTES = 32 * 1024 +TMA_NUM_STAGES = 6 +TMA_STORE_WARPS = 4 +# Each CTA is capped by its SM's write port, so the host link needs four of +# them; the register kernel keeps DEFAULT_BLOCK_QUOTA. +TMA_BLOCK_QUOTA = 4 + + +@cache_once +def _jit_hicache_tma_module(*, block_quota: int) -> Module: + args = make_cpp_args(TMA_STAGE_BYTES, TMA_NUM_STAGES, TMA_STORE_WARPS, block_quota) + return load_jit( + "hicache_tma", + *args, + cuda_files=["kvcacheio/hicache_tma.cuh"], + cuda_wrappers=[ + ("launch_one", f"&HiCacheTmaKernel<{args}>::run_one"), + ("launch_all", f"&HiCacheTmaKernel<{args}>::run_all"), + ("launch_one_mla", f"&HiCacheTmaKernel<{args}>::run_one_mla"), + ("launch_all_mla", f"&HiCacheTmaKernel<{args}>::run_all_mla"), + ("fits_device", f"&HiCacheTmaKernel<{args}>::fits_device"), + ], + ) + + +def hicache_tma_rows_per_chunk(element_size: int) -> int: + """Rows the TMA kernel moves per stage (mirrors hicache_tma_rows_per_chunk).""" + rows = 1 + while rows * 2 <= TMA_STAGE_BYTES // element_size and rows * 2 <= 128: + rows *= 2 + return rows + + +@cache_once +def use_hicache_tma_kernel( + *, element_size: int, block_quota: int, page_size: int | None = None +) -> bool: + """Whether transfers of `element_size`-byte rows go through the TMA kernel. + + A chunk that straddles pages degrades to one bulk op per row, so pages must + tile the chunk; `page_size=None` (caller unaware of paging) trusts the indices. + """ + if _is_hip or not envs.SGLANG_HICACHE_TMA_TRANSFER.get(): + return False + if element_size % 16 != 0 or torch.cuda.get_device_capability()[0] < 9: + return False + if page_size is not None and page_size % hicache_tma_rows_per_chunk(element_size): + return False + try: + module = _jit_hicache_tma_module(block_quota=block_quota) + except Exception as e: + logging.getLogger(__name__).warning( + f"Failed to load the TMA HiCache kernel, using the register kernel: {e}" + ) + return False + return bool(module.fits_device(torch.cuda.current_device())) + + def can_use_hicache_jit_kernel( *, element_size: int, unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference + page_size: int | None = None, ) -> bool: logger = logging.getLogger(__name__) + if use_hicache_tma_kernel( + element_size=element_size, + block_quota=block_quota or TMA_BLOCK_QUOTA, + page_size=page_size, + ): + return True unroll = unroll or _default_unroll(element_size) if not _tiles_across_lanes(element_size, unroll): logger.warning(f"Unsupported {element_size = } for JIT HiCache kernel") @@ -166,6 +236,7 @@ def transfer_hicache_one_layer( element_dim: int | None = None, unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference + page_size: int | None = None, ) -> None: element_dim = element_dim or k_cache_dst.size(-1) k_cache_src = k_cache_src.view(-1, element_dim) @@ -173,13 +244,19 @@ def transfer_hicache_one_layer( k_cache_dst = k_cache_dst.view(-1, element_dim) v_cache_dst = v_cache_dst.view(-1, element_dim) element_size = element_dim * k_cache_dst.element_size() - block_quota = block_quota or DEFAULT_BLOCK_QUOTA - unroll = unroll or _default_unroll(element_size) - module = _jit_hicache_module( - element_size=element_size, - unroll=unroll, - block_quota=block_quota, - ) + tma_quota = block_quota or TMA_BLOCK_QUOTA + if use_hicache_tma_kernel( + element_size=element_size, block_quota=tma_quota, page_size=page_size + ): + module = _jit_hicache_tma_module(block_quota=tma_quota) + else: + block_quota = block_quota or DEFAULT_BLOCK_QUOTA + unroll = unroll or _default_unroll(element_size) + module = _jit_hicache_module( + element_size=element_size, + unroll=unroll, + block_quota=block_quota, + ) module.launch_one( k_cache_dst, v_cache_dst, @@ -204,11 +281,28 @@ def transfer_hicache_all_layer( element_size: int | None = None, unroll: int | None = None, # can be tuned for performance block_quota: int | None = None, # can be tuned for less interference + page_size: int | None = None, ) -> None: if element_size is None: # assume both contiguous assert kv_cache_dst_stride_bytes == kv_cache_src_stride_bytes element_size = kv_cache_dst_stride_bytes + tma_quota = block_quota or TMA_BLOCK_QUOTA + if use_hicache_tma_kernel( + element_size=element_size, block_quota=tma_quota, page_size=page_size + ): + _jit_hicache_tma_module(block_quota=tma_quota).launch_all( + k_ptr_dst, + v_ptr_dst, + indices_dst, + k_ptr_src, + v_ptr_src, + indices_src, + kv_cache_src_stride_bytes, + kv_cache_dst_stride_bytes, + element_size, + ) + return block_quota = block_quota or DEFAULT_BLOCK_QUOTA unroll = unroll or _default_unroll(element_size) module = _jit_hicache_module( @@ -237,18 +331,25 @@ def transfer_hicache_one_layer_mla( element_dim: int | None = None, unroll: int | None = None, block_quota: int | None = None, + page_size: int | None = None, ) -> None: element_dim = element_dim or cache_dst.size(-1) cache_src = cache_src.view(-1, element_dim) cache_dst = cache_dst.view(-1, element_dim) element_size = element_dim * cache_dst.element_size() - block_quota = block_quota or DEFAULT_BLOCK_QUOTA - unroll = unroll or _default_unroll(element_size) - module = _jit_hicache_module( - element_size=element_size, - unroll=unroll, - block_quota=block_quota, - ) + tma_quota = block_quota or TMA_BLOCK_QUOTA + if use_hicache_tma_kernel( + element_size=element_size, block_quota=tma_quota, page_size=page_size + ): + module = _jit_hicache_tma_module(block_quota=tma_quota) + else: + block_quota = block_quota or DEFAULT_BLOCK_QUOTA + unroll = unroll or _default_unroll(element_size) + module = _jit_hicache_module( + element_size=element_size, + unroll=unroll, + block_quota=block_quota, + ) module.launch_one_mla( cache_dst, indices_dst, @@ -268,11 +369,26 @@ def transfer_hicache_all_layer_mla( element_size: int | None = None, unroll: int | None = None, block_quota: int | None = None, + page_size: int | None = None, ) -> None: if element_size is None: assert cache_dst_stride_bytes == cache_src_stride_bytes element_size = cache_dst_stride_bytes + tma_quota = block_quota or TMA_BLOCK_QUOTA + if use_hicache_tma_kernel( + element_size=element_size, block_quota=tma_quota, page_size=page_size + ): + _jit_hicache_tma_module(block_quota=tma_quota).launch_all_mla( + ptr_dst, + indices_dst, + ptr_src, + indices_src, + cache_src_stride_bytes, + cache_dst_stride_bytes, + element_size, + ) + return block_quota = block_quota or DEFAULT_BLOCK_QUOTA unroll = unroll or _default_unroll(element_size) module = _jit_hicache_module( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 19d91d52e..b497b0c3c 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -750,6 +750,9 @@ class Envs: # =================================================================== # Per-call cudaHostRegister limit in GB. SGLANG_HICACHE_HOST_REGISTER_CHUNK_GB = EnvInt(256) + # HiCache host<->device transfers use the TMA staging kernel when the GPU + # (sm_90+), row size and page size allow; set to 0 to force the register kernel. + SGLANG_HICACHE_TMA_TRANSFER = EnvBool(True) # Base token count for each MLA/DSA dedup broadcast chunk. SGLANG_MLA_DEDUP_CHUNK_TOKENS = EnvInt(2048) SGLANG_HICACHE_HF3FS_CONFIG_PATH = EnvStr(None) diff --git a/python/sglang/srt/mem_cache/pool_host/mha.py b/python/sglang/srt/mem_cache/pool_host/mha.py index f55904b61..371e4eec0 100644 --- a/python/sglang/srt/mem_cache/pool_host/mha.py +++ b/python/sglang/srt/mem_cache/pool_host/mha.py @@ -108,7 +108,8 @@ class MHATokenToKVPoolHost(HostKVCache): # write-back kernel has a ROCm path, so enable them on HIP too. This # keeps the ROCm write-back path consistent with CUDA. self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel( - element_size=self.element_dim * self.dtype.itemsize + page_size=self.page_size, + element_size=self.element_dim * self.dtype.itemsize, ) if self.layout == "page_first": @@ -276,6 +277,7 @@ class MHATokenToKVPoolHost(HostKVCache): if self.layout == "layer_first": if self.can_use_jit: jit_transfer_hicache_one_layer( + page_size=self.page_size, k_cache_dst=device_pool.k_buffer[device_layer_id], v_cache_dst=device_pool.v_buffer[device_layer_id], k_cache_src=self.k_buffer[host_layer_id], @@ -300,6 +302,7 @@ class MHATokenToKVPoolHost(HostKVCache): # index by layer_id to get a per-layer view with strided layout. # The kernel handles different src/dst strides automatically. jit_transfer_hicache_one_layer( + page_size=self.page_size, k_cache_dst=device_pool.k_buffer[device_layer_id], v_cache_dst=device_pool.v_buffer[device_layer_id], k_cache_src=self.k_data_refs[host_layer_id], @@ -437,6 +440,7 @@ class MHATokenToKVPoolHost(HostKVCache): if self.layout == "layer_first": if self.can_use_jit: jit_transfer_hicache_all_layer( + page_size=self.page_size, k_ptr_dst=self.k_data_ptrs, v_ptr_dst=self.v_data_ptrs, indices_dst=host_indices, @@ -783,7 +787,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): self.clear() self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel( - element_size=self.token_stride_size + page_size=self.page_size, element_size=self.token_stride_size ) self.k_device_ptrs = torch.tensor( [x.data_ptr() for x in self.device_pool.k_buffer], @@ -855,6 +859,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): if self.layout == "layer_first": if self.can_use_jit: jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=device_pool.k_buffer[layer_id], cache_src=self.k_buffer[layer_id], indices_dst=device_indices, @@ -872,6 +877,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): elif self.layout == "page_first": if self.can_use_jit: jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=device_pool.k_buffer[layer_id], cache_src=self.k_data_refs[layer_id], indices_dst=device_indices, @@ -921,6 +927,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): if self.can_use_jit: for layer_id in range(self.layer_num): jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=self.k_buffer[layer_id], cache_src=device_pool.k_buffer[layer_id], indices_dst=host_indices, @@ -939,6 +946,7 @@ class MHATokenToKOnlyPoolHost(HostKVCache): elif self.layout == "page_first": if self.can_use_jit: jit_transfer_hicache_all_layer_mla( + page_size=self.page_size, ptr_dst=self.k_data_ptrs, indices_dst=host_indices, ptr_src=self.k_device_ptrs, diff --git a/python/sglang/srt/mem_cache/pool_host/mla.py b/python/sglang/srt/mem_cache/pool_host/mla.py index 8898af3cc..86abe18c9 100644 --- a/python/sglang/srt/mem_cache/pool_host/mla.py +++ b/python/sglang/srt/mem_cache/pool_host/mla.py @@ -123,7 +123,8 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): # write-back kernel has a ROCm path, so enable them on HIP too. This # keeps the ROCm write-back path consistent with CUDA. self.can_use_jit = (_is_cuda or _is_hip) and can_use_hicache_jit_kernel( - element_size=self.kv_cache_dim * self.dtype.itemsize + page_size=self.page_size, + element_size=self.kv_cache_dim * self.dtype.itemsize, ) if self.layout in ("page_first", "page_first_kv_split"): @@ -668,6 +669,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): if self.layout == "layer_first": if self.can_use_jit: jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=device_pool.kv_buffer[device_layer_id], cache_src=self.kv_buffer[host_layer_id], indices_dst=device_indices, @@ -685,6 +687,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): elif self.layout == "page_first": if self.can_use_jit: jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=device_pool.kv_buffer[device_layer_id], cache_src=self.data_refs[host_layer_id], indices_dst=device_indices, @@ -794,6 +797,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): if self.layout == "layer_first": if self.can_use_jit: jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=self.kv_buffer[host_layer_id], cache_src=device_pool.kv_buffer[device_layer_id], indices_dst=host_indices, @@ -811,6 +815,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): elif self.layout == "page_first": if self.can_use_jit: jit_transfer_hicache_one_layer_mla( + page_size=self.page_size, cache_dst=self.data_refs[host_layer_id], cache_src=device_pool.kv_buffer[device_layer_id], indices_dst=host_indices, @@ -893,6 +898,7 @@ class MLATokenToKVPoolHost(HiSparseHostPoolMixin, HostKVCache): if self.layout == "layer_first": if self.can_use_jit: jit_transfer_hicache_all_layer_mla( + page_size=self.page_size, ptr_dst=self.data_ptrs, indices_dst=host_indices, ptr_src=device_data_ptrs, diff --git a/test/registered/kernels/benchmark/kvcache/bench_hicache.py b/test/registered/kernels/benchmark/kvcache/bench_hicache.py index 6d3aec687..c5de7ec21 100644 --- a/test/registered/kernels/benchmark/kvcache/bench_hicache.py +++ b/test/registered/kernels/benchmark/kvcache/bench_hicache.py @@ -24,8 +24,11 @@ from sgl_kernel import transfer_kv_all_layer, transfer_kv_per_layer from sglang.kernels.jit.benchmark import marker from sglang.kernels.jit.benchmark.utils import get_benchmark_range from sglang.kernels.ops.kvcache.hicache import ( - transfer_hicache_all_layer, - transfer_hicache_one_layer, + DEFAULT_BLOCK_QUOTA, + TMA_BLOCK_QUOTA, + _default_unroll, + _jit_hicache_module, + _jit_hicache_tma_module, ) from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci @@ -108,15 +111,19 @@ def sglang_jit_transfer_one( indices_src: torch.Tensor, element_dim: int, ) -> None: - """SGL JIT Kernel for single layer transfer.""" - transfer_hicache_one_layer( - k_cache_dst, - v_cache_dst, + """SGL JIT register kernel for single layer transfer (bypasses TMA routing).""" + element_size = element_dim * k_cache_dst.element_size() + _jit_hicache_module( + element_size=element_size, + unroll=_default_unroll(element_size), + block_quota=DEFAULT_BLOCK_QUOTA, + ).launch_one( + k_cache_dst.view(-1, element_dim), + v_cache_dst.view(-1, element_dim), indices_dst, - k_cache_src, - v_cache_src, + k_cache_src.view(-1, element_dim), + v_cache_src.view(-1, element_dim), indices_src, - element_dim=element_dim, ) @@ -153,17 +160,58 @@ def sglang_jit_transfer_all( stride_bytes: int, element_size: int, ) -> None: - """SGL JIT Kernel for all layer transfer.""" - transfer_hicache_all_layer( + """SGL JIT register kernel for all layer transfer (bypasses TMA routing).""" + _jit_hicache_module( + element_size=element_size, + unroll=_default_unroll(element_size), + block_quota=DEFAULT_BLOCK_QUOTA, + ).launch_all( k_ptrs_dst, v_ptrs_dst, indices_dst, k_ptrs_src, v_ptrs_src, indices_src, - kv_cache_src_stride_bytes=stride_bytes, - kv_cache_dst_stride_bytes=stride_bytes, - element_size=element_size, + stride_bytes, + stride_bytes, + ) + + +def sglang_tma_transfer_one( + k_cache_dst: torch.Tensor, + v_cache_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_cache_src: torch.Tensor, + v_cache_src: torch.Tensor, + indices_src: torch.Tensor, +) -> None: + """SGL TMA staging kernel for single layer transfer.""" + _jit_hicache_tma_module(block_quota=TMA_BLOCK_QUOTA).launch_one( + k_cache_dst, v_cache_dst, indices_dst, k_cache_src, v_cache_src, indices_src + ) + + +def sglang_tma_transfer_all( + k_ptrs_dst: torch.Tensor, + v_ptrs_dst: torch.Tensor, + indices_dst: torch.Tensor, + k_ptrs_src: torch.Tensor, + v_ptrs_src: torch.Tensor, + indices_src: torch.Tensor, + stride_bytes: int, + element_size: int, +) -> None: + """SGL TMA staging kernel for all layer transfer.""" + _jit_hicache_tma_module(block_quota=TMA_BLOCK_QUOTA).launch_all( + k_ptrs_dst, + v_ptrs_dst, + indices_dst, + k_ptrs_src, + v_ptrs_src, + indices_src, + stride_bytes, + stride_bytes, + element_size, ) @@ -191,6 +239,13 @@ ELEMENT_SIZE_RANGE = get_benchmark_range( LINE_VALS = ["aot", "jit", "torch"] if DISABLE_TORCH: LINE_VALS.remove("torch") +# The TMA staging kernel needs sm_90+ (cp.async.bulk); skip the line elsewhere. +if ( + torch.cuda.is_available() + and torch.version.hip is None + and torch.cuda.get_device_capability()[0] >= 9 +): + LINE_VALS.insert(2, "tma") # ============================================================================= @@ -246,6 +301,17 @@ def benchmark_one_layer_h2d(element_size: int, batch_size: int, provider: str): ) for i in range(NUM_LAYERS) ], + "tma": lambda: [ + sglang_tma_transfer_one( + k_cache_dst[i], + v_cache_dst[i], + indices_dst_gpu, + k_cache_src[i], + v_cache_src[i], + indices_src_gpu, + ) + for i in range(NUM_LAYERS) + ], "torch": lambda: [ pytorch_transfer( k_cache_dst[i], @@ -329,6 +395,16 @@ def benchmark_all_layer_d2h(element_size: int, batch_size: int, provider: str): element_bytes, element_bytes, ), + "tma": lambda: sglang_tma_transfer_all( + k_ptrs_dst, + v_ptrs_dst, + indices_dst_gpu, + k_ptrs_src, + v_ptrs_src, + indices_src_gpu, + element_bytes, + element_bytes, + ), "torch": lambda: [ pytorch_transfer( k_caches_dst[i], diff --git a/test/registered/kernels/ops/kvcache/test_hicache_tma.py b/test/registered/kernels/ops/kvcache/test_hicache_tma.py new file mode 100644 index 000000000..1e30dc6ab --- /dev/null +++ b/test/registered/kernels/ops/kvcache/test_hicache_tma.py @@ -0,0 +1,173 @@ +import sys + +import pytest +import torch + +from sglang.kernels.ops.kvcache.hicache import _jit_hicache_tma_module +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() + or torch.version.hip is not None + or torch.cuda.get_device_capability()[0] < 9, + reason="HiCache TMA kernel requires SM90+", +) + +POOL_TOKENS = 8192 +NUM_LAYERS = 3 +ROW_DIM = 256 # 512-byte bf16 rows: below the register kernel's 128 B unit width x 4 + + +def _token_indices(num_tokens: int, page_size: int, dtype: torch.dtype, seed: int): + gen = torch.Generator().manual_seed(seed) + pages = torch.randperm(POOL_TOKENS // page_size, generator=gen)[ + : num_tokens // page_size + ] + idx = (pages[:, None] * page_size + torch.arange(page_size)).reshape(-1) + return idx.to(device="cuda", dtype=dtype) + + +def _fill(t: torch.Tensor, seed: int) -> None: + t.view(torch.int16).copy_( + torch.randint( + 0, + 30000, + t.shape, + dtype=torch.int16, + generator=torch.Generator().manual_seed(seed), + ) + ) + + +def _host_view(layout: str, layer: int): + if layout == "layer_first": + return torch.empty(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True) + # page_first: [tokens, layers, dim]; a per-layer view has strided rows + return torch.empty( + POOL_TOKENS, NUM_LAYERS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True + )[:, layer] + + +@pytest.mark.parametrize("host_layout", ["layer_first", "page_first"]) +@pytest.mark.parametrize("index_dtype", [torch.int64, torch.int32]) +@pytest.mark.parametrize("page_size", [128, 1]) +def test_one_layer_roundtrip( + host_layout: str, index_dtype: torch.dtype, page_size: int +) -> None: + """H2D then D2H of one layer; page runs take the single-op paths (bulk copy, + tensor-map box, bulk store), scattered rows take the per-row paths, and the + odd token count leaves a partial tail chunk.""" + module = _jit_hicache_tma_module(block_quota=2) + num_tokens = 2048 + (96 if page_size == 1 else 0) + k_host, v_host = _host_view(host_layout, 1), _host_view(host_layout, 2) + k_dev = torch.zeros(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, device="cuda") + v_dev = torch.zeros_like(k_dev) + _fill(k_host, 1) + _fill(v_host, 2) + host_idx = _token_indices(num_tokens, page_size, index_dtype, seed=3) + dev_idx = _token_indices(num_tokens, page_size, index_dtype, seed=4) + + module.launch_one(k_dev, v_dev, dev_idx, k_host, v_host, host_idx) + torch.cuda.synchronize() + assert torch.equal(k_dev[dev_idx.long()].cpu(), k_host[host_idx.cpu().long()]) + assert torch.equal(v_dev[dev_idx.long()].cpu(), v_host[host_idx.cpu().long()]) + untouched = torch.ones(POOL_TOKENS, dtype=torch.bool, device="cuda") + untouched[dev_idx.long()] = False + assert not k_dev[untouched].any() and not v_dev[untouched].any() + + _fill(k_dev, 5) + _fill(v_dev, 6) + k_host.zero_() + v_host.zero_() + module.launch_one(k_host, v_host, host_idx, k_dev, v_dev, dev_idx) + torch.cuda.synchronize() + assert torch.equal(k_host[host_idx.cpu().long()], k_dev[dev_idx.long()].cpu()) + assert torch.equal(v_host[host_idx.cpu().long()], v_dev[dev_idx.long()].cpu()) + + +def _ptr_table(tensors) -> torch.Tensor: + return torch.tensor( + [t.data_ptr() for t in tensors], dtype=torch.uint64, device="cuda" + ) + + +def test_all_layer_tables_lf_to_pf() -> None: + """All-layer D2H through per-layer pointer tables into a page-first host pool + (strided destination rows), the write-back shape.""" + module = _jit_hicache_tma_module(block_quota=2) + k_dev = [ + torch.empty(POOL_TOKENS, ROW_DIM, dtype=torch.bfloat16, device="cuda") + for _ in range(NUM_LAYERS) + ] + v_dev = [torch.empty_like(k_dev[0]) for _ in range(NUM_LAYERS)] + for i, t in enumerate(k_dev + v_dev): + _fill(t, 10 + i) + k_host = torch.zeros( + POOL_TOKENS, NUM_LAYERS, ROW_DIM, dtype=torch.bfloat16, pin_memory=True + ) + v_host = torch.zeros_like(k_host).pin_memory() + host_idx = _token_indices(2048, 128, torch.int64, seed=7) + dev_idx = _token_indices(2048, 128, torch.int64, seed=8) + row_bytes = ROW_DIM * 2 + + module.launch_all( + _ptr_table([k_host[:, l] for l in range(NUM_LAYERS)]), + _ptr_table([v_host[:, l] for l in range(NUM_LAYERS)]), + host_idx, + _ptr_table(k_dev), + _ptr_table(v_dev), + dev_idx, + row_bytes, + NUM_LAYERS * row_bytes, + row_bytes, + ) + torch.cuda.synchronize() + for l in range(NUM_LAYERS): + assert torch.equal(k_host[host_idx.cpu(), l], k_dev[l][dev_idx].cpu()) + assert torch.equal(v_host[host_idx.cpu(), l], v_dev[l][dev_idx].cpu()) + untouched = torch.ones(POOL_TOKENS, dtype=torch.bool) + untouched[host_idx.cpu()] = False + assert not k_host[untouched].any() and not v_host[untouched].any() + + +def test_mla_single_buffer() -> None: + """MLA rows (576 x bf16 = 1152 B, not a multiple of 128 B) through the + single-buffer entry points, one layer and all layers.""" + module = _jit_hicache_tma_module(block_quota=2) + dim = 576 + dev = [ + torch.empty(POOL_TOKENS, dim, dtype=torch.bfloat16, device="cuda") + for _ in range(NUM_LAYERS) + ] + for i, t in enumerate(dev): + _fill(t, 20 + i) + host = torch.zeros( + POOL_TOKENS, NUM_LAYERS, dim, dtype=torch.bfloat16, pin_memory=True + ) + host_idx = _token_indices(1024, 128, torch.int64, seed=9) + dev_idx = _token_indices(1024, 128, torch.int64, seed=10) + row_bytes = dim * 2 + + module.launch_all_mla( + _ptr_table([host[:, l] for l in range(NUM_LAYERS)]), + host_idx, + _ptr_table(dev), + dev_idx, + row_bytes, + NUM_LAYERS * row_bytes, + row_bytes, + ) + torch.cuda.synchronize() + for l in range(NUM_LAYERS): + assert torch.equal(host[host_idx.cpu(), l], dev[l][dev_idx].cpu()) + + dev[0].zero_() + module.launch_one_mla(dev[0], dev_idx, host[:, 0], host_idx) + torch.cuda.synchronize() + assert torch.equal(dev[0][dev_idx].cpu(), host[host_idx.cpu(), 0]) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v", "-s"]))