diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 5d62bacb0..6eff09e17 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -1123,6 +1123,116 @@ class MHATokenToKVPool(KVCache): ) +class NoOpMHATokenToKVPool(MHATokenToKVPool): + """KV cache pool that skips physical K/V buffer allocation. + + Used in embedding-mode prefill-only workloads with the FA + fa_skip_kv_cache path, where no layer reads or writes KV cache because + attention uses raw K/V via flash_attn_varlen_func. Other prefill-only paths + such as scoring/MIS may benefit from the same idea later, but some still + stage K/V through paged cache today. + + This class keeps the scheduler's view of pool capacity (self.size is + honored for admission) but allocates only (page_size, head_num, head_dim) + placeholder tensors per layer to satisfy any code paths that dereference + the buffers. + + Callers MUST ensure no real set_kv_buffer/get_*_buffer calls happen against + this pool; those paths raise loudly so misuse is visible. + """ + + def _create_buffers(self): + # Allocate minimal placeholder buffers. They exist purely so that code + # paths holding `k_buffer` / `v_buffer` references (pointer tables, + # layer-transfer counters, stride arithmetic) keep working without + # None-guards scattered across the codebase. Shape is + # [page_size, head_num, head_dim] per layer so that the unconditional + # `key_cache.view(-1, page_size, head_num, head_dim)` in the FA backend + # at the top of forward_extend succeeds regardless of --page-size. + # Total footprint is still on the order of KB vs GBs for a real pool. + with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE): + self.k_buffer = [ + torch.zeros( + (self.page_size, self.head_num, self.head_dim), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + self.v_buffer = [ + torch.zeros( + (self.page_size, self.head_num, self.v_head_dim), + dtype=self.store_dtype, + device=self.device, + ) + for _ in range(self.layer_num) + ] + + self.k_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.k_buffer], + dtype=torch.uint64, + device=self.device, + ) + self.v_data_ptrs = torch.tensor( + [x.data_ptr() for x in self.v_buffer], + dtype=torch.uint64, + device=self.device, + ) + self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0) + self.data_strides = torch.tensor( + [ + np.prod(x.shape[1:]) * x.dtype.itemsize + for x in self.k_buffer + self.v_buffer + ], + device=self.device, + ) + + def _finalize_allocation_log(self, num_tokens: int): + self.mem_usage = 0.0 + placeholder_bytes = ( + 2 + * self.layer_num + * self.page_size + * self.head_num + * max(self.head_dim, self.v_head_dim) + * self.store_dtype.itemsize + ) + logger.info( + f"KV Cache skipped (no-op pool). Logical #tokens: {num_tokens}, " + f"physical K/V size: ~{placeholder_bytes / 1024:.1f} KB placeholder" + ) + + def get_kv_size_bytes(self): + # Report zero so downstream memory accounting matches reality. + return (0, 0) + + def set_kv_buffer(self, *args, **kwargs): + raise RuntimeError( + "NoOpMHATokenToKVPool.set_kv_buffer was called. This pool is only " + "valid in prefill-only modes (e.g. --is-embedding, scoring) with " + "the FA backend's fa_skip_kv_cache path active; the attention " + "backend must never write to it. Check that the workload truly " + "performs no decode and that the FA backend's fa_skip_kv_cache " + "preconditions are met." + ) + + def get_key_buffer(self, layer_id: int): + # Return the placeholder. The FA backend reads this before taking the + # fa_skip_kv_cache branch (which does not use it); the placeholder shape + # is (page_size, head_num, head_dim) so downstream .view() calls succeed. + return self.k_buffer[layer_id - self.start_layer] + + def get_value_buffer(self, layer_id: int): + return self.v_buffer[layer_id - self.start_layer] + + def get_kv_buffer(self, layer_id: int): + return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # no-op; embedding mode has no KV cache to move + return + + class MHATokenToKVPoolFP4(MHATokenToKVPool): def _create_buffers(self): diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index dd612d020..1098b10a8 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py @@ -30,6 +30,7 @@ from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPoolFP4, MLATokenToKVPool, MLATokenToKVPoolFP4, + NoOpMHATokenToKVPool, NSATokenToKVPool, ReqToTokenPool, ) @@ -196,6 +197,44 @@ class ModelRunnerKVCacheMixin: return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio + def _validate_prefill_only_disable_kv_cache_pool_family( + self: ModelRunner, + is_nsa_model: bool, + is_dsv4_model: bool, + current_platform, + ): + if not self.server_args.prefill_only_disable_kv_cache or self.is_draft_worker: + return + + unsupported_pool_family = None + if is_dsv4_model: + unsupported_pool_family = "DeepSeekV4TokenToKVPool" + elif current_platform.is_out_of_tree() and not self.mambaish_config: + unsupported_pool_family = "out-of-tree platform KV pool" + elif ( + self.server_args.attention_backend == "ascend" and not self.mambaish_config + ): + unsupported_pool_family = "NPU/Ascend KV pool" + elif self.use_mla_backend and is_nsa_model: + unsupported_pool_family = "NSA/MLA KV pool" + elif self.use_mla_backend and not self.mambaish_config: + unsupported_pool_family = "MLA KV pool" + elif self.is_hybrid_swa: + unsupported_pool_family = "SWA KV pool" + elif self.mambaish_config: + unsupported_pool_family = "hybrid linear/Mamba KV pool" + elif is_float4_e2m1fn_x2(self.kv_cache_dtype): + unsupported_pool_family = "FP4 MHA KV pool" + + if unsupported_pool_family is not None: + raise RuntimeError( + "--prefill-only-disable-kv-cache is not supported for " + f"{unsupported_pool_family}. Supported configurations today: plain MHA " + "models on CUDA with the FA (fa3/fa4) prefill backend, --is-embedding, " + "--chunked-prefill-size=-1, --disable-radix-cache, no context-parallel " + "attention, no HiSparse, and --kv-cache-dtype != fp4_e2m1." + ) + def _init_pools(self: ModelRunner): """Initialize the memory pools.""" max_num_reqs = self.max_running_requests @@ -292,6 +331,10 @@ class ModelRunnerKVCacheMixin: # Out-of-tree platform plugin system — used by elif below from sglang.srt.platforms import current_platform + self._validate_prefill_only_disable_kv_cache_pool_family( + is_nsa_model, is_dsv4_model, current_platform + ) + if is_dsv4_model: swa_page_size = self.page_size assert swa_page_size == 256, "In paged swa mode, page_size must be 256." @@ -591,7 +634,12 @@ class ModelRunnerKVCacheMixin: ), ) else: - self.token_to_kv_pool = MHATokenToKVPool( + pool_cls = ( + NoOpMHATokenToKVPool + if self.server_args.prefill_only_disable_kv_cache + else MHATokenToKVPool + ) + self.token_to_kv_pool = pool_cls( self.max_total_num_tokens, page_size=self.page_size, dtype=self.kv_cache_dtype, @@ -719,6 +767,24 @@ class ModelRunnerKVCacheMixin: swa_allocator.full_to_swa_index_mapping ) + # Defensive check: the explicit validation above should reject known + # unsupported pool families before allocation. Keep this guard here so + # future pool-selection refactors fail at boot instead of on first use. + if ( + self.server_args.prefill_only_disable_kv_cache + and not self.is_draft_worker + and not isinstance(self.token_to_kv_pool, NoOpMHATokenToKVPool) + ): + raise RuntimeError( + "--prefill-only-disable-kv-cache expected NoOpMHATokenToKVPool but the " + f"runtime pool is {type(self.token_to_kv_pool).__name__}. This pool " + "family is not yet supported by --prefill-only-disable-kv-cache. " + "Supported configurations today: plain MHA models on CUDA with the FA " + "(fa3/fa4) prefill backend, --is-embedding, --chunked-prefill-size=-1, " + "--disable-radix-cache, no context-parallel attention, no HiSparse, " + "and --kv-cache-dtype != fp4_e2m1." + ) + def _apply_token_constraints(self: ModelRunner, token_capacity: int) -> int: """Apply external constraints to token capacity: user cap, PP sync. diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1660a2417..ec7791337 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -365,6 +365,7 @@ class ServerArgs: trust_remote_code: bool = False context_length: Optional[int] = None is_embedding: bool = False + prefill_only_disable_kv_cache: bool = False enable_multimodal: Optional[bool] = None revision: Optional[str] = None model_impl: str = "auto" @@ -846,6 +847,10 @@ class ServerArgs: # Validate PD disaggregation flags early (before dummy-model short-circuit). self._handle_pd_disaggregation() + # Validate --prefill-only-disable-kv-cache args early (before dummy-model + # short-circuit). The backend check is run later after backends settle. + self._validate_prefill_only_disable_kv_cache_args() + if self.model_path.lower() in ["none", "dummy"]: # Skip for dummy models return @@ -914,6 +919,13 @@ class ServerArgs: # the final attention backend and chunked_prefill_size are in effect. self._handle_multi_item_scoring() + # Backend-dependent half of --prefill-only-disable-kv-cache validation. + # Must stay after _handle_attention_backend_compatibility() (above) and + # _handle_multi_item_scoring() so the resolved prefill backend is final; + # the flag/precondition half runs earlier in + # _validate_prefill_only_disable_kv_cache_args(). + self._handle_prefill_only_disable_kv_cache() + # Handle Hicache settings. self._handle_hicache() @@ -3243,6 +3255,101 @@ class ServerArgs: "Pipeline parallelism is incompatible with overlap schedule." ) + def _validate_prefill_only_disable_kv_cache_args(self): + """Validate --prefill-only-disable-kv-cache flag/precondition constraints. + + Runs before the dummy-model short-circuit so misuse is rejected even + for dummy models. Backend resolution is checked separately by + _handle_prefill_only_disable_kv_cache after backends settle. + """ + if not self.prefill_only_disable_kv_cache: + return + + # This flag is intentionally scoped to embedding mode for now. Other + # prefill-only paths (for example scoring and MIS) can benefit from + # the same idea later, but some of them still stage K/V through the + # paged cache today. + if not self.is_embedding: + raise ValueError( + "--prefill-only-disable-kv-cache currently requires --is-embedding. " + "Other prefill-only workloads may be supported in a future change once " + "their attention paths stop reading or writing the paged KV cache." + ) + if self.kv_cache_dtype == "fp4_e2m1": + raise ValueError( + "--prefill-only-disable-kv-cache does not currently support " + "--kv-cache-dtype=fp4_e2m1 because the FP4 pool uses a separate " + "allocation path." + ) + + # Structural preconditions for the FA backend's fa_skip_kv_cache path, + # which is the only embedding path that doesn't read or write the pool: + # - chunked_prefill_size == -1 keeps a request in a single forward, + # so K/V never has to be reused across prefill chunks. + # - disable_radix_cache stops the prefix cache from indexing pool + # slots that no longer hold real data. + if self.chunked_prefill_size != -1: + raise ValueError( + "--prefill-only-disable-kv-cache requires --chunked-prefill-size=-1 so the FA " + "backend takes the fa_skip_kv_cache path; otherwise the pool would be touched " + "between prefill chunks." + ) + if not self.disable_radix_cache: + raise ValueError( + "--prefill-only-disable-kv-cache requires --disable-radix-cache because the " + "radix cache indexes KV pool slots that no longer hold real data." + ) + + # Context-parallel prefill stages K/V through cp_allgather_and_save_kv_cache, + # which writes to the pool via set_kv_buffer. NoOpMHATokenToKVPool intentionally + # raises on writes, so the engine would boot fine but fail on the first request. + if self.attn_cp_size > 1: + raise ValueError( + "--prefill-only-disable-kv-cache is incompatible with --attn-cp-size > 1: " + "the context-parallel attention path writes K/V to the pool via set_kv_buffer, " + "which the no-op pool intentionally rejects." + ) + if self.enable_prefill_context_parallel: + raise ValueError( + "--prefill-only-disable-kv-cache is incompatible with " + "--enable-prefill-context-parallel: the prefill-CP path stages K/V through " + "the paged cache, which the no-op pool does not support." + ) + + # HiSparse selects a different pool class (HiSparseNSATokenToKVPool / + # HiSparseTokenToKVPoolAllocator) that is not the no-op pool. + if self.enable_hisparse: + raise ValueError( + "--prefill-only-disable-kv-cache is incompatible with --enable-hisparse: " + "HiSparse uses a dedicated pool family that is not the no-op MHA pool." + ) + + def _handle_prefill_only_disable_kv_cache(self): + """Validate --prefill-only-disable-kv-cache backend constraint. + + Must run after _handle_attention_backend_compatibility() (which fills + the default attention_backend if unset) and _handle_multi_item_scoring() + (which may further mutate it). The assertion below guards against + accidental call-site reordering: if attention_backend is still None, + backends haven't settled yet and get_attention_backends() would return + a stale (None, None). + """ + if not self.prefill_only_disable_kv_cache: + return + + assert self.attention_backend is not None, ( + "_handle_prefill_only_disable_kv_cache must run after " + "_handle_attention_backend_compatibility() so the prefill backend is resolved." + ) + + prefill_backend, _ = self.get_attention_backends() + if prefill_backend not in ("fa3", "fa4"): + raise ValueError( + "--prefill-only-disable-kv-cache currently requires the FA prefill backend " + f"(fa3/fa4), but got prefill backend {prefill_backend!r}. Other prefill-only " + "workloads and backends may be supported in a future change." + ) + def _handle_hicache(self): """Normalize hicache-related knobs into a valid runtime configuration. @@ -4383,6 +4490,11 @@ class ServerArgs: action="store_true", help="Whether to use a CausalLM as an embedding model.", ) + parser.add_argument( + "--prefill-only-disable-kv-cache", + action="store_true", + help="Skip the physical KV cache allocation for embedding-mode prefill-only workloads. Currently only valid with --is-embedding, --chunked-prefill-size=-1, --disable-radix-cache, an FA prefill backend, and non-FP4 KV cache so the fa_skip_kv_cache path is active (no layer reads or writes the cache). Other prefill-only workloads such as scoring/MIS may benefit from this later once their attention paths stop using paged KV. Scheduler admission accounting is unchanged; per-layer K/V tensors are sized to (page_size, head_num, head_dim) placeholders so GPU memory is not wasted.", + ) parser.add_argument( "--enable-multimodal", default=ServerArgs.enable_multimodal, diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 16b9ac415..a24fd5425 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -515,5 +515,63 @@ class TestNgramExternalSamArgs(CustomTestCase): self.assertIn("external-corpus-max-tokens", str(context.exception)) +class TestPrefillOnlyDisableKvCache(unittest.TestCase): + """Validation for --prefill-only-disable-kv-cache. + + The flag wires NoOpMHATokenToKVPool, which is only safe when: + - the engine is in embedding mode (fa_skip_kv_cache active in FA backend), + - chunked_prefill_size == -1 (no inter-chunk K/V reuse), + - disable_radix_cache (radix cache otherwise indexes empty pool slots), + - no context-parallel attention (CP writes to the pool via set_kv_buffer), + - no HiSparse (uses a different pool family), + - kv_cache_dtype != fp4_e2m1 (FP4 pool is a separate allocation path). + All other configurations must be rejected at __post_init__ time so users + get a clear error before model load. + """ + + def _base_kwargs(self, **overrides): + kwargs = dict( + model_path="dummy", + is_embedding=True, + chunked_prefill_size=-1, + disable_radix_cache=True, + prefill_only_disable_kv_cache=True, + ) + kwargs.update(overrides) + return kwargs + + def test_valid_minimal_config_constructs(self): + sa = ServerArgs(**self._base_kwargs()) + self.assertTrue(sa.prefill_only_disable_kv_cache) + + def test_rejects_when_not_embedding(self): + with self.assertRaisesRegex(ValueError, "requires --is-embedding"): + ServerArgs(**self._base_kwargs(is_embedding=False)) + + def test_rejects_when_chunked_prefill_size_not_minus_one(self): + with self.assertRaisesRegex(ValueError, "--chunked-prefill-size=-1"): + ServerArgs(**self._base_kwargs(chunked_prefill_size=8192)) + + def test_rejects_when_radix_cache_enabled(self): + with self.assertRaisesRegex(ValueError, "--disable-radix-cache"): + ServerArgs(**self._base_kwargs(disable_radix_cache=False)) + + def test_rejects_attn_cp_size_greater_than_one(self): + with self.assertRaisesRegex(ValueError, "--attn-cp-size"): + ServerArgs(**self._base_kwargs(attn_cp_size=2, tp_size=2)) + + def test_rejects_prefill_context_parallel(self): + with self.assertRaisesRegex(ValueError, "--enable-prefill-context-parallel"): + ServerArgs(**self._base_kwargs(enable_prefill_context_parallel=True)) + + def test_rejects_hisparse(self): + with self.assertRaisesRegex(ValueError, "--enable-hisparse"): + ServerArgs(**self._base_kwargs(enable_hisparse=True)) + + def test_rejects_fp4_kv_cache(self): + with self.assertRaisesRegex(ValueError, "fp4_e2m1"): + ServerArgs(**self._base_kwargs(kv_cache_dtype="fp4_e2m1")) + + if __name__ == "__main__": unittest.main()