From 6365d6faeec335ad2e14ca72267d880ee3a32465 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 7 Jun 2026 21:27:46 -0700 Subject: [PATCH] [spec] Misc defensive guards for EAGLE draft KV indexing (#27486) --- .../sglang/jit_kernel/csrc/elementwise/kvcache.cuh | 12 ++++++++++-- python/sglang/jit_kernel/kvcache.py | 8 ++++++++ python/sglang/srt/mem_cache/memory_pool.py | 5 +++++ .../srt/speculative/multi_layer_eagle_worker.py | 5 +++++ .../srt/speculative/multi_layer_eagle_worker_v2.py | 5 +++++ python/sglang/srt/speculative/spec_utils.py | 3 +++ 6 files changed, 36 insertions(+), 2 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh b/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh index fa17cbf88..50133517d 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/kvcache.cuh @@ -8,6 +8,7 @@ #include #include +#include #include namespace { @@ -23,6 +24,7 @@ struct StoreKVCacheParams { int64_t stride_cache_bytes; int64_t stride_indices; uint32_t batch_size; + int64_t size_limit; }; constexpr uint32_t kNumWarps = 4; @@ -94,7 +96,8 @@ __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) const uint32_t split_id = warp_id % kSplit; const auto& [ k_input, v_input, k_cache, v_cache, indices, // ptr - stride_k, stride_v, stride_cache, stride_indices, batch_size // size + stride_k, stride_v, stride_cache, stride_indices, batch_size, // size + size_limit // bound ] = params; if (item_id >= batch_size) return; @@ -102,6 +105,9 @@ __global__ void store_kvcache(const __grid_constant__ StoreKVCacheParams params) PDLWaitPrimary(); const auto index = *index_ptr; + // A stale/OOB slot id would cause an illegal memory access in the store below; + // fail fast at the culprit instead. always-on (kvcache JIT compiles without NDEBUG). + assert(index >= 0 && index < size_limit); const auto k_src = pointer::offset(k_input, item_id * stride_k, split_id * kSplitSize); const auto v_src = pointer::offset(v_input, item_id * stride_v, split_id * kSplitSize); const auto k_dst = pointer::offset(k_cache, index * stride_cache, split_id * kSplitSize); @@ -138,7 +144,8 @@ struct StoreKVCacheKernel { const tvm::ffi::TensorView k_cache, const tvm::ffi::TensorView v_cache, const tvm::ffi::TensorView indices, - const int num_split) { + const int num_split, + const int64_t size_limit) { using namespace host; auto B = SymbolicSize{"batch_size"}; auto D = SymbolicSize{"element_size"}; @@ -188,6 +195,7 @@ struct StoreKVCacheKernel { .stride_cache_bytes = S.unwrap() * dtype_size, .stride_indices = I.unwrap(), .batch_size = static_cast(B.unwrap()), + .size_limit = size_limit, }; // select kernel and update num_split if needed const auto use_int32 = indice_dtype.is_type(); diff --git a/python/sglang/jit_kernel/kvcache.py b/python/sglang/jit_kernel/kvcache.py index 542d1866e..b611e7657 100644 --- a/python/sglang/jit_kernel/kvcache.py +++ b/python/sglang/jit_kernel/kvcache.py @@ -57,6 +57,7 @@ def store_cache( *, row_bytes: int = 0, num_split: int = 0, # can be tuned for performance + size_limit: int = 0, ) -> None: """Store key and value tensors into KV cache at specified indices. @@ -66,6 +67,10 @@ def store_cache( k_cache (torch.Tensor): Key cache tensor of shape (num_pages, H * D). v_cache (torch.Tensor): Value cache tensor of shape (num_pages, H * D). indices (torch.Tensor): Indices tensor of shape (batch_size,). + size_limit (int): Valid slot bound (cache row count = real slots + the + reserved padding slot); an index outside [0, size_limit) fails fast + (device assert) instead of an illegal memory access. Defaults to the + cache row count when 0. """ row_bytes = row_bytes or k.shape[-1] * k.element_size() module = _jit_kvcache_module(row_bytes) @@ -76,6 +81,8 @@ def store_cache( num_split = 2 else: num_split = 1 + if size_limit <= 0: + size_limit = k_cache.shape[0] module.store_cache( k, v, @@ -83,4 +90,5 @@ def store_cache( v_cache, indices, num_split, + size_limit, ) diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 24abecec0..8efe9aae9 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -103,6 +103,7 @@ def _set_kv_buffer_impl( row_dim: int, # head_num * head_dim store_dtype: torch.dtype, device_module: Any, + size_limit: int, alt_stream: Optional[torch.cuda.Stream] = None, same_kv_dim: bool = True, ) -> None: @@ -115,6 +116,7 @@ def _set_kv_buffer_impl( v_cache.view(-1, row_dim), indices, row_bytes=row_bytes, + size_limit=size_limit, ) if _is_cpu and _cpu_has_amx_support: @@ -1254,6 +1256,9 @@ class MHATokenToKVPool(KVCache): row_dim=self.row_dim, store_dtype=self.store_dtype, device_module=self.device_module, + # size + page_size = real slots + the reserved padding slot (padded / + # dummy tokens write there); valid index range is [0, size + page_size). + size_limit=self.size + self.page_size, alt_stream=self.alt_stream, same_kv_dim=self.same_kv_dim, ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker.py b/python/sglang/srt/speculative/multi_layer_eagle_worker.py index a7a325dad..3712c8460 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker.py @@ -89,6 +89,11 @@ class MultiLayerEagleWorker(TpModelWorker): self.topk = server_args.speculative_eagle_topk self.speculative_num_steps = server_args.speculative_num_steps self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + assert self.speculative_num_draft_tokens == self.speculative_num_steps + 1, ( + "multi-layer EAGLE requires speculative_num_draft_tokens == " + "speculative_num_steps + 1, " + f"got {self.speculative_num_draft_tokens} and {self.speculative_num_steps}" + ) self.gpu_id = gpu_id self.device = server_args.device self.target_worker = target_worker diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 59bf7c287..04b3841a2 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -106,6 +106,11 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): self.topk = server_args.speculative_eagle_topk self.speculative_num_steps = server_args.speculative_num_steps self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + assert self.speculative_num_draft_tokens == self.speculative_num_steps + 1, ( + "multi-layer EAGLE requires speculative_num_draft_tokens == " + "speculative_num_steps + 1, " + f"got {self.speculative_num_draft_tokens} and {self.speculative_num_steps}" + ) self.speculative_algorithm = SpeculativeAlgorithm.from_string( server_args.speculative_algorithm ) diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 18ff077d1..80c742d66 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -86,6 +86,9 @@ def draft_kv_indices_buffer_width( num_seqs * topk branches each attend up to max_context_len KV slots; the topk factor is mandatory -- dropping it under-allocates and overflows the row (#27338, #27460). """ + assert ( + num_seqs * topk * max_context_len < 2**31 + ), "kv_indices flat offset would overflow int32; reduce batch/topk/context" return num_seqs * topk * max_context_len