[HiCache] TMA-staged host<->device KV transfer kernel (sm_90+) (#40278)

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 <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/mbarrier.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/warp.cuh>
+
+#include <dlpack/dlpack.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cuda.h>
+#include <cudaTypedefs.h>
+
+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 <uint32_t kStageBytes, uint32_t kNumStages>
+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 <typename T, uint32_t kStageBytes, uint32_t kNumStages, uint32_t kStoreWarps>
+__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<kStageBytes, kNumStages>;
+  extern __shared__ __align__(128) uint8_t smem_raw[];
+  auto& smem = *reinterpret_cast<Smem*>(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<const T*>(p.indices_src);
+  const T* idx_dst = static_cast<const T*>(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<const void*>(static_cast<const uint64_t*>(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<T>(r);
+        run_dst &= r >= c.rows || cur.dst[k] == first_dst + static_cast<T>(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<int64_t>(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<int64_t>(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<int32_t>(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<int64_t>(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<uint32_t>((static_cast<uint64_t>(u) * p.units_per_row_magic) >> 32);
+      int32_t rem = static_cast<int32_t>(u - row * units_per_row);
+      if (rem < 0) {
+        --row;
+        rem += units_per_row;
+      }
+      col = static_cast<uint32_t>(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<uint4*>(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<void*>(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<const uint4*>(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 <uint32_t kStageBytes, uint32_t kNumStages, uint32_t kStoreWarps, uint32_t kBlockQuota>
+struct HiCacheTmaKernel {
+  using Smem = HicacheTmaSmem<kStageBytes, kNumStages>;
+  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 <typename T>
+  static constexpr auto kernel = hicache_tma_transfer_kernel<T, kStageBytes, kNumStages, kStoreWarps>;
+
+  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<int>(device_id)));
+    return static_cast<std::size_t>(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<PFN_cuTensorMapEncodeTiled_v12000>(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<cuuint64_t>(num_rows)};
+    const cuuint64_t gstride[1] = {static_cast<cuuint64_t>(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<void*>(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<int>(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<int32_t>, kernel<int64_t>}) {
+        RuntimeDeviceCheck(
+            cudaFuncSetAttribute(fn, cudaFuncAttributeMaxDynamicSharedMemorySize, static_cast<int>(kSmemBytes)));
+      }
+      return true;
+    }();
+    (void)attr_set;
+    LaunchKernel(std::min(chunks, kBlockQuota), kThreads, device, kSmemBytes)(
+        use_int32 ? kernel<int32_t> : kernel<int64_t>, 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<kDLGPU, kDLGPUHost, kDLCPU>()
+        .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<int32_t, int64_t>(dtype)
+        .with_device<kDLGPU>(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<uint32_t>(D.unwrap() * dtype_size);
+    const auto src_stride = static_cast<int64_t>(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<int64_t>(M.unwrap() * dtype_size),
+        .row_bytes = row_bytes,
+        .length = static_cast<uint32_t>(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(&params.src_map[0], params.k_src, k_cache_src.shape()[0], row_bytes, src_stride);
+      if (has_v) encode_src_map(&params.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<uint64_t>().with_device<kDLGPU>(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<uint32_t>(row_bytes),
+        .length = static_cast<uint32_t>(L.unwrap()),
+        .num_layers = static_cast<uint32_t>(N.unwrap()),
+        .units_per_row_magic = units_per_row_magic(static_cast<uint32_t>(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

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(

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)

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,

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,

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],

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"]))