diff --git a/python/sglang/kernels/jit/csrc/elementwise/adler32_checksum.cuh b/python/sglang/kernels/jit/csrc/elementwise/adler32_checksum.cuh new file mode 100644 index 000000000..aa6cb8c85 --- /dev/null +++ b/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 + +#include + +#include +#include + +#include +#include + +// 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(ptr + offset) & 15) != 0) { + uint64_t value = ptr[offset]; + sum += value; + weighted += (nbytes - offset) * value; + ++offset; + } + const uint4* vectors = reinterpret_cast(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((1 + sum) % MOD_ADLER), + static_cast((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(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(states.data_ptr()); + uint32_t* out_b = out_a + num_blocks; + uint64_t* out_len = reinterpret_cast(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(num_bytes), + static_cast(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(ptrs[global_tid]); + uint32_t nbytes = static_cast(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(ptrs.data_ptr()); + const int32_t* lens_ptr = static_cast(lens.data_ptr()); + + int64_t num_blocks = (num_items + kBlockSize - 1) / kBlockSize; + + uint32_t* out_a = static_cast(states.data_ptr()); + uint32_t* out_b = out_a + num_blocks; + uint64_t* out_len = reinterpret_cast(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(states.data_ptr()); + uint32_t* in_b = in_a + num_blocks_to_reduce; + uint64_t* in_len = reinterpret_cast(in_b + num_blocks_to_reduce); + + int64_t* out_ptr = static_cast(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 diff --git a/python/sglang/kernels/ops/memory/adler32.py b/python/sglang/kernels/ops/memory/adler32.py new file mode 100644 index 000000000..36b48f72a --- /dev/null +++ b/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) diff --git a/python/sglang/srt/arg_groups/fields/disagg.py b/python/sglang/srt/arg_groups/fields/disagg.py index bbd6c4da1..8414e4894 100644 --- a/python/sglang/srt/arg_groups/fields/disagg.py +++ b/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 diff --git a/python/sglang/srt/disaggregation/checksum.py b/python/sglang/srt/disaggregation/checksum.py new file mode 100644 index 000000000..a32036fab --- /dev/null +++ b/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) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a3ef89846..f4af0635f 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/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""" diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 837b0c014..1a7d6788a 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/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, diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py index 3498b99e6..0762faa44 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/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] diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 7811e4477..6c4e7b9e8 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 442eba64d..cecc87100 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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( diff --git a/test/registered/kernel/disaggregation/test_kv_checksum.py b/test/registered/kernel/disaggregation/test_kv_checksum.py new file mode 100644 index 000000000..d7809c778 --- /dev/null +++ b/test/registered/kernel/disaggregation/test_kv_checksum.py @@ -0,0 +1,345 @@ +"""Tests for KV checksum integration in PD disaggregation.""" + +import unittest +import zlib +from types import SimpleNamespace +from unittest.mock import Mock, patch + +import torch + +from sglang.srt.disaggregation.checksum import ( + KvChecksumComputer, + is_health_check_req, +) +from sglang.srt.disaggregation.decode import SchedulerDisaggregationDecodeMixin +from sglang.srt.disaggregation.prefill import SchedulerDisaggregationPrefillMixin +from sglang.srt.disaggregation.utils import MetadataBuffers +from sglang.srt.environ import envs +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +def _ref_strided_adler32(tensors, indices, strides) -> int: + parts = [] + for tensor, idx, stride in zip(tensors, indices, strides): + raw = tensor.cpu().contiguous().flatten().view(torch.uint8) + for i in idx.cpu().tolist(): + parts.append(raw[i * stride : (i + 1) * stride].numpy().tobytes()) + return zlib.adler32(b"".join(parts)) + + +def _make_buf(size=4, *, kv_checksum_enabled=True, output_dsa_topk_indices_dim=0): + return MetadataBuffers( + size=size, + hidden_size=16, + hidden_states_dtype=torch.float32, + max_sampling_mask_tokens=128, + output_dsa_topk_indices_dim=output_dsa_topk_indices_dim, + kv_checksum_enabled=kv_checksum_enabled, + ) + + +class TestMetadataBuffers(unittest.TestCase): + def test_extends_aux_with_kv_checksum(self): + for sampling_mask in (False, True): + for seed_dim in (0, 32): + with ( + self.subTest(sampling_mask=sampling_mask, seed_dim=seed_dim), + envs.SGLANG_ENABLE_DISAGG_SAMPLING_MASK.override(sampling_mask), + ): + buf = _make_buf(size=8, output_dsa_topk_indices_dim=seed_dim) + disabled = _make_buf( + size=8, + kv_checksum_enabled=False, + output_dsa_topk_indices_dim=seed_dim, + ) + ptrs, data_lens, item_lens = buf.get_buf_infos() + base_ptrs, base_data_lens, base_item_lens = disabled.get_buf_infos() + self.assertEqual(len(ptrs), len(base_ptrs) + 1) + self.assertEqual(data_lens[:-1], base_data_lens) + self.assertEqual(item_lens[:-1], base_item_lens) + self.assertEqual(ptrs[-1], buf.kv_checksum.data_ptr()) + self.assertEqual(data_lens[-1], buf.kv_checksum.nbytes) + self.assertEqual(item_lens[-1], buf.kv_checksum[0].nbytes) + self.assertEqual(item_lens[-1], item_lens[-2]) + self.assertEqual(len(buf.get_buf(2)), len(disabled.get_buf(2))) + + def test_set_get_kv_checksum_roundtrip(self): + buf = _make_buf() + buf.set_kv_checksum(SimpleNamespace(metadata_buffer_index=1), 0xDEADBEEF) + self.assertEqual(buf.get_kv_checksum(1), 0xDEADBEEF) + self.assertEqual(buf.get_kv_checksum(0), 0) + + +class TestKvChecksumComputerConfig(unittest.TestCase): + def test_flattens_nested_state_descriptor_components(self): + computer = KvChecksumComputer( + torch.device("cpu"), + kv_data_ptrs=[11, 22], + kv_item_lens=[33, 44], + state_data_ptrs=[[55, 66], [77]], + state_item_lens=[[88, 99], [111]], + ) + + self.assertEqual(computer._state_data_ptrs, [55, 66, 77]) + self.assertEqual(computer._state_item_lens, [88, 99, 111]) + + +class TestKvChecksumHealthCheck(unittest.TestCase): + def test_detects_health_check_request(self): + self.assertTrue(is_health_check_req(SimpleNamespace(rid="HEALTH_CHECK_1"))) + self.assertFalse(is_health_check_req(SimpleNamespace(rid="user_req"))) + self.assertFalse(is_health_check_req(SimpleNamespace(rid=None))) + + +def _make_kv(num_layers, num_pages, page_elems, dtype=torch.float16): + return [ + torch.randn(num_pages, page_elems, dtype=dtype, device="cuda:0") + for _ in range(2 * num_layers) + ] + + +def _make_computer(kv, item_len, state=None, state_item_lens=None): + return KvChecksumComputer( + torch.device("cuda:0"), + kv_data_ptrs=[t.data_ptr() for t in kv], + kv_item_lens=[item_len] * len(kv), + state_data_ptrs=[t.data_ptr() for t in (state or [])], + state_item_lens=state_item_lens or [], + ) + + +class TestKvChecksumComputer(unittest.TestCase): + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA not available") + self.device = torch.device("cuda:0") + + def test_kv_only_matches_reference(self): + kv = _make_kv(num_layers=4, num_pages=32, page_elems=128, dtype=torch.bfloat16) + idx = torch.tensor([3, 7, 8, 15, 31], dtype=torch.int64, device=self.device) + item_len = 128 * 2 + value = _make_computer(kv, item_len).compute(idx) + expected = _ref_strided_adler32(kv, [idx] * len(kv), [item_len] * len(kv)) + self.assertEqual(value, expected) + + def test_kv_corruption_detected(self): + kv = _make_kv(num_layers=2, num_pages=16, page_elems=64) + idx = torch.tensor([5], dtype=torch.int64, device=self.device) + computer = _make_computer(kv, 64 * 2) + v1 = computer.compute(idx) + kv[0][5, 0] += 1 + self.assertNotEqual(v1, computer.compute(idx)) + + def test_kv_plus_state_matches_reference(self): + kv = _make_kv(num_layers=2, num_pages=8, page_elems=32) + state = [ + torch.randn(4, 16, dtype=torch.float16, device=self.device) + for _ in range(2) + ] + kv_idx = torch.tensor([0, 1, 2], dtype=torch.int64, device=self.device) + state_idx = torch.tensor([1], dtype=torch.int64, device=self.device) + kv_len, state_lens = 32 * 2, [16 * 2, 16 * 2] + computer = _make_computer(kv, kv_len, state, state_lens) + value = computer.compute(kv_idx, state_idx) + expected = _ref_strided_adler32( + kv + state, + [kv_idx] * len(kv) + [state_idx] * len(state), + [kv_len] * len(kv) + state_lens, + ) + self.assertEqual(value, expected) + state[0][1, 0] += 1 + self.assertNotEqual(value, computer.compute(kv_idx, state_idx)) + + def test_nested_state_components_match_reference(self): + kv = _make_kv(num_layers=1, num_pages=8, page_elems=32) + state_components = [ + [torch.randn(4, 16, dtype=torch.float16, device=self.device)], + [ + torch.randn(4, 8, dtype=torch.float16, device=self.device), + torch.randn(4, 12, dtype=torch.float16, device=self.device), + ], + ] + kv_idx = torch.tensor([0, 3, 7], dtype=torch.int64, device=self.device) + state_idx = torch.tensor([1, 2], dtype=torch.int64, device=self.device) + state_tensors = [tensor for comp in state_components for tensor in comp] + kv_len = 32 * 2 + state_lens = [[16 * 2], [8 * 2, 12 * 2]] + computer = KvChecksumComputer( + self.device, + kv_data_ptrs=[t.data_ptr() for t in kv], + kv_item_lens=[kv_len] * len(kv), + state_data_ptrs=[ + [tensor.data_ptr() for tensor in comp] for comp in state_components + ], + state_item_lens=state_lens, + ) + + value = computer.compute(kv_idx, state_idx) + expected = _ref_strided_adler32( + kv + state_tensors, + [kv_idx] * len(kv) + [state_idx] * len(state_tensors), + [kv_len] * len(kv) + [item for comp in state_lens for item in comp], + ) + self.assertEqual(value, expected) + + +class _FakeScheduler(SchedulerDisaggregationDecodeMixin): + def __init__(self, computer, req_to_token): + self.kv_checksum_computer = computer + self.waiting_queue = [] + self.token_to_kv_pool_allocator = SimpleNamespace( + page_size=1, get_kvcache=lambda: SimpleNamespace() + ) + self.req_to_token_pool = SimpleNamespace(req_to_token=req_to_token) + self.tree_cache = None + self.output_streamer = SimpleNamespace(stream_output=self.stream_output) + self.metrics_reporter = SimpleNamespace(enable_metrics=True) + self.metrics_collector = Mock() + self.streamed_aborts = [] + + def stream_output(self, reqs, return_logprob): + self.streamed_aborts.extend(reqs) + + +class _FakePrefillScheduler(SchedulerDisaggregationPrefillMixin): + def __init__(self): + self.kv_checksum_computer = object() + self.disagg_metadata_buffers = SimpleNamespace(set_kv_checksum=Mock()) + + +def _make_req(expected_chksum, num_input_tokens, rid="r0"): + return SimpleNamespace( + rid=rid, + bootstrap_room=12345, + kv=SimpleNamespace(req_pool_idx=0), + origin_input_ids=list(range(num_input_tokens)), + fill_ids=list(range(num_input_tokens)), + expected_kv_checksum=expected_chksum, + return_logprob=False, + ) + + +class TestPrefillHealthCheckChecksum(unittest.TestCase): + def test_health_check_clears_metadata_checksum(self): + sched = _FakePrefillScheduler() + req = _make_req(0xDEADBEEF, 1, rid="HEALTH_CHECK_1") + with patch.object( + SchedulerDisaggregationPrefillMixin, + "_send_kv_chunk", + lambda *args, **kwargs: None, + ): + sched.send_kv_chunk(req, last_chunk=True) + sched.disagg_metadata_buffers.set_kv_checksum.assert_called_once_with(req, 0) + + +class TestGetNewPrebuiltBatchChecksum(unittest.TestCase): + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA not available") + self.device = torch.device("cuda:0") + self.num_pages = 8 + self.kv = _make_kv(num_layers=2, num_pages=self.num_pages, page_elems=32) + self.item_len = 32 * 2 + self.req_to_token = torch.arange( + self.num_pages, dtype=torch.int64, device=self.device + ).view(1, self.num_pages) + self.true_chksum = _ref_strided_adler32( + self.kv, + [torch.arange(self.num_pages, dtype=torch.int64, device=self.device)] + * len(self.kv), + [self.item_len] * len(self.kv), + ) + + def _make_sched(self, computer=None): + if computer is _SENTINEL: + computer = _make_computer(self.kv, self.item_len) + return _FakeScheduler(computer, self.req_to_token) + + def _run_once(self, sched, batch_ret=None): + running_batch = SimpleNamespace() + with patch.object( + SchedulerDisaggregationDecodeMixin, + "_get_new_prebuilt_batch", + lambda s, rb: batch_ret, + ): + return sched.get_new_prebuilt_batch(running_batch) + + def test_match_keeps_req(self): + sched = self._make_sched(_SENTINEL) + sched.waiting_queue = [_make_req(self.true_chksum, self.num_pages)] + self._run_once(sched) + self.assertEqual(len(sched.waiting_queue), 1) + self.assertEqual(sched.streamed_aborts, []) + + def test_mismatch_aborts(self): + sched = self._make_sched(_SENTINEL) + req = _make_req(0xDEADBEEF, self.num_pages) + sched.waiting_queue = [req] + with ( + envs.SGLANG_IS_IN_CI.override(False), + patch("sglang.srt.disaggregation.decode.prepare_abort") as mock_abort, + patch("sglang.srt.disaggregation.decode.release_kv_cache") as mock_release, + ): + self._run_once(sched) + self._run_once(sched) + self.assertEqual(sched.waiting_queue, []) + self.assertEqual(sched.streamed_aborts, [req]) + mock_abort.assert_called_once() + mock_release.assert_called_once() + sched.metrics_collector.increment_transfer_failed_reqs.assert_called_once_with() + + def test_mismatch_raises_in_ci(self): + sched = self._make_sched(_SENTINEL) + sched.waiting_queue = [_make_req(0xDEADBEEF, self.num_pages)] + with ( + envs.SGLANG_IS_IN_CI.override(True), + patch("sglang.srt.disaggregation.decode.prepare_abort") as mock_abort, + self.assertRaisesRegex(RuntimeError, "KV checksum mismatch"), + ): + self._run_once(sched) + mock_abort.assert_not_called() + sched.metrics_collector.increment_transfer_failed_reqs.assert_not_called() + + def test_health_check_skips_checksum(self): + sched = self._make_sched(_SENTINEL) + req = _make_req(0xDEADBEEF, self.num_pages, rid="HEALTH_CHECK_1") + sched.waiting_queue = [req] + with ( + patch("sglang.srt.disaggregation.decode.prepare_abort") as mock_abort, + patch("sglang.srt.disaggregation.decode.release_kv_cache") as mock_release, + ): + self._run_once(sched) + self.assertEqual(sched.waiting_queue, [req]) + self.assertEqual(sched.streamed_aborts, []) + mock_abort.assert_not_called() + mock_release.assert_not_called() + + def test_retract_re_verifies(self): + sched = self._make_sched(_SENTINEL) + req = _make_req(self.true_chksum, self.num_pages) + sched.waiting_queue = [req] + self._run_once(sched) + self._run_once(sched) + self.assertEqual(sched.waiting_queue, [req]) + self.assertEqual(sched.streamed_aborts, []) + + def test_disabled_delegates_to_batch_builder(self): + sched = self._make_sched(computer=None) + sched.waiting_queue = [_make_req(0xABCD, 4)] + sentinel = object() + self.assertIs(self._run_once(sched, batch_ret=sentinel), sentinel) + + def test_zero_expected_skips_checksum(self): + sched = self._make_sched(_SENTINEL) + sched.waiting_queue = [_make_req(0, self.num_pages)] + self._run_once(sched) + self.assertEqual(len(sched.waiting_queue), 1) + + +_SENTINEL = object() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/kernel/memory/test_adler32_checksum.py b/test/registered/kernel/memory/test_adler32_checksum.py new file mode 100644 index 000000000..7f39dd7fb --- /dev/null +++ b/test/registered/kernel/memory/test_adler32_checksum.py @@ -0,0 +1,95 @@ +"""Tests for Adler-32 GPU checksum against Python zlib.adler32.""" + +import unittest +import zlib + +import torch + +from sglang.kernels.ops.memory.adler32 import ( + adler32_checksum, + adler32_regions_checksum, + adler32_strided_checksum, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large") + + +def _ref_adler32(tensor: torch.Tensor) -> int: + return zlib.adler32(tensor.cpu().contiguous().view(torch.uint8).numpy().tobytes()) + + +def _ref_strided_adler32(tensors, indices, strides) -> int: + parts = [] + for tensor, idx, stride in zip(tensors, indices, strides): + raw = tensor.cpu().contiguous().flatten().view(torch.uint8) + for i in idx.cpu().tolist(): + parts.append(raw[i * stride : (i + 1) * stride].numpy().tobytes()) + return zlib.adler32(b"".join(parts)) + + +class TestAdler32(unittest.TestCase): + def setUp(self) -> None: + if not torch.cuda.is_available(): + self.skipTest("CUDA not available") + torch.manual_seed(42) + + def _check_whole(self, tensor): + self.assertEqual(adler32_checksum(tensor), _ref_adler32(tensor)) + + def test_whole_small(self): + self._check_whole(torch.tensor([1.0, 2, 3, 4], device="cuda")) + + def test_whole_single_element(self): + self._check_whole(torch.tensor([42.0], device="cuda")) + + def test_whole_dtypes(self): + for dtype, shape in [ + (torch.bfloat16, (1024, 128)), + (torch.float16, (512, 64)), + (torch.float32, (4096, 256)), + ]: + self._check_whole(torch.randn(*shape, dtype=dtype, device="cuda")) + + def _check_strided(self, tensors, indices, strides): + actual = adler32_strided_checksum( + [t.data_ptr() for t in tensors], strides, indices + ) + expected = _ref_strided_adler32(tensors, indices, strides) + self.assertEqual(actual, expected) + + def test_strided_single_tensor(self): + t = torch.randn(100, 64, device="cuda") + idx = torch.tensor([0, 5, 10, 50, 99], dtype=torch.int64, device="cuda") + self._check_strided([t], [idx], [64 * 4]) + + def test_strided_multi_tensor_different_strides(self): + t1 = torch.randn(50, 32, dtype=torch.float16, device="cuda") + t2 = torch.randn(80, 64, dtype=torch.float16, device="cuda") + idx1 = torch.tensor([0, 10, 49], dtype=torch.int64, device="cuda") + idx2 = torch.tensor([30, 79], dtype=torch.int64, device="cuda") + self._check_strided([t1, t2], [idx1, idx2], [32 * 2, 64 * 2]) + + def test_strided_many_items(self): + t = torch.randn(1000, 128, dtype=torch.bfloat16, device="cuda") + idx = torch.arange(1000, dtype=torch.int64, device="cuda") + self._check_strided([t], [idx], [128 * 2]) + + def test_regions(self): + first = torch.randint( + 0, 256, (5 * 1024 * 1024,), dtype=torch.uint8, device="cuda" + ) + second = torch.randint(0, 256, (12345,), dtype=torch.uint8, device="cuda") + actual = adler32_regions_checksum( + [first.data_ptr(), second.data_ptr()], + [first.numel(), second.numel()], + first.device, + ) + expected = zlib.adler32( + first.cpu().numpy().tobytes() + second.cpu().numpy().tobytes() + ) + self.assertEqual(actual, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index fa25f2548..6fc1a88c4 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -117,7 +117,9 @@ def _commit_disagg_handoff( replayed_boundary: bool = False, ) -> None: queue = DecodeTransferQueue.__new__(DecodeTransferQueue) - queue.scheduler = SimpleNamespace(batch_result_processor=processor) + queue.scheduler = SimpleNamespace( + batch_result_processor=processor, kv_checksum_computer=None + ) queue.spec_algorithm = SimpleNamespace(is_none=lambda: True) queue.metadata_buffers = SimpleNamespace( get_buf=lambda _: (