From 0911a9627950693c3790c06a7a6794b6cd967de8 Mon Sep 17 00:00:00 2001 From: Pengyu Chen Date: Sun, 24 May 2026 06:40:36 +0000 Subject: [PATCH] feat(yoco): add --kv-sharing-fast-prefill server arg + validators (PR-A/3) Introduce an opt-in server flag for the YOCO-style fast-prefill optimization that vLLM ships for Gemma-4 (Gemma4SelfDecoderLayers / Gemma4CrossDecoderLayers in vllm/model_executor/models/gemma4.py). When enabled, the last num_kv_shared_layers decoder layers will eventually run only on the per-request last-extend-token rows during EXTEND-mode forwards (model-side implementation lands in the follow-up PR). This first PR in a stack of three contains only the flag plumbing and the cross-cutting validators that don't depend on the live model. Changes: * ServerArgs.kv_sharing_fast_prefill: bool dataclass field (default False) * argparse --kv-sharing-fast-prefill (store_true) * ModelConfig.__init__ accepts kv_sharing_fast_prefill keyword * ModelConfig.from_server_args propagates the flag * Gemma-4 branch of _handle_model_specific_adjustments validates: - triton backend (prefill + decode) when the flag is on - num_kv_shared_layers > 0 on the served model's text config - logs 'KV-sharing fast prefill enabled for ' on success * ServerArgs._validate_kv_sharing_fast_prefill_combos (new helper) rejects --speculative-algorithm EAGLE / EAGLE3 with the flag on. Called from check_server_args(); factored out so it's unit-testable without a live HF config. * test/registered/unit/server_args/test_kv_sharing_fast_prefill.py: 8 unit tests covering the dataclass default, CLI parsing, argparse registration, ModelConfig plumbing, and all four validator branches (EAGLE rejected, EAGLE3 rejected, NEXTN accepted, flag-off skipped). Default-off, opt-in only; no behavior change for existing deployments. The diff in server_args.py contains some auto-format noise (assert (cond), 'msg' -> assert cond, ('msg')) from a ruff format rule that triggers on save. The real changes are the four blocks cited above. Stack base: pyc/sota-gemma4-31b-mm-disabled @ 3a3195b30 Next: PR-B (model-side YOCO branch in Gemma4TextModel.forward) and PR-C (benchmark + tuning). Plan: .humanize/yoco-gemma4/refined-plan.md Co-authored-by: Claude --- python/sglang/srt/configs/model_config.py | 6 + python/sglang/srt/server_args.py | 401 +++++++++++------- .../test_kv_sharing_fast_prefill.py | 114 +++++ 3 files changed, 375 insertions(+), 146 deletions(-) create mode 100644 test/registered/unit/server_args/test_kv_sharing_fast_prefill.py diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 3d87838bb324..3b510fae352e 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -175,6 +175,7 @@ def __init__( language_only: bool = False, disable_hybrid_swa_memory: bool = False, model_config_parser: str = "auto", + kv_sharing_fast_prefill: bool = False, ) -> None: # Parse args self.model_path = model_path @@ -187,6 +188,10 @@ def __init__( self.is_multi_layer_eagle = is_multi_layer_eagle self.disable_hybrid_swa_memory = disable_hybrid_swa_memory self.model_config_parser = model_config_parser + # YOCO fast-prefill opt-in flag; consumed by models with + # ``num_kv_shared_layers > 0`` (e.g. Gemma-4). See + # ``ServerArgs.kv_sharing_fast_prefill`` for full semantics. + self.kv_sharing_fast_prefill = kv_sharing_fast_prefill # Validate quantize_and_serve configuration self._validate_quantize_and_serve_config() @@ -408,6 +413,7 @@ def from_server_args( is_draft_model=is_draft_model, disable_hybrid_swa_memory=server_args.disable_hybrid_swa_memory, model_config_parser=server_args.model_config_parser, + kv_sharing_fast_prefill=server_args.kv_sharing_fast_prefill, **kwargs, ) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index d4192f947744..1975cf4232f3 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -681,6 +681,13 @@ class ServerArgs: # placeholder token. See MIS_DELIMITER_TOKEN_ID for details. enable_mis: bool = False + # KV-sharing fast-prefill (YOCO) — opt-in optimization for models with + # ``num_kv_shared_layers > 0`` (Gemma-4 family). When enabled, the back + # half of the decoder runs only on the per-request last-extend-token rows + # during EXTEND-mode forwards, skipping redundant compute on the KV-shared + # layers. See ``gemma4_causal.py`` for the model-side implementation. + kv_sharing_fast_prefill: bool = False + # Optimization/debug options disable_radix_cache: bool = False cuda_graph_max_bs: Optional[int] = None @@ -1487,9 +1494,9 @@ def _handle_gpu_memory_settings(self, gpu_mem): ) self.cuda_graph_bs = self._generate_cpu_graph_batch_sizes() - assert ( - self.torch_compile_max_bs > 0 - ), "cuda_graph_bs should contain positive batch sizes" + assert self.torch_compile_max_bs > 0, ( + "cuda_graph_bs should contain positive batch sizes" + ) self.cuda_graph_max_bs = self.torch_compile_max_bs if self.piecewise_cuda_graph_max_tokens is None: @@ -1815,12 +1822,12 @@ def _handle_model_specific_adjustments(self): else: self.enable_dp_attention = True self.moe_dense_tp_size = 1 - assert ( - self.dp_size == 1 - ), "For round-robin split mode, dp attention is not supported." - assert ( - self.tp_size <= 8 - ), "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + assert self.dp_size == 1, ( + "For round-robin split mode, dp attention is not supported." + ) + assert self.tp_size <= 8, ( + "Context parallel only supports single machine (tp_size <= 8). Cross-machine CP has precision issues." + ) self.attn_cp_size = self.tp_size // self.dp_size logger.warning( @@ -1861,9 +1868,9 @@ def _handle_model_specific_adjustments(self): self._set_default_dsa_backends(self.kv_cache_dtype, major) if self.enable_dsa_prefill_context_parallel: - assert ( - self.disaggregation_mode != "decode" - ), "CP is only supported for prefill when PD disaggregation, please remove --enable-dsa-prefill-context-parallel." + assert self.disaggregation_mode != "decode", ( + "CP is only supported for prefill when PD disaggregation, please remove --enable-dsa-prefill-context-parallel." + ) else: # DeepSeek V3/R1/V3.1 @@ -2098,9 +2105,9 @@ def _handle_model_specific_adjustments(self): ) if self.moe_runner_backend == "triton_kernel": - assert ( - self.ep_size == 1 - ), "Triton kernel MoE is only supported when ep_size == 1" + assert self.ep_size == 1, ( + "Triton kernel MoE is only supported when ep_size == 1" + ) elif model_arch in MIMO_V2_MODEL_ARCHS: if model_arch == "MiMoV2ForCausalLM" and not self.encoder_only: @@ -2182,7 +2189,9 @@ def _handle_model_specific_adjustments(self): "ascend", "trtllm_mha", "intel_xpu", - }, f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {self.attention_backend}" + }, ( + f"fa3, aiter, triton, ascend, trtllm_mha or intel_xpu is required for Llama4 model but got {self.attention_backend}" + ) if is_sm100_supported() and self.moe_runner_backend == "auto": if self.quantization in {"fp8", "modelopt_fp8"}: self.moe_runner_backend = "flashinfer_trtllm" @@ -2231,6 +2240,41 @@ def _handle_model_specific_adjustments(self): f"got prefill={prefill_backend}, decode={decode_backend}" ) + # YOCO fast-prefill compatibility validation. When the user opts + # into ``--kv-sharing-fast-prefill``, require the triton attention + # backend (the v0 model code uses the triton metadata builder + # directly), and the model must actually declare KV-shared layers. + if self.kv_sharing_fast_prefill: + if prefill_backend != "triton" or decode_backend != "triton": + raise ValueError( + "--kv-sharing-fast-prefill currently requires the " + "triton attention backend on both prefill and decode " + f"sides, got prefill={prefill_backend}, " + f"decode={decode_backend}. Pass " + "``--attention-backend triton`` (or omit the override " + "on H100/SM90, which already defaults to triton)." + ) + _hf_text_config_yoco = None + try: + _hf_text_config_yoco = self.get_model_config().hf_text_config + except Exception: + pass + _num_kv_shared = int( + getattr(_hf_text_config_yoco, "num_kv_shared_layers", 0) or 0 + ) + if _num_kv_shared <= 0: + raise ValueError( + "--kv-sharing-fast-prefill requires the served model " + "to declare ``num_kv_shared_layers > 0`` on its text " + f"config; got num_kv_shared_layers={_num_kv_shared} " + f"for {model_arch}. Remove the flag for this model." + ) + logger.info( + "KV-sharing fast prefill enabled for %s (num_kv_shared_layers=%d).", + model_arch, + _num_kv_shared, + ) + if is_sm100_supported() and self.moe_runner_backend == "auto": if self.get_model_config().quantization == "modelopt_fp4": self.quantization = "modelopt_fp4" @@ -2316,9 +2360,9 @@ def _handle_model_specific_adjustments(self): self.disable_hybrid_swa_memory = True # https://docs.sglang.ai/advanced_features/attention_backend.html accepted_backends = ["fa3", "triton", "trtllm_mha"] - assert ( - self.attention_backend in accepted_backends - ), f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {self.attention_backend}" + assert self.attention_backend in accepted_backends, ( + f"One of the attention backends in {accepted_backends} is required for {model_arch}, but got {self.attention_backend}" + ) elif model_arch in ["Olmo2ForCausalLM"]: # FIXME: https://github.com/sgl-project/sglang/pull/7367 is not compatible with Olmo3 model. logger.warning( @@ -2337,9 +2381,9 @@ def _handle_model_specific_adjustments(self): # Flashinfer appears to degrade performance when sliding window attention # is used for the Olmo2 architecture. Olmo2 does not use sliding window attention # but Olmo3 does. - assert ( - self.attention_backend != "flashinfer" - ), "FlashInfer backend can significantly degrade the performance of Olmo3 models." + assert self.attention_backend != "flashinfer", ( + "FlashInfer backend can significantly degrade the performance of Olmo3 models." + ) logger.info( f"Using {self.attention_backend} as attention backend for {model_arch}." @@ -2576,9 +2620,9 @@ def _handle_mamba_radix_cache( return if not support_mamba_cache_extra_buffer: - assert ( - not self.enable_mamba_extra_buffer() - ), f"mamba extra_buffer is not supported for {model_arch} model" + assert not self.enable_mamba_extra_buffer(), ( + f"mamba extra_buffer is not supported for {model_arch} model" + ) if self.enable_mamba_extra_buffer(): # extra_buffer if self.disable_radix_cache: @@ -2588,23 +2632,25 @@ def _handle_mamba_radix_cache( "Please use --mamba-scheduler-strategy no_buffer instead." ) - assert ( - is_cuda() or is_musa() or is_npu() - ), "Mamba extra_buffer is only supported on CUDA and MUSA and NPU devices with FLA backend" + assert is_cuda() or is_musa() or is_npu(), ( + "Mamba extra_buffer is only supported on CUDA and MUSA and NPU devices with FLA backend" + ) if self.speculative_num_draft_tokens is not None: - assert ( - self.mamba_track_interval >= self.speculative_num_draft_tokens - ), f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}" + assert self.mamba_track_interval >= self.speculative_num_draft_tokens, ( + f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}" + ) if self.page_size is not None: - assert ( - self.mamba_track_interval % self.page_size == 0 - ), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}" + assert self.mamba_track_interval % self.page_size == 0, ( + f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}" + ) assert ( max(FLA_CHUNK_SIZE, self.page_size) % min(FLA_CHUNK_SIZE, self.page_size) == 0 - ), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}" + ), ( + f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}" + ) elif not self.disable_radix_cache: # no_buffer if self.page_size is not None and self.page_size != 1: logger.warning( @@ -2749,9 +2795,9 @@ def _handle_attention_backend_compatibility(self): "Cuda graph is disabled because of using torch Flex Attention backend" ) self.disable_cuda_graph = True - assert ( - self.speculative_algorithm is None - ), "Speculative decoding is currently not supported with Flex Attention backend" + assert self.speculative_algorithm is None, ( + "Speculative decoding is currently not supported with Flex Attention backend" + ) # Whisper's encoder token padding conflicts with prefix caching. # Only disable for Whisper; other encoder-decoder models (e.g., mllama) use radix cache. @@ -3097,40 +3143,40 @@ def _handle_linear_attn_backend(self): def _handle_context_parallelism(self): if self.attn_cp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - self.tp_size % self.attn_cp_size == 0 - ), "tp_size must be divisible by attn_cp_size" - assert ( - self.tp_size % (self.dp_size * self.attn_cp_size) == 0 - ), "tp_size must be divisible by dp_size * attn_cp_size" + assert self.tp_size % self.attn_cp_size == 0, ( + "tp_size must be divisible by attn_cp_size" + ) + assert self.tp_size % (self.dp_size * self.attn_cp_size) == 0, ( + "tp_size must be divisible by dp_size * attn_cp_size" + ) - assert ( - not self.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not self.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if self.moe_dp_size > 1: # The tp_size is the world size, not the real tensor parallel size - assert ( - self.tp_size % self.moe_dp_size == 0 - ), "tp_size must be divisible by moe_dp_size" - assert ( - self.ep_size * self.moe_dp_size <= self.tp_size - ), "ep_size * moe_dp_size must be less than or equal to tp_size" + assert self.tp_size % self.moe_dp_size == 0, ( + "tp_size must be divisible by moe_dp_size" + ) + assert self.ep_size * self.moe_dp_size <= self.tp_size, ( + "ep_size * moe_dp_size must be less than or equal to tp_size" + ) assert self.pp_size == 1, "PP is not supported with context parallelism" if self.ep_size > 1: - assert ( - self.ep_size * self.moe_dp_size == self.tp_size - ), "ep_size * moe_dp_size must be equal to tp_size" + assert self.ep_size * self.moe_dp_size == self.tp_size, ( + "ep_size * moe_dp_size must be equal to tp_size" + ) - assert ( - not self.enable_aiter_allreduce_fusion - ), "Aiter allreduce fusion is not supported with context parallelism" + assert not self.enable_aiter_allreduce_fusion, ( + "Aiter allreduce fusion is not supported with context parallelism" + ) if self.attn_cp_size != self.moe_dp_size: - assert ( - self.moe_dp_size == 1 - ), "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + assert self.moe_dp_size == 1, ( + "attn_cp_size != moe_dp_size is only supported when moe_dp_size == 1" + ) def _handle_data_parallelism(self): if self.dp_size == 1: @@ -3146,9 +3192,9 @@ def _handle_data_parallelism(self): ) if self.enable_dp_lm_head: - assert ( - self.enable_dp_attention - ), "Please enable dp attention when setting enable_dp_lm_head. " + assert self.enable_dp_attention, ( + "Please enable dp attention when setting enable_dp_lm_head. " + ) def _handle_moe_kernel_config(self): if self.quantization == "mxfp8": @@ -3172,20 +3218,26 @@ def _handle_moe_kernel_config(self): "modelopt_fp8", "modelopt_mixed", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." + ], ( + f"Invalid quantization '{self.quantization}'. \nFlashInfer Cutlass MOE supports only: 'modelopt_fp4', 'modelopt_fp8', 'modelopt_mixed', or bfloat16 (None)." + ) assert self.ep_size in [ 1, self.tp_size, - ], "The expert parallel size must be 1 or the same as the tensor parallel size" + ], ( + "The expert parallel size must be 1 or the same as the tensor parallel size" + ) if self.moe_runner_backend == "flashinfer_cutedsl": - assert self.quantization in [ - "modelopt_fp4" - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4'." + assert self.quantization in ["modelopt_fp4"], ( + f"Invalid quantization '{self.quantization}'. \nFlashInfer CuteDSL MOE currently supports only: 'modelopt_fp4'." + ) assert self.ep_size in [ 1, self.tp_size, - ], "The expert parallel size must be 1 or the same as the tensor parallel size" + ], ( + "The expert parallel size must be 1 or the same as the tensor parallel size" + ) assert self.moe_a2a_backend in [ "none", "deepep", @@ -3208,7 +3260,9 @@ def _handle_moe_kernel_config(self): "modelopt_mixed", "compressed-tensors", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." + ], ( + f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM MOE supports only: 'modelopt_fp4', 'fp8', 'modelopt_fp8', 'modelopt_mixed', 'compressed-tensors', or bfloat16 (None)." + ) self.disable_shared_experts_fusion = True logger.warning( "FlashInfer TRTLLM MoE is enabled. --disable-shared-experts-fusion is automatically set." @@ -3220,7 +3274,9 @@ def _handle_moe_kernel_config(self): "mxfp8", "modelopt_fp4", None, - ], f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', or bfloat16 (None)." + ], ( + f"Invalid quantization '{self.quantization}'. \nFlashInfer TRTLLM routed MOE supports only: 'fp8', 'mxfp8', 'modelopt_fp4', or bfloat16 (None)." + ) self.disable_shared_experts_fusion = True logger.warning( "FlashInfer TRTLLM routed MoE is enabled. --disable-shared-experts-fusion is automatically set." @@ -3239,9 +3295,9 @@ def _handle_moe_kernel_config(self): "fp8", "mxfp8", ]: - assert ( - self.ep_size == 1 - ), "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + assert self.ep_size == 1, ( + "FP8/MXFP8 Cutlass MoE is only supported with ep_size == 1" + ) # TODO(yuwei): Fix piecewise cuda graph support for bypassed topk MoE backends. # Exception: GptOssForCausalLM wraps the entire MoE block in its own @@ -3328,13 +3384,13 @@ def _handle_a2a_moe(self): f"Wrong value of {fuse_mode=}, the NPU only support 1 or 2." ) elif fuse_mode == 2: - assert ( - self.quantization == "modelslim" - ), "When fuse_mode is set to 2, the NPU supports only ModelSlim quantization." + assert self.quantization == "modelslim", ( + "When fuse_mode is set to 2, the NPU supports only ModelSlim quantization." + ) if self.moe_a2a_backend == "flashinfer": - assert ( - self.enable_dp_attention and self.dp_size == self.tp_size - ), "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + assert self.enable_dp_attention and self.dp_size == self.tp_size, ( + "Flashinfer MoE A2A is only supported with dp_size == tp_size and --enable-dp-attention" + ) self.ep_size = self.tp_size logger.warning( f"Flashinfer MoE A2A is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{self.tp_size}]." @@ -3353,7 +3409,9 @@ def _handle_a2a_moe(self): assert self.moe_runner_backend in [ "flashinfer_cutlass", "flashinfer_cutedsl", - ], "Flashinfer MoE A2A is only supported with flashinfer_cutlass or flashinfer_cutedsl moe runner backend" + ], ( + "Flashinfer MoE A2A is only supported with flashinfer_cutlass or flashinfer_cutedsl moe runner backend" + ) if ( self.moe_runner_backend == "flashinfer_cutedsl" and self.max_prefill_tokens is not None @@ -3401,7 +3459,9 @@ def _handle_a2a_moe(self): if self.chunked_prefill_size > 0 and self.disaggregation_mode != "decode": assert ( self.chunked_prefill_size - ) <= envs.SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get(), "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK (default 4096) must be larger or equal to chunked_prefill_size" + ) <= envs.SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK.get(), ( + "SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK (default 4096) must be larger or equal to chunked_prefill_size" + ) def _handle_eplb_and_dispatch(self): if self.enable_eplb and (self.expert_distribution_recorder_mode is None): @@ -3426,7 +3486,9 @@ def _handle_elastic_ep(self): assert self.eplb_algorithm in [ "elasticity_aware", "elasticity_aware_hierarchical", - ], "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ], ( + "Elastic EP requires eplb_algorithm to be set to 'auto' or 'elasticity_aware(_hierarchical)'." + ) assert self.pp_size == 1, "PP size should be set to 1 under elastic EP" @@ -3435,9 +3497,9 @@ def _handle_elastic_ep(self): self.mooncake_ib_device ) if self.elastic_ep_rejoin: - assert ( - self.elastic_ep_backend is not None - ), "Elastic EP rejoin requires elastic_ep_backend to be set." + assert self.elastic_ep_backend is not None, ( + "Elastic EP rejoin requires elastic_ep_backend to be set." + ) def _handle_expert_distribution_metrics(self): if self.enable_expert_distribution_metrics and ( @@ -3806,9 +3868,9 @@ def _handle_pd_disaggregation(self): self.mamba_scheduler_strategy = "no_buffer" elif self.disaggregation_mode == "prefill": - assert ( - self.disaggregation_transfer_backend != "fake" - ), "Prefill server does not support 'fake' as the transfer backend" + assert self.disaggregation_transfer_backend != "fake", ( + "Prefill server does not support 'fake' as the transfer backend" + ) self.disable_cuda_graph = True @@ -6183,6 +6245,21 @@ def add_cli_args(parser: argparse.ArgumentParser): "radix cache, and chunked prefill.", ) + # KV-sharing fast-prefill (YOCO) opt-in flag. + parser.add_argument( + "--kv-sharing-fast-prefill", + action="store_true", + help=( + "Enable the YOCO-style fast-prefill optimization for models " + "with KV-shared layers (currently Gemma-4 family). When set, " + "the last ``num_kv_shared_layers`` decoder layers process only " + "the per-request last-extend-token rows during EXTEND-mode " + "forwards, skipping redundant compute. Requires the triton " + "attention backend and is incompatible with EAGLE/EAGLE3 " + "speculative decoding. Default off; opt-in only." + ), + ) + # Optimization/debug options parser.add_argument( "--disable-radix-cache", @@ -7029,11 +7106,35 @@ def mamba_cache_chunk_size(self) -> int: # It is used to determine the caching point in a sequence during prefill. return max(FLA_CHUNK_SIZE, self.page_size) + def _validate_kv_sharing_fast_prefill_combos(self): + """Cross-cutting validators for the YOCO opt-in flag. + + Architecture-specific validation (Gemma-4 detection, triton backend + pin, ``num_kv_shared_layers > 0``) lives in + ``_handle_model_specific_adjustments``; this function captures only + the model-agnostic constraints so they can be unit-tested in + isolation. + """ + if not self.kv_sharing_fast_prefill: + return + spec_algo = (self.speculative_algorithm or "").upper() + if spec_algo in ("EAGLE", "EAGLE3"): + raise ValueError( + "--kv-sharing-fast-prefill is incompatible with " + f"--speculative-algorithm {spec_algo}. EAGLE-family " + "algorithms read logits for every prompt token; YOCO " + "fast-prefill only computes logits at the per-request " + "last-extend-token position. Use NEXTN/FROZEN_KV_MTP " + "(Gemma-4 assistant draft) instead, or drop the flag." + ) + def check_server_args(self): # Check parallel size constraints - assert ( - self.tp_size * self.pp_size - ) % self.nnodes == 0, "tp_size must be divisible by number of nodes" + assert (self.tp_size * self.pp_size) % self.nnodes == 0, ( + "tp_size must be divisible by number of nodes" + ) + + self._validate_kv_sharing_fast_prefill_combos() assert ( self.pp_max_micro_batch_size is None or self.pp_max_micro_batch_size >= 1 @@ -7053,7 +7154,9 @@ def check_server_args(self): if self.pp_size > 1: assert ( self.disable_overlap_schedule and self.speculative_algorithm is None - ), "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + ), ( + "Pipeline parallelism is not compatible with overlap schedule, speculative decoding" + ) assert not ( self.dp_size > 1 and self.nnodes != 1 and not self.enable_dp_attention @@ -7080,32 +7183,32 @@ def check_server_args(self): # Check speculative decoding if self.speculative_algorithm is not None: - assert ( - not self.enable_mixed_chunk - ), "enable_mixed_chunk is required for speculative decoding" + assert not self.enable_mixed_chunk, ( + "enable_mixed_chunk is required for speculative decoding" + ) # Check chunked prefill # Skip validation if chunked prefill is disabled (i.e., size <= 0). # Skip validation if disaggregation mode is decode. if self.chunked_prefill_size > 0 and self.disaggregation_mode != "decode": - assert ( - self.chunked_prefill_size % self.page_size == 0 - ), "chunked_prefill_size must be divisible by page_size" + assert self.chunked_prefill_size % self.page_size == 0, ( + "chunked_prefill_size must be divisible by page_size" + ) # Check pdmux if self.enable_pdmux: - assert ( - self.pp_size == 1 - ), "PD-Multiplexing is only supported with pipeline parallelism disabled (pp_size=1)." - assert ( - self.chunked_prefill_size == -1 - ), "PD-Multiplexing is not compatible with chunked prefill." - assert ( - self.disaggregation_mode == "null" - ), "PD-Multiplexing is not compatible with disaggregation mode." - assert ( - self.disable_overlap_schedule - ), "PD-Multiplexing is not compatible with overlap schedule." + assert self.pp_size == 1, ( + "PD-Multiplexing is only supported with pipeline parallelism disabled (pp_size=1)." + ) + assert self.chunked_prefill_size == -1, ( + "PD-Multiplexing is not compatible with chunked prefill." + ) + assert self.disaggregation_mode == "null", ( + "PD-Multiplexing is not compatible with disaggregation mode." + ) + assert self.disable_overlap_schedule, ( + "PD-Multiplexing is not compatible with overlap schedule." + ) # NOTE: CUDA Green Context may encounter potential issues with CudaGraph on torch 2.7.x – 2.8.x, leading to performance degradation. import torch @@ -7131,7 +7234,9 @@ def check_server_args(self): assert self.schedule_policy in [ "fcfs", "lof", - ], f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{self.schedule_policy}' is not supported." + ], ( + f"To use priority scheduling, schedule_policy must be 'fcfs' or 'lof'. '{self.schedule_policy}' is not supported." + ) if self.default_priority_value is None: logger.warning( "--default-priority-value is not set while --enable-priority-scheduling is enabled. " @@ -7153,9 +7258,9 @@ def check_server_args(self): validate_hisparse(self) - assert ( - self.schedule_conservativeness >= 0 - ), "schedule_conservativeness must be non-negative" + assert self.schedule_conservativeness >= 0, ( + "schedule_conservativeness must be non-negative" + ) if self.model_impl == "mindspore": assert is_npu(), "MindSpore model impl is only supported on Ascend npu." @@ -7274,9 +7379,9 @@ def check_lora_server_args(self): pinned=False, ) elif isinstance(lora_path, dict): - assert ( - "lora_name" in lora_path and "lora_path" in lora_path - ), f"When providing LoRA paths as a list of dict, each dict should contain 'lora_name' and 'lora_path' keys. Got: {lora_path}" + assert "lora_name" in lora_path and "lora_path" in lora_path, ( + f"When providing LoRA paths as a list of dict, each dict should contain 'lora_name' and 'lora_path' keys. Got: {lora_path}" + ) lora_ref = LoRARef( lora_id=LoRARef.deterministic_id( lora_path["lora_name"], lora_path["lora_path"] @@ -7314,14 +7419,16 @@ def check_lora_server_args(self): if self.lora_target_modules: self.lora_target_modules = set(self.lora_target_modules) if "all" in self.lora_target_modules: - assert ( - len(self.lora_target_modules) == 1 - ), "If 'all' is specified in --lora-target-modules, it should be the only module specified." + assert len(self.lora_target_modules) == 1, ( + "If 'all' is specified in --lora-target-modules, it should be the only module specified." + ) # Ensure sufficient information is provided for LoRA initialization. assert self.lora_paths or ( self.max_lora_rank and self.lora_target_modules - ), "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + ), ( + "When no initial --lora-paths is provided, you need to specify both --max-lora-rank and --lora-target-modules for LoRA initialization." + ) # Validate max_loaded_loras if self.max_loaded_loras is not None: @@ -7343,9 +7450,9 @@ def check_lora_server_args(self): if self.lora_use_virtual_experts: logger.info("Virtual expert computation enabled.") - assert ( - self.lora_drain_wait_threshold >= 0.0 - ), "--lora-drain-wait-threshold must be non-negative." + assert self.lora_drain_wait_threshold >= 0.0, ( + "--lora-drain-wait-threshold must be non-negative." + ) def validate_buckets_rule(self, arg_name: str, buckets_rule: List[str]): if not buckets_rule: @@ -7357,43 +7464,45 @@ def validate_buckets_rule(self, arg_name: str, buckets_rule: List[str]): "tse", "default", "custom", - ], f"Unsupported {arg_name} rule type: '{rule}'. Must be one of: 'tse', 'default', 'custom'" + ], ( + f"Unsupported {arg_name} rule type: '{rule}'. Must be one of: 'tse', 'default', 'custom'" + ) if rule == "tse": - assert ( - len(buckets_rule) == 4 - ), f"{arg_name} TSE rule requires exactly 4 parameters: ['tse', middle, base, count], got {len(buckets_rule)}" + assert len(buckets_rule) == 4, ( + f"{arg_name} TSE rule requires exactly 4 parameters: ['tse', middle, base, count], got {len(buckets_rule)}" + ) try: middle = float(buckets_rule[1]) base = float(buckets_rule[2]) count = int(buckets_rule[3]) except (ValueError, IndexError): - assert ( - False - ), f"{arg_name} TSE rule parameters must be: ['tse', , , ]" + assert False, ( + f"{arg_name} TSE rule parameters must be: ['tse', , , ]" + ) assert base > 1, f"{arg_name} TSE base must be larger than 1, got: {base}" assert count > 0, f"{arg_name} TSE count must be positive, got: {count}" assert middle > 0, f"{arg_name} TSE middle must be positive, got: {middle}" elif rule == "default": - assert ( - len(buckets_rule) == 1 - ), f"{arg_name} default rule should only have one parameter: ['default'], got {len(buckets_rule)}" + assert len(buckets_rule) == 1, ( + f"{arg_name} default rule should only have one parameter: ['default'], got {len(buckets_rule)}" + ) elif rule == "custom": - assert ( - len(buckets_rule) >= 2 - ), f"{arg_name} custom rule requires at least one bucket value: ['custom', value1, ...]" + assert len(buckets_rule) >= 2, ( + f"{arg_name} custom rule requires at least one bucket value: ['custom', value1, ...]" + ) try: bucket_values = [float(x) for x in buckets_rule[1:]] except ValueError: assert False, f"{arg_name} custom rule bucket values must be numeric" - assert len(set(bucket_values)) == len( - bucket_values - ), f"{arg_name} custom rule bucket values should not contain duplicates" - assert all( - val >= 0 for val in bucket_values - ), f"{arg_name} custom rule bucket values should be non-negative" + assert len(set(bucket_values)) == len(bucket_values), ( + f"{arg_name} custom rule bucket values should not contain duplicates" + ) + assert all(val >= 0 for val in bucket_values), ( + f"{arg_name} custom rule bucket values should be non-negative" + ) def adjust_mem_fraction_for_vlm(self, model_config): vision_config = getattr(model_config.hf_config, "vision_config", None) diff --git a/test/registered/unit/server_args/test_kv_sharing_fast_prefill.py b/test/registered/unit/server_args/test_kv_sharing_fast_prefill.py new file mode 100644 index 000000000000..4c9a2a43114f --- /dev/null +++ b/test/registered/unit/server_args/test_kv_sharing_fast_prefill.py @@ -0,0 +1,114 @@ +"""Unit tests for the ``--kv-sharing-fast-prefill`` server flag (YOCO opt-in). + +These tests cover the dataclass field, argparse registration, ``ModelConfig`` +plumbing, and the cross-cutting validator that runs inside +``ServerArgs.check_server_args``. Architecture-specific validation +(Gemma-4 arch + triton attention backend + ``num_kv_shared_layers > 0``) +runs inside ``_handle_model_specific_adjustments`` and requires a real +Hugging Face config; that path is exercised by integration tests against +the live ``google/gemma-4-31B-it`` checkpoint and is not duplicated here. +""" + +import unittest +from unittest.mock import patch + +from sglang.srt.server_args import ServerArgs, prepare_server_args +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + +# Mock get_device so all tests run on CPU-only CI runners. +_mock_device = patch("sglang.srt.server_args.get_device", return_value="cuda") +_mock_device.start() + + +class TestKvSharingFastPrefillFlag(CustomTestCase): + """Plumbing-level tests: dataclass field, argparse, ModelConfig.""" + + def test_default_is_disabled(self): + server_args = ServerArgs(model_path="dummy") + self.assertFalse(server_args.kv_sharing_fast_prefill) + + def test_cli_flag_parses_to_true(self): + server_args = prepare_server_args( + [ + "--model-path", + "dummy", + "--kv-sharing-fast-prefill", + ] + ) + self.assertTrue(server_args.kv_sharing_fast_prefill) + + def test_cli_help_lists_flag(self): + # The presence of the help text is the simplest way to assert the + # arg is registered without invoking the full parser. + import argparse + + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + # ``--kv-sharing-fast-prefill`` should map to dest ``kv_sharing_fast_prefill``. + kvsp = next( + ( + a + for a in parser._actions + if "--kv-sharing-fast-prefill" in a.option_strings + ), + None, + ) + self.assertIsNotNone(kvsp) + self.assertEqual(kvsp.dest, "kv_sharing_fast_prefill") + self.assertFalse(kvsp.default) + + def test_model_config_plumbing(self): + # ``ModelConfig.__init__`` must accept the new keyword. + from sglang.srt.configs.model_config import ModelConfig + + self.assertIn( + "kv_sharing_fast_prefill", + ModelConfig.__init__.__code__.co_varnames, + ) + + +class TestKvSharingFastPrefillValidators(CustomTestCase): + """Cross-cutting validators that fire in ``check_server_args``.""" + + def _server_args(self, **kwargs): + return ServerArgs(model_path="dummy", **kwargs) + + def test_eagle_combo_rejected(self): + sa = self._server_args( + kv_sharing_fast_prefill=True, + speculative_algorithm="EAGLE", + ) + with self.assertRaisesRegex(ValueError, "incompatible with"): + sa._validate_kv_sharing_fast_prefill_combos() + + def test_eagle3_combo_rejected(self): + sa = self._server_args( + kv_sharing_fast_prefill=True, + speculative_algorithm="EAGLE3", + ) + with self.assertRaisesRegex(ValueError, "incompatible with"): + sa._validate_kv_sharing_fast_prefill_combos() + + def test_nextn_combo_accepted(self): + # NEXTN (the algorithm that promotes to FROZEN_KV_MTP for Gemma-4 + # assistant drafts) is allowed and must not raise. + sa = self._server_args( + kv_sharing_fast_prefill=True, + speculative_algorithm="NEXTN", + ) + sa._validate_kv_sharing_fast_prefill_combos() # should not raise + + def test_flag_off_does_not_validate(self): + # When the flag is off, even EAGLE is fine for this validator. + sa = self._server_args( + kv_sharing_fast_prefill=False, + speculative_algorithm="EAGLE", + ) + sa._validate_kv_sharing_fast_prefill_combos() # should not raise + + +if __name__ == "__main__": + unittest.main()