[PD] Add optional KV transfer checksums (#39500)

python/sglang/kernels/jit/csrc/elementwise/adler32_checksum.cuh

@@ -0,0 +1,362 @@
+// Adler-32 checksum CUDA kernels for tensor verification.
+//
+// Two modes:
+//   1. Whole-tensor: checksum all bytes of a contiguous tensor
+//   2. Strided: checksum selected items across multiple tensors,
+//      given precomputed byte pointers and per-item lengths
+//
+// Adler-32:
+//   A = 1 + sum(bytes)              (mod 65521)
+//   B = sum of running A values     (mod 65521)
+//   checksum = (B << 16) | A
+//
+// Combine (left || right):
+//   A = A_left + A_right - 1             (mod 65521)
+//   B = B_left + B_right + len_right * (A_left - 1)  (mod 65521)
+
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstddef>
+#include <cstdint>
+
+// ROCm's masked __shfl_*_sync traps on partial wavefronts (fewer than 64 lanes
+// launched), so use the maskless intrinsic on ROCm.
+#ifdef USE_ROCM
+#define ADLER32_SHFL_DOWN(mask, var, delta) __shfl_down((var), (delta))
+#else
+#define ADLER32_SHFL_DOWN(mask, var, delta) __shfl_down_sync((mask), (var), (delta))
+#endif
+
+namespace sglang {
+
+namespace {
+
+constexpr uint32_t MOD_ADLER = 65521;
+constexpr size_t kBlockSize = 256;
+
+struct Adler32State {
+  uint32_t a;
+  uint32_t b;
+  uint64_t len;
+};
+
+__device__ __forceinline__ Adler32State adler32_identity() {
+  return {1, 0, 0};
+}
+
+__device__ Adler32State adler32_compute(const uint8_t* ptr, uint64_t nbytes) {
+  uint64_t sum = 0;
+  uint64_t weighted = 0;
+  uint64_t offset = 0;
+  while (offset < nbytes && (reinterpret_cast<uintptr_t>(ptr + offset) & 15) != 0) {
+    uint64_t value = ptr[offset];
+    sum += value;
+    weighted += (nbytes - offset) * value;
+    ++offset;
+  }
+  const uint4* vectors = reinterpret_cast<const uint4*>(ptr + offset);
+  uint64_t num_vectors = (nbytes - offset) / 16;
+  for (uint64_t i = 0; i < num_vectors; ++i) {
+    uint4 value = vectors[i];
+    uint32_t words[4] = {value.x, value.y, value.z, value.w};
+    uint64_t word_sums[4];
+    uint64_t within = 0;
+#pragma unroll
+    for (uint32_t word = 0; word < 4; ++word) {
+      uint32_t value_word = words[word];
+      uint64_t x0 = value_word & 0xffu;
+      uint64_t x1 = (value_word >> 8) & 0xffu;
+      uint64_t x2 = (value_word >> 16) & 0xffu;
+      uint64_t x3 = value_word >> 24;
+      word_sums[word] = x0 + x1 + x2 + x3;
+      within += 3 * x0 + 2 * x1 + x2;
+    }
+    uint64_t vector_sum = word_sums[0] + word_sums[1] + word_sums[2] + word_sums[3];
+    within += 12 * word_sums[0] + 8 * word_sums[1] + 4 * word_sums[2];
+    uint64_t base = offset + i * 16;
+    sum += vector_sum;
+    weighted += (nbytes - base - 15) * vector_sum + within;
+    if ((i & 4095) == 4095) {
+      sum %= MOD_ADLER;
+      weighted %= MOD_ADLER;
+    }
+  }
+  offset += num_vectors * 16;
+  while (offset < nbytes) {
+    uint64_t value = ptr[offset];
+    sum += value;
+    weighted += (nbytes - offset) * value;
+    ++offset;
+  }
+  return {
+      static_cast<uint32_t>((1 + sum) % MOD_ADLER),
+      static_cast<uint32_t>((nbytes % MOD_ADLER + weighted) % MOD_ADLER),
+      nbytes};
+}
+
+__device__ __forceinline__ Adler32State adler32_combine(Adler32State left, Adler32State right) {
+  uint64_t a = ((uint64_t)left.a + (uint64_t)right.a - 1 + MOD_ADLER) % MOD_ADLER;
+  uint64_t len_mod = (uint64_t)right.len % MOD_ADLER;
+  uint64_t a_minus_1 = ((uint64_t)left.a - 1 + MOD_ADLER) % MOD_ADLER;
+  uint64_t b = ((uint64_t)left.b + (uint64_t)right.b + len_mod * a_minus_1) % MOD_ADLER;
+  return {(uint32_t)a, (uint32_t)b, left.len + right.len};
+}
+
+__device__ Adler32State adler32_block_reduce(Adler32State state, Adler32State* smem, uint32_t num_valid) {
+#ifdef USE_ROCM
+  // ROCm wavefront is 64 lanes; __shfl_*_sync requires a 64-bit mask.
+  constexpr uint64_t kFullMask = 0xffffffffffffffffULL;
+#else
+  constexpr uint32_t kFullMask = 0xffffffffu;
+#endif
+  uint32_t lane = threadIdx.x % warpSize;
+  uint32_t warp_id = threadIdx.x / warpSize;
+  uint32_t num_warps = (num_valid + warpSize - 1) / warpSize;
+
+  // Phase 1: intra-warp ordered reduction via __shfl_down_sync
+  uint32_t valid_in_warp = warpSize;
+  if (warp_id == num_warps - 1) {
+    uint32_t remainder = num_valid % warpSize;
+    if (remainder != 0) valid_in_warp = remainder;
+  }
+
+#pragma unroll
+  for (uint32_t delta = 1; delta < warpSize; delta *= 2) {
+    Adler32State right;
+    right.a = ADLER32_SHFL_DOWN(kFullMask, state.a, delta);
+    right.b = ADLER32_SHFL_DOWN(kFullMask, state.b, delta);
+    uint32_t len_lo = ADLER32_SHFL_DOWN(kFullMask, (uint32_t)state.len, delta);
+    uint32_t len_hi = ADLER32_SHFL_DOWN(kFullMask, (uint32_t)(state.len >> 32), delta);
+    right.len = (uint64_t)len_lo | ((uint64_t)len_hi << 32);
+    if (lane + delta < valid_in_warp) {
+      state = adler32_combine(state, right);
+    }
+  }
+
+  // Lane 0 of each warp writes result to shared memory
+  if (lane == 0) {
+    smem[warp_id] = state;
+  }
+  __syncthreads();
+
+  // Phase 2: warp 0 reduces across warp results
+  if (warp_id == 0) {
+    state = (lane < num_warps) ? smem[lane] : adler32_identity();
+#pragma unroll
+    for (uint32_t delta = 1; delta < warpSize; delta *= 2) {
+      Adler32State right;
+      right.a = ADLER32_SHFL_DOWN(kFullMask, state.a, delta);
+      right.b = ADLER32_SHFL_DOWN(kFullMask, state.b, delta);
+      uint32_t len_lo = ADLER32_SHFL_DOWN(kFullMask, (uint32_t)state.len, delta);
+      uint32_t len_hi = ADLER32_SHFL_DOWN(kFullMask, (uint32_t)(state.len >> 32), delta);
+      right.len = (uint64_t)len_lo | ((uint64_t)len_hi << 32);
+      if (lane + delta < num_warps) {
+        state = adler32_combine(state, right);
+      }
+    }
+  }
+
+  return state;
+}
+
+// --- Kernel 1: Whole tensor ---
+
+__global__ void adler32_whole_kernel(
+    const uint8_t* __restrict__ data,
+    uint32_t* __restrict__ out_a,
+    uint32_t* __restrict__ out_b,
+    uint64_t* __restrict__ out_len,
+    uint64_t num_bytes,
+    uint64_t chunk_size) {
+  extern __shared__ Adler32State smem[];
+
+  uint64_t global_tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x;
+  uint64_t start = global_tid * chunk_size;
+
+  Adler32State state;
+  if (start < num_bytes) {
+    uint64_t end = start + chunk_size;
+    if (end > num_bytes) end = num_bytes;
+    state = adler32_compute(data + start, end - start);
+  } else {
+    state = adler32_identity();
+  }
+
+  uint32_t threads_in_block = blockDim.x;
+  uint64_t block_start = (uint64_t)blockIdx.x * blockDim.x;
+  uint64_t total_threads = (num_bytes + chunk_size - 1) / chunk_size;
+  if (block_start + threads_in_block > total_threads) {
+    threads_in_block = (uint32_t)(total_threads - block_start);
+  }
+
+  Adler32State result = adler32_block_reduce(state, smem, threads_in_block);
+
+  if (threadIdx.x == 0) {
+    out_a[blockIdx.x] = result.a;
+    out_b[blockIdx.x] = result.b;
+    out_len[blockIdx.x] = result.len;
+  }
+}
+
+struct Adler32WholeKernel {
+  static void run(tvm::ffi::TensorView data, tvm::ffi::TensorView states, int64_t num_bytes, int64_t chunk_size) {
+    using namespace host;
+
+    const uint8_t* data_ptr = static_cast<const uint8_t*>(data.data_ptr());
+
+    int64_t total_threads = (num_bytes + chunk_size - 1) / chunk_size;
+    int64_t num_blocks = (total_threads + kBlockSize - 1) / kBlockSize;
+
+    uint32_t* out_a = static_cast<uint32_t*>(states.data_ptr());
+    uint32_t* out_b = out_a + num_blocks;
+    uint64_t* out_len = reinterpret_cast<uint64_t*>(out_b + num_blocks);
+
+    DLDevice device;
+    device.device_type = kDLCUDA;
+    device.device_id = data.device().device_id;
+
+    size_t smem_bytes = (kBlockSize / 32) * sizeof(Adler32State);
+    LaunchKernel(num_blocks, kBlockSize, device, smem_bytes)(
+        adler32_whole_kernel,
+        data_ptr,
+        out_a,
+        out_b,
+        out_len,
+        static_cast<uint64_t>(num_bytes),
+        static_cast<uint64_t>(chunk_size));
+  }
+};
+
+// --- Kernel 2: Strided items ---
+
+__global__ void adler32_strided_kernel(
+    const int64_t* __restrict__ ptrs,
+    const int32_t* __restrict__ lens,
+    uint32_t* __restrict__ out_a,
+    uint32_t* __restrict__ out_b,
+    uint64_t* __restrict__ out_len,
+    int64_t num_items) {
+  extern __shared__ Adler32State smem[];
+
+  int64_t global_tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+
+  Adler32State state;
+  if (global_tid < num_items) {
+    const uint8_t* ptr = reinterpret_cast<const uint8_t*>(ptrs[global_tid]);
+    uint32_t nbytes = static_cast<uint32_t>(lens[global_tid]);
+    state = adler32_compute(ptr, nbytes);
+  } else {
+    state = adler32_identity();
+  }
+
+  uint32_t threads_in_block = blockDim.x;
+  int64_t block_start = (int64_t)blockIdx.x * blockDim.x;
+  if (block_start + threads_in_block > num_items) {
+    threads_in_block = (uint32_t)(num_items - block_start);
+  }
+
+  Adler32State result = adler32_block_reduce(state, smem, threads_in_block);
+
+  if (threadIdx.x == 0) {
+    out_a[blockIdx.x] = result.a;
+    out_b[blockIdx.x] = result.b;
+    out_len[blockIdx.x] = result.len;
+  }
+}
+
+struct Adler32StridedKernel {
+  static void
+  run(tvm::ffi::TensorView ptrs, tvm::ffi::TensorView lens, tvm::ffi::TensorView states, int64_t num_items) {
+    using namespace host;
+
+    const int64_t* ptrs_ptr = static_cast<const int64_t*>(ptrs.data_ptr());
+    const int32_t* lens_ptr = static_cast<const int32_t*>(lens.data_ptr());
+
+    int64_t num_blocks = (num_items + kBlockSize - 1) / kBlockSize;
+
+    uint32_t* out_a = static_cast<uint32_t*>(states.data_ptr());
+    uint32_t* out_b = out_a + num_blocks;
+    uint64_t* out_len = reinterpret_cast<uint64_t*>(out_b + num_blocks);
+
+    DLDevice device;
+    device.device_type = kDLCUDA;
+    device.device_id = ptrs.device().device_id;
+
+    size_t smem_bytes = (kBlockSize / 32) * sizeof(Adler32State);
+    LaunchKernel(num_blocks, kBlockSize, device, smem_bytes)(
+        adler32_strided_kernel, ptrs_ptr, lens_ptr, out_a, out_b, out_len, num_items);
+  }
+};
+
+// --- Kernel 3: Final reduction ---
+
+__global__ void adler32_reduce_kernel(
+    const uint32_t* __restrict__ in_a,
+    const uint32_t* __restrict__ in_b,
+    const uint64_t* __restrict__ in_len,
+    int64_t* __restrict__ output,
+    int64_t num_blocks_to_reduce) {
+  extern __shared__ Adler32State smem[];
+
+  // Each thread loads one block result; if more blocks than threads,
+  // sequentially combine multiple entries first.
+  Adler32State state = adler32_identity();
+  int64_t items_per_thread = (num_blocks_to_reduce + blockDim.x - 1) / blockDim.x;
+  int64_t start = (int64_t)threadIdx.x * items_per_thread;
+  int64_t end = start + items_per_thread;
+  if (end > num_blocks_to_reduce) end = num_blocks_to_reduce;
+
+  for (int64_t i = start; i < end; ++i) {
+    Adler32State s = {in_a[i], in_b[i], in_len[i]};
+    state = adler32_combine(state, s);
+  }
+
+  uint32_t num_valid = blockDim.x;
+  int64_t total_active = (num_blocks_to_reduce + items_per_thread - 1) / items_per_thread;
+  if (total_active < num_valid) num_valid = (uint32_t)total_active;
+
+  Adler32State result = adler32_block_reduce(state, smem, num_valid);
+
+  if (threadIdx.x == 0) {
+    int64_t checksum = ((int64_t)result.b << 16) | (int64_t)result.a;
+    output[0] = checksum;
+  }
+}
+
+struct Adler32ReduceKernel {
+  static void run(tvm::ffi::TensorView states, tvm::ffi::TensorView output, int64_t num_blocks_to_reduce) {
+    using namespace host;
+
+    uint32_t* in_a = static_cast<uint32_t*>(states.data_ptr());
+    uint32_t* in_b = in_a + num_blocks_to_reduce;
+    uint64_t* in_len = reinterpret_cast<uint64_t*>(in_b + num_blocks_to_reduce);
+
+    int64_t* out_ptr = static_cast<int64_t*>(output.data_ptr());
+
+    uint32_t threads = kBlockSize;
+    if (num_blocks_to_reduce < (int64_t)threads) threads = (uint32_t)num_blocks_to_reduce;
+    // Round up to next power of 2 for reduction
+    uint32_t t = 1;
+    while (t < threads)
+      t *= 2;
+    threads = t;
+    if (threads > kBlockSize) threads = kBlockSize;
+
+    DLDevice device;
+    device.device_type = kDLCUDA;
+    device.device_id = states.device().device_id;
+
+    size_t smem_bytes = ((threads + 31) / 32) * sizeof(Adler32State);
+    LaunchKernel(1, threads, device, smem_bytes)(
+        adler32_reduce_kernel, in_a, in_b, in_len, out_ptr, num_blocks_to_reduce);
+  }
+};
+
+}  // namespace
+
+}  // namespace sglang

python/sglang/kernels/ops/memory/adler32.py

@@ -0,0 +1,136 @@
+"""Adler-32 checksum for GPU tensors.
+
+Three entry points:
+  - adler32_checksum(tensor) -> int : checksum all bytes of a contiguous tensor
+  - adler32_regions_checksum(data_ptrs, lengths, device) -> int : checksum
+    multiple raw GPU regions in order
+  - adler32_strided_checksum(data_ptrs, strides, indices) -> int : checksum
+    selected items across multiple tensors in a single kernel launch
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_adler32_module() -> Module:
+    return load_jit(
+        "adler32_checksum",
+        cuda_files=["elementwise/adler32_checksum.cuh"],
+        cuda_wrappers=[
+            ("adler32_whole", "Adler32WholeKernel::run"),
+            ("adler32_strided", "Adler32StridedKernel::run"),
+            ("adler32_reduce", "Adler32ReduceKernel::run"),
+        ],
+    )
+
+
+_BLOCK_SIZE = 256
+_MAX_BLOCKS = 1024
+_REGION_CHUNK_SIZE = 512 * 1024
+
+
+def _alloc_states(num_blocks: int, device: torch.device) -> torch.Tensor:
+    state_bytes = num_blocks * (4 + 4 + 8)
+    return torch.empty((state_bytes + 7) // 8, dtype=torch.int64, device=device)
+
+
+def _reduce_states(module: Module, states: torch.Tensor, num_blocks: int) -> int:
+    result = torch.empty(1, dtype=torch.int64, device=states.device)
+    module.adler32_reduce(states, result, num_blocks)
+    return result.item()
+
+
+def adler32_checksum(tensor: torch.Tensor) -> int:
+    """Compute Adler-32 checksum of entire tensor (all bytes)."""
+    assert tensor.is_contiguous(), "tensor must be contiguous"
+
+    module = _jit_adler32_module()
+    num_bytes = tensor.numel() * tensor.element_size()
+    if num_bytes == 0:
+        return 1
+
+    total_threads = min(num_bytes, _BLOCK_SIZE * _MAX_BLOCKS)
+    chunk_size = max(1, (num_bytes + total_threads - 1) // total_threads)
+    num_threads_needed = (num_bytes + chunk_size - 1) // chunk_size
+    num_blocks = (num_threads_needed + _BLOCK_SIZE - 1) // _BLOCK_SIZE
+
+    states = _alloc_states(num_blocks, tensor.device)
+    data_view = tensor.view(torch.uint8)
+    module.adler32_whole(data_view, states, num_bytes, chunk_size)
+    return _reduce_states(module, states, num_blocks)
+
+
+def adler32_regions_checksum(
+    data_ptrs: list[int], lengths: list[int], device: torch.device
+) -> int:
+    assert len(data_ptrs) == len(lengths)
+    module = _jit_adler32_module()
+    chunk_ptrs = []
+    chunk_lengths = []
+    for data_ptr, length in zip(data_ptrs, lengths):
+        assert data_ptr > 0 and length >= 0
+        offset = 0
+        while offset < length:
+            chunk_length = min(_REGION_CHUNK_SIZE, length - offset)
+            chunk_ptrs.append(data_ptr + offset)
+            chunk_lengths.append(chunk_length)
+            offset += chunk_length
+    if not chunk_ptrs:
+        return 1
+    ptrs = torch.tensor(chunk_ptrs, dtype=torch.int64, device=device)
+    lens = torch.tensor(chunk_lengths, dtype=torch.int32, device=device)
+    num_blocks = (len(chunk_ptrs) + _BLOCK_SIZE - 1) // _BLOCK_SIZE
+    states = _alloc_states(num_blocks, torch.device(device))
+    module.adler32_strided(ptrs, lens, states, len(chunk_ptrs))
+    return _reduce_states(module, states, num_blocks)
+
+
+def adler32_strided_checksum(
+    data_ptrs: list[int],
+    strides: list[int],
+    indices: list[torch.Tensor],
+) -> int:
+    """Compute Adler-32 checksum of selected items across multiple tensors.
+
+    Args:
+        data_ptrs: list of raw device pointers (int), one per tensor
+        strides: list of ints, bytes per item for each tensor
+        indices: list of 1D index tensors (one per tensor, on GPU)
+    """
+    module = _jit_adler32_module()
+    device = indices[0].device
+
+    all_ptrs = []
+    all_lens = []
+    total_items = 0
+    for base, idx, stride in zip(data_ptrs, indices, strides):
+        assert idx.is_contiguous(), "index tensor must be contiguous"
+        if idx.numel() == 0:
+            continue
+        offsets = idx.to(torch.int64) * stride + base
+        all_ptrs.append(offsets)
+        all_lens.append(
+            torch.full((idx.numel(),), stride, dtype=torch.int32, device=device)
+        )
+        total_items += idx.numel()
+
+    if total_items == 0:
+        return 1
+
+    ptrs = torch.cat(all_ptrs)
+    lens = torch.cat(all_lens)
+
+    num_blocks = (total_items + _BLOCK_SIZE - 1) // _BLOCK_SIZE
+    states = _alloc_states(num_blocks, device)
+
+    module.adler32_strided(ptrs, lens, states, total_items)
+    return _reduce_states(module, states, num_blocks)

python/sglang/srt/arg_groups/fields/disagg.py

@@ -70,6 +70,12 @@ class Disagg(msgspec.Struct):
             choices=DISAGG_TRANSFER_BACKEND_CHOICES,
         ),
     ] = "mooncake"
+    disaggregation_enable_kv_checksum: A[
+        bool,
+        "Compute an Adler-32 checksum over each request's KV pages on prefill "
+        "and verify it on decode. Enable on both prefill and decode engines. "
+        "A mismatch aborts the request, or raises in CI. Disabled by default.",
+    ] = False
     disaggregation_bootstrap_port: A[
         int, "Bootstrap server port on the prefill server. Default is 8998."
     ] = 8998

python/sglang/srt/disaggregation/checksum.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+from typing import List, Optional, Sequence
+
+import torch
+
+from sglang.kernels.ops.memory.adler32 import adler32_strided_checksum
+from sglang.srt.constants import HEALTH_CHECK_RID_PREFIX
+from sglang.srt.managers.schedule_batch import Req
+from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, HybridLinearKVPool
+from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
+
+NestedInts = Sequence[int] | Sequence[Sequence[int]]
+
+
+def _flatten_ints(values: NestedInts) -> list[int]:
+    flattened: list[int] = []
+    for value in values:
+        if isinstance(value, (list, tuple)):
+            flattened.extend(int(item) for item in value)
+        else:
+            flattened.append(int(value))
+    return flattened
+
+
+def is_health_check_req(req: Req) -> bool:
+    rid = req.rid
+    return isinstance(rid, str) and rid.startswith(HEALTH_CHECK_RID_PREFIX)
+
+
+def _to_page_indices_gpu(idx: torch.Tensor, page_size: int) -> torch.Tensor:
+    idx = idx.to(torch.int64).contiguous().reshape(-1)
+    if page_size == 1:
+        return idx
+    return (idx[::page_size] // page_size).contiguous()
+
+
+def page_indices_for_request(scheduler, req: Req, end_idx: int) -> torch.Tensor:
+    page_size = scheduler.token_to_kv_pool_allocator.page_size
+    kv_indices = scheduler.req_to_token_pool.req_to_token[
+        req.kv.req_pool_idx, 0:end_idx
+    ]
+    return _to_page_indices_gpu(kv_indices, page_size)
+
+
+def state_indices_for_request(
+    scheduler, req: Req, seq_len: int
+) -> Optional[torch.Tensor]:
+    pool = scheduler.token_to_kv_pool_allocator.get_kvcache()
+    if isinstance(pool, HybridLinearKVPool):
+        return (
+            scheduler.req_to_token_pool.req_index_to_mamba_index_mapping[
+                req.kv.req_pool_idx
+            ]
+            .to(torch.int64)
+            .contiguous()
+            .reshape(-1)
+        )
+    if isinstance(pool, SWAKVPool):
+        page_size = scheduler.token_to_kv_pool_allocator.page_size
+        window_size = scheduler.sliding_window_size
+        window_start = max(0, seq_len - window_size)
+        window_start = (window_start // page_size) * page_size
+        window_full = scheduler.req_to_token_pool.req_to_token[
+            req.kv.req_pool_idx, window_start:seq_len
+        ]
+        window_swa = (
+            scheduler.token_to_kv_pool_allocator.translate_loc_from_full_to_swa(
+                window_full
+            )
+        )
+        return _to_page_indices_gpu(window_swa, page_size)
+    if isinstance(pool, DSATokenToKVPool):
+        device_page_size = pool.page_size
+        kv_full = scheduler.req_to_token_pool.req_to_token[
+            req.kv.req_pool_idx, :seq_len
+        ]
+        return _to_page_indices_gpu(kv_full, device_page_size)
+    return None
+
+
+class KvChecksumComputer:
+    def __init__(
+        self,
+        device: torch.device,
+        kv_data_ptrs: Sequence[int],
+        kv_item_lens: Sequence[int],
+        state_data_ptrs: NestedInts = (),
+        state_item_lens: NestedInts = (),
+    ):
+        assert len(kv_data_ptrs) == len(kv_item_lens)
+        assert len(kv_data_ptrs) > 0
+        self._device = torch.device(device)
+        self._kv_data_ptrs = [int(ptr) for ptr in kv_data_ptrs]
+        self._kv_item_lens = [int(item_len) for item_len in kv_item_lens]
+        self._state_data_ptrs = _flatten_ints(state_data_ptrs)
+        self._state_item_lens = _flatten_ints(state_item_lens)
+        assert len(self._state_data_ptrs) == len(self._state_item_lens)
+
+    def compute(
+        self,
+        kv_page_indices_gpu: torch.Tensor,
+        state_indices_gpu: Optional[torch.Tensor] = None,
+    ) -> int:
+        assert kv_page_indices_gpu.is_cuda and kv_page_indices_gpu.is_contiguous()
+        all_ptrs = list(self._kv_data_ptrs)
+        all_lens = list(self._kv_item_lens)
+        all_indices: List[torch.Tensor] = [kv_page_indices_gpu] * len(
+            self._kv_data_ptrs
+        )
+        if self._state_data_ptrs:
+            assert state_indices_gpu is not None
+            assert state_indices_gpu.is_cuda and state_indices_gpu.is_contiguous()
+            all_ptrs += self._state_data_ptrs
+            all_lens += self._state_item_lens
+            all_indices += [state_indices_gpu] * len(self._state_data_ptrs)
+        return adler32_strided_checksum(all_ptrs, all_lens, all_indices)

python/sglang/srt/disaggregation/decode.py

@@ -37,6 +37,12 @@ from sglang.srt.configs.mamba_utils import Mamba2CacheParams
 from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.disaggregation.base.conn import StateType
+from sglang.srt.disaggregation.checksum import (
+    KvChecksumComputer,
+    is_health_check_req,
+    page_indices_for_request,
+    state_indices_for_request,
+)
 from sglang.srt.disaggregation.common.conn import CommonKVManager, CommonKVReceiver
 from sglang.srt.disaggregation.decode_hicache_mixin import (
     DecodeHiCachePreallocMixin,
@@ -113,6 +119,7 @@ from sglang.srt.runtime_context import (
 from sglang.srt.utils import ceil_align, get_num_new_pages, is_npu
 from sglang.srt.utils.network import NetworkAddress
 from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+from sglang.utils import is_in_ci
 
 logger = logging.getLogger(__name__)
 
@@ -430,6 +437,18 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
                 self.scheduler.tp_worker.model_runner.swa_max_total_num_tokens,
             )
 
+        if get_disagg().disaggregation_enable_kv_checksum:
+            kv_args = self.kv_manager.kv_args
+            self.scheduler.kv_checksum_computer = KvChecksumComputer(
+                device=torch.device(f"cuda:{self.scheduler.ps.gpu_id}"),
+                kv_data_ptrs=kv_args.kv_data_ptrs,
+                kv_item_lens=kv_args.kv_item_lens,
+                state_data_ptrs=kv_args.state_data_ptrs,
+                state_item_lens=kv_args.state_item_lens,
+            )
+        else:
+            self.scheduler.kv_checksum_computer = None
+
     def _uses_swa_tail_prealloc(self) -> bool:
         return (
             isinstance(self.token_to_kv_pool, (SWAKVPool, DeepSeekV4TokenToKVPool))
@@ -2137,7 +2156,13 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
             prealloc_queue.note_destinations_queued(len(decode_reqs))
 
     def _commit_transfer_to_req(self, decode_req: DecodeRequest):
+        # Preserve the checksum before the metadata slot is freed so it can be
+        # re-verified when the request enters a batch, including after retraction.
         idx = decode_req.metadata_buffer_index
+        if self.scheduler.kv_checksum_computer is not None:
+            decode_req.req.expected_kv_checksum = self.metadata_buffers.get_kv_checksum(
+                idx
+            )
         (
             output_id,
             cached_tokens,
@@ -2685,6 +2710,57 @@ class SchedulerDisaggregationDecodeMixin:
         return NextBatchPlan(batch_to_run=ret, running_batch=running_batch)
 
     def get_new_prebuilt_batch(
+        self, running_batch: ScheduleBatch
+    ) -> Optional[ScheduleBatch]:
+        computer: Optional[KvChecksumComputer] = self.kv_checksum_computer
+        if computer is None:
+            return self._get_new_prebuilt_batch(running_batch)
+
+        verified: List[Req] = []
+        for req in self.waiting_queue:
+            if is_health_check_req(req):
+                verified.append(req)
+                continue
+            expected = req.expected_kv_checksum
+            if expected == 0:
+                verified.append(req)
+                continue
+            seq_len = len(req.origin_input_ids)
+            page_indices_gpu = page_indices_for_request(self, req, seq_len)
+            state_indices = state_indices_for_request(self, req, seq_len)
+            actual = computer.compute(page_indices_gpu, state_indices)
+            if actual == expected:
+                verified.append(req)
+                continue
+            msg = (
+                f"KV checksum mismatch req={req.rid} "
+                f"bootstrap_room={req.bootstrap_room} "
+                f"expected={expected:#x} got={actual:#x}"
+            )
+            logger.error(msg)
+            self._handle_kv_checksum_mismatch(req, msg)
+        self.waiting_queue = verified
+
+        return self._get_new_prebuilt_batch(running_batch)
+
+    def _handle_kv_checksum_mismatch(self, req: Req, msg: str) -> None:
+        # A mismatch means the KV this worker received is not what prefill sent,
+        # so the cause is hardware or transport rather than the request. Serving
+        # keeps going and drops just this request; CI fails instead, because a
+        # single aborted request is easy to miss in a passing run.
+        if is_in_ci():
+            raise RuntimeError(msg)
+        prepare_abort(
+            req,
+            "KV checksum mismatch",
+            status_code=HTTPStatus.INTERNAL_SERVER_ERROR,
+        )
+        self.output_streamer.stream_output([req], req.return_logprob)
+        release_kv_cache(req, self.tree_cache, is_insert=False)
+        if self.metrics_reporter.enable_metrics:
+            self.metrics_collector.increment_transfer_failed_reqs()
+
+    def _get_new_prebuilt_batch(
         self: Scheduler, running_batch: ScheduleBatch
     ) -> Optional[ScheduleBatch]:
         """Create a schedulebatch for fake completed prefill"""

python/sglang/srt/disaggregation/prefill.py

@@ -31,6 +31,12 @@ import torch
 
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.disaggregation.base.conn import StateType
+from sglang.srt.disaggregation.checksum import (
+    KvChecksumComputer,
+    is_health_check_req,
+    page_indices_for_request,
+    state_indices_for_request,
+)
 from sglang.srt.disaggregation.common.conn import CommonKVManager
 from sglang.srt.disaggregation.common.staging_buffer import (
     compute_grid_segments,
@@ -205,6 +211,17 @@ class PrefillBootstrapQueue:
                     "supported by Mooncake."
                 )
         self.kv_manager = self._init_kv_manager()
+        if get_disagg().disaggregation_enable_kv_checksum:
+            kv_args = self.kv_manager.kv_args
+            self.scheduler.kv_checksum_computer = KvChecksumComputer(
+                device=torch.device(f"cuda:{self.scheduler.ps.gpu_id}"),
+                kv_data_ptrs=kv_args.kv_data_ptrs,
+                kv_item_lens=kv_args.kv_item_lens,
+                state_data_ptrs=kv_args.state_data_ptrs,
+                state_item_lens=kv_args.state_item_lens,
+            )
+        else:
+            self.scheduler.kv_checksum_computer = None
 
     def _init_kv_manager(self) -> CommonKVManager:
         kv_args_class = get_kv_class(self.transfer_backend, KVClassType.KVARGS)
@@ -1273,6 +1290,25 @@ class SchedulerDisaggregationPrefillMixin:
         self.send_kv_chunk(req, last_chunk=False, end_idx=cached_end)
 
     def send_kv_chunk(
+        self,
+        req: Req,
+        last_chunk: bool = False,
+        end_idx: Optional[int] = None,
+    ) -> None:
+        computer: Optional[KvChecksumComputer] = self.kv_checksum_computer
+        if last_chunk and computer is not None:
+            if is_health_check_req(req):
+                value = 0
+            else:
+                if end_idx is None:
+                    end_idx = min(req.extend_range.end, len(req.origin_input_ids))
+                page_indices_gpu = page_indices_for_request(self, req, end_idx)
+                state_indices = state_indices_for_request(self, req, end_idx)
+                value = computer.compute(page_indices_gpu, state_indices)
+            self.disagg_metadata_buffers.set_kv_checksum(req, value)
+        self._send_kv_chunk(req, last_chunk=last_chunk, end_idx=end_idx)
+
+    def _send_kv_chunk(
         self: Scheduler,
         req: Req,
         last_chunk: bool = False,

python/sglang/srt/disaggregation/utils.py

@@ -336,6 +336,8 @@ class MetadataBuffers:
         max_top_logprobs_num: int = 128,
         custom_mem_pool: torch.cuda.MemPool = None,
         output_dsa_topk_indices_dim: int = 0,
+        *,
+        kv_checksum_enabled: bool = False,
     ):
         self.custom_mem_pool = custom_mem_pool
         self.output_dsa_topk_indices_dim = output_dsa_topk_indices_dim
@@ -414,6 +416,26 @@ class MetadataBuffers:
                 (size, 8), dtype=bootstrap_room_dtype, device=device
             )
 
+        self.kv_checksum: torch.Tensor | None = None
+        if kv_checksum_enabled:
+            with (
+                torch.cuda.use_mem_pool(self.custom_mem_pool)
+                if self.custom_mem_pool
+                else nullcontext()
+            ):
+                # Width 8 (uint64) keeps the per-row size at the 64 B RDMA minimum.
+                self.kv_checksum = torch.zeros(
+                    (self.output_ids.shape[0], 8),
+                    dtype=self.bootstrap_room.dtype,
+                    device=self.bootstrap_room.device,
+                )
+
+    def set_kv_checksum(self, req: Req, value: int) -> None:
+        self.kv_checksum[req.metadata_buffer_index, 0] = value
+
+    def get_kv_checksum(self, idx: int) -> int:
+        return int(self.kv_checksum[idx, 0].item())
+
     def get_buf_infos(self):
         bufs = [
             self.output_ids,
@@ -432,6 +454,8 @@ class MetadataBuffers:
         if self.output_dsa_topk_indices is not None:
             bufs.append(self.output_dsa_topk_indices)
         bufs.append(self.bootstrap_room)
+        if self.kv_checksum is not None:
+            bufs.append(self.kv_checksum)
         bufs = [buf for buf in bufs if buf is not None]
         ptrs = [buf.data_ptr() for buf in bufs]
         data_lens = [buf.nbytes for buf in bufs]

python/sglang/srt/managers/schedule_batch.py

@@ -1313,6 +1313,7 @@ class Req(ReqDllmMixin):
         # first prefill batch; the cached-prefix early-send never goes past it.
         self.early_send_prefix_end: Optional[int] = None
         self.metadata_buffer_index: int = -1
+        self.expected_kv_checksum: int = 0
         # Used in overlap sequence to signal that an optimistic request should
         # abort chunking. Set in create_sender, consumed in process_batch_result.
         self.pending_bootstrap = False

python/sglang/srt/managers/scheduler.py

@@ -74,6 +74,7 @@ from sglang.srt.configs.model_config import (
 )
 from sglang.srt.constrained.grammar_manager import GrammarManager
 from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix
+from sglang.srt.disaggregation.checksum import KvChecksumComputer
 from sglang.srt.disaggregation.decode import (
     DecodePreallocQueue,
     DecodeTransferQueue,
@@ -438,6 +439,7 @@ class Scheduler(
     # Class-level default so on_idle's stall gate works even if a fork
     # overrides init_load_publisher (which would otherwise not set it).
     _last_stall_publish_ts: float = float("-inf")
+    kv_checksum_computer: Optional[KvChecksumComputer] = None
 
     def __init__(
         self,
@@ -1522,6 +1524,7 @@ class Scheduler(
                 max_sampling_mask_tokens=self.server_args.sampling_mask_max_tokens,
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
                 output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
+                kv_checksum_enabled=get_disagg().disaggregation_enable_kv_checksum,
             )
 
             # The decode requests polling kv cache
@@ -1569,6 +1572,7 @@ class Scheduler(
                 max_sampling_mask_tokens=self.server_args.sampling_mask_max_tokens,
                 custom_mem_pool=self.token_to_kv_pool_allocator.get_kvcache().maybe_get_custom_mem_pool(),
                 output_dsa_topk_indices_dim=output_dsa_topk_indices_dim,
+                kv_checksum_enabled=get_disagg().disaggregation_enable_kv_checksum,
             )
 
             self.disagg_prefill_bootstrap_queue = PrefillBootstrapQueue(