From 0c065671c94c17851f691da1551e3cab8bd99a13 Mon Sep 17 00:00:00 2001 From: cctry Date: Sun, 21 Jun 2026 20:45:16 -0700 Subject: [PATCH] [Spec] Redo: split init_backends; account draft weights in --mem-fraction-static (#28855) Co-authored-by: cctry --- python/sglang/bench_one_batch.py | 3 +- python/sglang/srt/managers/scheduler.py | 19 ++++--- python/sglang/srt/managers/tp_worker.py | 16 ++++-- .../sglang/srt/model_executor/model_runner.py | 56 ++++++++----------- .../model_runner_kv_cache_mixin.py | 53 ++++++++++++------ .../srt/speculative/base_spec_worker.py | 22 +++++--- .../srt/speculative/dflash_worker_v2.py | 17 ++++-- .../sglang/srt/speculative/eagle_worker_v2.py | 25 ++++++--- .../speculative/frozen_kv_mtp_worker_v2.py | 16 ++++-- .../multi_layer_eagle_worker_v2.py | 20 +++++-- .../srt/speculative/standalone_worker_v2.py | 10 +++- .../server_fixtures/spec_eagle_fixture.py | 2 +- .../8-gpu-models/test_deepseek_v32.py | 2 +- .../cp/test_deepseek_v32_cp_single_node.py | 4 +- .../models_e2e/test_deepseek_v3_mtp.py | 2 +- .../models_e2e/test_dsa_dsv32_dp_mtp.py | 2 +- .../models_e2e/test_dsa_dsv32_tp_mtp.py | 2 +- .../test_unified_radix_cache_kl_dsv4.py | 2 +- 18 files changed, 167 insertions(+), 106 deletions(-) diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index ad154fd50..aec431267 100644 --- a/python/sglang/bench_one_batch.py +++ b/python/sglang/bench_one_batch.py @@ -324,7 +324,8 @@ def load_model(server_args, port_args, gpu_id, tp_rank): else: model_runner = ModelRunner(**runner_kwargs) model_runner.alloc_memory_pool() - model_runner.init_backends() + model_runner.init_attention_backends() + model_runner.init_cuda_graphs() rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}") tokenizer = get_tokenizer( server_args.tokenizer_path, diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index d32d44cbc..c76a86e15 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -844,11 +844,17 @@ class Scheduler( token_to_kv_pool_allocator=allocator, ) - def init_all_backends(self): - """Initialize attention backends and capture cuda graphs for all workers.""" - self.tp_worker.init_backends() + def init_all_attention_backends(self): + """Initialize attention backends for all workers.""" + self.tp_worker.init_attention_backends() if self.draft_worker is not None: - self.draft_worker.init_backends() + self.draft_worker.init_attention_backends() + + def init_all_cuda_graphs(self): + """Capture cuda graphs for all workers.""" + self.tp_worker.init_cuda_graphs() + if self.draft_worker is not None: + self.draft_worker.init_cuda_graphs() def init_model_worker(self): # Load model weights. @@ -859,12 +865,11 @@ class Scheduler( self.maybe_init_draft_worker() # Allocate KV cache pools for all workers. - # Memory profiling now sees all loaded weights. self.init_memory_pools() - # Initialize attention backends and capture cuda graphs. # TODO: make memory profile consider cuda graph memory as well - self.init_all_backends() + self.init_all_attention_backends() + self.init_all_cuda_graphs() # Dispatch the model worker if self.spec_algorithm.is_none(): diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index e3afad69c..d1654fe6a 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -333,11 +333,19 @@ class TpModelWorker(BaseTpWorker): ) assert max_req_len > 0, "Memory pool size is too small" - def init_backends(self, disable_cuda_graph: bool = False): - """Initialize attention backends and capture cuda graphs.""" - self.model_runner.init_backends(disable_cuda_graph=disable_cuda_graph) + def init_attention_backends(self): + """Initialize attention backends for all model runners.""" + self.model_runner.init_attention_backends() for mr in self.model_runner_list[1:]: - mr.init_backends(disable_cuda_graph=disable_cuda_graph) + mr.init_attention_backends() + + def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True): + """Capture cuda graphs for all model runners.""" + self.model_runner.init_cuda_graphs( + capture_decode_cuda_graph=capture_decode_cuda_graph + ) + for mr in self.model_runner_list[1:]: + mr.init_cuda_graphs(capture_decode_cuda_graph=capture_decode_cuda_graph) def _init_model_config(self): from sglang.srt.configs.model_config import ModelConfig diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index f3fe4daa0..3becbca04 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -788,18 +788,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Deduce KV cache dtype self.configure_kv_cache_dtype() - # Snapshot free memory at the end of the weight-load phase. KV-pool - # profiling uses this instead of measuring at alloc_memory_pool() - # time: draft-model weights load between the two phases and must stay - # outside the --mem-fraction-static budget (deployments tune the - # fraction assuming draft weights live in the non-static slack). - self.post_model_load_memory = get_available_gpu_memory( - self.device, - self.gpu_id, - distributed=get_world_group().world_size > 1, - cpu_group=get_world_group().cpu_group, - ) - def get_pp_proxy_topk_size(self) -> Optional[int]: hf_config = self.model_config.hf_text_config if ( @@ -862,18 +850,13 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.graph_mem_usage = 0 self.prefill_cuda_graph_runner = None - def init_backends(self, disable_cuda_graph: bool = False): - """Initialize attention backends and capture cuda graphs.""" - server_args = self.server_args - + def init_attention_backends(self): + """Initialize attention backends only (no cuda graph capture).""" # TODO: Refactor device-specific init branches into platform interface (separate PR). # Must be called BEFORE init_decode_cuda_graph() so CUDA graph capture # runs with aux hidden state capture enabled. self.init_aux_hidden_state_capture() - # Device-specific attention-backend init. cg_supported gates decode - # cuda-graph capture (out-of-tree / unknown platforms may not support it). - cg_supported = True if self.device == "cuda" or self.device == "musa": self.init_cublas() self.init_attention_backend() @@ -892,12 +875,15 @@ class ModelRunner(ModelRunnerKVCacheMixin): get_world_group().world_size, get_world_group().cpu_group, ) - elif current_platform.is_out_of_tree(): - self.init_attention_backend() - cg_supported = current_platform.support_cuda_graph() else: self.init_attention_backend() - cg_supported = False + + def init_cuda_graphs(self, capture_decode_cuda_graph: bool = True): + """Capture cuda graphs. Requires init_attention_backends() to have run. + + Spec draft runners pass capture_decode_cuda_graph=False + because they capture their own decode-style graphs separately. + """ # The eager (no-cuda-graph) phase runner, built AFTER the attention # backend so its __init__ can warm up kernels (run-once) and allocate the @@ -912,23 +898,27 @@ class ModelRunner(ModelRunnerKVCacheMixin): # eager buffer allocated above. (init_prefill_cuda_graph routes prefill # to the eager runner when the prefill graph is disabled.) self.init_prefill_cuda_graph() - if not disable_cuda_graph and cg_supported: - self.init_decode_cuda_graph() + + self.decode_cuda_graph_runner = None + self.graph_mem_usage = 0 + + if capture_decode_cuda_graph: + if self.device in ("cuda", "musa", "cpu", "npu"): + self.init_decode_cuda_graph() + elif ( + current_platform.is_out_of_tree() + and current_platform.support_cuda_graph() + ): + self.init_decode_cuda_graph() else: - self.decode_cuda_graph_runner = None - self.graph_mem_usage = 0 - if disable_cuda_graph: - # Decode cuda graph disabled: route eager decode through the - # EagerRunner (the dispatch gate isinstance(..., EagerRunner) keeps - # _forward_raw off any replay branch). self.decode_cuda_graph_runner = self.eager_runner # Register forward hooks AFTER cuda-graph capture so their tensor ops are # not traced into any captured graph — capture stays hook-free and hooks # fire only on the eager forward path (capture replay never runs Python # hooks anyway). - if server_args.forward_hooks: - register_forward_hooks(self.model, server_args.forward_hooks) + if self.server_args.forward_hooks: + register_forward_hooks(self.model, self.server_args.forward_hooks) self.prealloc_symmetric_memory_pool() 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 685f1d405..db4c8ad8e 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 @@ -1,6 +1,7 @@ from __future__ import annotations import logging +import math from typing import TYPE_CHECKING import torch @@ -82,24 +83,40 @@ _is_hip = is_hip() class ModelRunnerKVCacheMixin: def _profile_available_bytes(self: ModelRunner, pre_model_load_memory: int) -> int: - # Use the snapshot taken at the end of this runner's weight-load phase, - # not the current free memory: draft-model weights loaded after that - # point are charged to the non-static slack, not the static budget. - post_model_load_memory = getattr(self, "post_model_load_memory", None) - if post_model_load_memory is None: - post_model_load_memory = get_available_gpu_memory( - self.device, - self.gpu_id, - distributed=get_world_group().world_size > 1, - cpu_group=get_world_group().cpu_group, - ) + # KV pool budget = currently-free GPU memory minus the non-static runtime + # slack (pre_model_load_memory * (1 - mem_fraction_static)). Whatever is + # already resident (model weights, etc.) is thus charged against it. + available_gpu_memory = get_available_gpu_memory( + self.device, + self.gpu_id, + distributed=get_world_group().world_size > 1, + cpu_group=get_world_group().cpu_group, + ) - rest_memory = post_model_load_memory - pre_model_load_memory * ( + rest_memory = available_gpu_memory - pre_model_load_memory * ( 1 - self.mem_fraction_static ) if self.mambaish_config is not None: rest_memory = self.handle_max_mamba_cache(rest_memory) + # Loaded weights (target + draft) can exceed the static budget + if rest_memory <= 0: + minimum_mem_fraction_static = ( + 1 - available_gpu_memory / pre_model_load_memory + ) + suggested_mem_fraction_static = ( + math.ceil(minimum_mem_fraction_static * 1000) / 1000 + ) + raise ValueError( + f"Loaded weights leave no GPU memory for the KV cache under " + f"--mem-fraction-static={self.mem_fraction_static}. " + f"Raise --mem-fraction-static above " + f"{suggested_mem_fraction_static:.3f} " + f"(minimum viable = 1 - available/pre = " + f"{minimum_mem_fraction_static:.4f}). If using speculative " + f"decoding, draft weights are now counted." + ) + return int(rest_memory * (1 << 30)) # return in bytes def handle_max_mamba_cache(self: ModelRunner, total_rest_memory): @@ -560,9 +577,9 @@ class ModelRunnerKVCacheMixin: self.model_config.hf_text_config.swa_num_key_value_heads // get_attention_tp_size(), ), - "swa_head_dim": self.model_config.hf_text_config.swa_head_dim, - "swa_v_head_dim": self.model_config.hf_text_config.swa_v_head_dim, - "v_head_dim": self.model_config.hf_text_config.v_head_dim, + "swa_head_dim": self.model_config.swa_head_dim, + "swa_v_head_dim": self.model_config.swa_v_head_dim, + "v_head_dim": self.model_config.v_head_dim, } self.token_to_kv_pool = SWAKVPool( size=self.full_max_total_num_tokens, @@ -683,9 +700,9 @@ class ModelRunnerKVCacheMixin: self.model_config.hf_text_config.swa_num_key_value_heads // get_attention_tp_size(), ), - "swa_head_dim": self.model_config.hf_text_config.swa_head_dim, - "swa_v_head_dim": self.model_config.hf_text_config.swa_v_head_dim, - "v_head_dim": self.model_config.hf_text_config.v_head_dim, + "swa_head_dim": self.model_config.swa_head_dim, + "swa_v_head_dim": self.model_config.swa_v_head_dim, + "v_head_dim": self.model_config.v_head_dim, } self.token_to_kv_pool = SWAKVPool( size=self.full_max_total_num_tokens, diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index acaf61332..f70bd3e01 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -78,15 +78,16 @@ class EagleDraftWorkerBase(ABC): def alloc_memory_pool(self, **kwargs): pass - def init_backends(self): - """Initialize standard backends (no cuda graphs) then draft-specific backends. - - Subclasses should wrap this with their context managers (draft_tp_context, - speculative_moe_backend_context, etc.) rather than reimplementing the logic. - """ - self.draft_worker.init_backends(disable_cuda_graph=True) + def init_attention_backends(self): + """Subclasses wrap this with their context managers (draft_tp_context, + speculative_moe_backend_context, etc.) rather than reimplementing it.""" + self.draft_worker.init_attention_backends() self.init_attention_backend() - self.init_cuda_graphs() + + def init_cuda_graphs(self): + """Capture draft graphs (decode disabled on the draft TpModelWorker).""" + self.draft_worker.init_cuda_graphs(capture_decode_cuda_graph=False) + self._capture_cuda_graphs() def prepare_for_draft_extend( self, @@ -295,7 +296,10 @@ class BaseSpecWorker(ABC): def alloc_memory_pool(self, **kwargs): pass - def init_backends(self): + def init_attention_backends(self): + pass + + def init_cuda_graphs(self): pass def on_verify_complete_cpu( diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index f9eb180f2..c92a5d702 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -293,18 +293,23 @@ class DFlashWorkerV2(BaseSpecWorker): token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) - def init_backends(self): - disable_cuda_graph = False - if is_cuda() and not self.server_args.disable_cuda_graph: + def init_attention_backends(self): + self._draft_worker.init_attention_backends() + + def init_cuda_graphs(self): + capture_decode_cuda_graph = not self.server_args.disable_cuda_graph + if is_cuda() and capture_decode_cuda_graph: available_mem = get_available_gpu_memory(self.device, self.gpu_id) - disable_cuda_graph = available_mem < 1.0 - if disable_cuda_graph: + if available_mem < 1.0: + capture_decode_cuda_graph = False logger.warning( "Disable DFLASH draft cuda graph because only %.2f GB GPU " "memory is available after target backend initialization.", available_mem, ) - self._draft_worker.init_backends(disable_cuda_graph=disable_cuda_graph) + self._draft_worker.init_cuda_graphs( + capture_decode_cuda_graph=capture_decode_cuda_graph + ) def _init_fused_kv_helper(self) -> None: """Initialize the fused KV materialization helper with pre-stacked weights.""" diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 86af77e7d..2a1951686 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -243,15 +243,21 @@ class EagleDraftWorker(EagleDraftWorkerBase): f"({draft_vocab_size}) != target vocab ({target_vocab_size})." ) - def init_backends(self): + def init_attention_backends(self): with self.draft_tp_context( self.draft_runner.tp_group ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): - self.draft_worker.init_backends(disable_cuda_graph=True) + self.draft_worker.init_attention_backends() self.init_attention_backend() + + def init_cuda_graphs(self): + with self.draft_tp_context( + self.draft_runner.tp_group + ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): + self.draft_worker.init_cuda_graphs(capture_decode_cuda_graph=False) if check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE): self.draft_runner.init_prefill_cuda_graph(force_for_draft_worker=True) - self.init_cuda_graphs() + self._capture_cuda_graphs() if (c := self.draft_runner.canary_manager) is not None: c.mark_init_finished() @@ -359,8 +365,8 @@ class EagleDraftWorker(EagleDraftWorkerBase): self.draft_runner.attn_backend = self.draft_extend_attn_backend self.tree_mask_mode = TreeMaskMode.FULL_MASK - def init_cuda_graphs(self): - """Capture cuda graphs.""" + def _capture_cuda_graphs(self): + """Capture the draft worker's own cuda graphs (decode + draft-extend).""" self.cuda_graph_runner = None self.cuda_graph_runner_for_draft_extend = None @@ -992,8 +998,11 @@ class EAGLEWorkerV2(BaseSpecWorker): self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - def init_backends(self): - self._draft_worker.init_backends() + def init_attention_backends(self): + self._draft_worker.init_attention_backends() + + def init_cuda_graphs(self): + self._draft_worker.init_cuda_graphs() # Build adaptive runtime states after target and draft backends exist. if self.adaptive_controller is not None: with ( @@ -1241,7 +1250,7 @@ class EAGLEWorkerV2(BaseSpecWorker): cuda_graph_bs=cuda_graph_bs, ): self._draft_worker.init_attention_backend() - self._draft_worker.init_cuda_graphs() + self._draft_worker._capture_cuda_graphs() # Build target attention backend and CUDA graph runner target_model_runner = self._target_worker.model_runner diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index 3be413951..d8d4b62ac 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -197,16 +197,24 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): ), ) - def init_backends(self): + def init_attention_backends(self): with ( self.draft_tp_context(self.draft_model_runner.tp_group), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(), ): - TpModelWorker.init_backends(self, disable_cuda_graph=True) + TpModelWorker.init_attention_backends(self) self.draft_attn_backend = self._init_draft_attn_backend() self.draft_model_runner.draft_attn_backend = self.draft_attn_backend - self.init_cuda_graphs() + + def init_cuda_graphs(self): + with ( + self.draft_tp_context(self.draft_model_runner.tp_group), + speculative_moe_backend_context(), + speculative_moe_a2a_backend_context(), + ): + TpModelWorker.init_cuda_graphs(self, capture_decode_cuda_graph=False) + self._capture_cuda_graphs() @property def draft_model_runner(self): @@ -342,7 +350,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): with self._frozen_kv_target_view(forward_batch): self.draft_attn_backend.init_forward_metadata_out_graph(fb_view) - def init_cuda_graphs(self) -> None: + def _capture_cuda_graphs(self) -> None: if cuda_graph_fully_disabled() or self.speculative_num_steps <= 1: return if self.target_worker.device != "cuda": 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 a0c884eb3..0cb0d7847 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -200,11 +200,17 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): device=self.device, ) - def init_backends(self): + def init_attention_backends(self): with self.draft_tp_context( self.draft_runner_list[0].tp_group ), speculative_moe_backend_context(): - super().init_backends() + super().init_attention_backends() + + def init_cuda_graphs(self): + with self.draft_tp_context( + self.draft_runner_list[0].tp_group + ), speculative_moe_backend_context(): + super().init_cuda_graphs() def mtp_model_runner(self, step: int): return self.draft_runner_list[step] @@ -233,8 +239,7 @@ class MultiLayerEagleDraftWorker(EagleDraftWorkerBase): self.draft_extend_attn_backend_list[-1] ) - def init_cuda_graphs(self): - """Capture cuda graphs.""" + def _capture_cuda_graphs(self): self.cuda_graph_runner = None self.cuda_graph_runner_for_draft_extend = None @@ -721,8 +726,11 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): self.req_to_token_pool = req_to_token_pool self.token_to_kv_pool_allocator = token_to_kv_pool_allocator - def init_backends(self): - self._draft_worker.init_backends() + def init_attention_backends(self): + self._draft_worker.init_attention_backends() + + def init_cuda_graphs(self): + self._draft_worker.init_cuda_graphs() @property def target_worker(self): diff --git a/python/sglang/srt/speculative/standalone_worker_v2.py b/python/sglang/srt/speculative/standalone_worker_v2.py index 3bccdd475..dd6e9770d 100644 --- a/python/sglang/srt/speculative/standalone_worker_v2.py +++ b/python/sglang/srt/speculative/standalone_worker_v2.py @@ -131,11 +131,17 @@ class StandaloneDraftWorker(EagleDraftWorker): self.init_token_map() self.init_lm_head() - def init_backends(self): + def init_attention_backends(self): with self.draft_tp_context( self.draft_runner.tp_group ), speculative_moe_backend_context(): - super().init_backends() + super().init_attention_backends() + + def init_cuda_graphs(self): + with self.draft_tp_context( + self.draft_runner.tp_group + ), speculative_moe_backend_context(): + super().init_cuda_graphs() def init_lm_head(self): """Override to prevent sharing embeddings and lm_head with target model.""" diff --git a/python/sglang/test/server_fixtures/spec_eagle_fixture.py b/python/sglang/test/server_fixtures/spec_eagle_fixture.py index ddb36c704..10426ecbd 100644 --- a/python/sglang/test/server_fixtures/spec_eagle_fixture.py +++ b/python/sglang/test/server_fixtures/spec_eagle_fixture.py @@ -58,7 +58,7 @@ class SpecEagleServerBase(CustomTestCase): attention_backend = "flashinfer" # Primary axis: False -> overlap scheduler; True -> synchronous (non-overlap). disable_overlap = False - mem_fraction_static = 0.75 + mem_fraction_static = 0.85 max_running_requests = 8 chunked_prefill_size = 128 # bf16 rather than fp16: fp16 activations can overflow (-> Inf -> NaN) on diff --git a/test/registered/8-gpu-models/test_deepseek_v32.py b/test/registered/8-gpu-models/test_deepseek_v32.py index c4b184bac..0cf1314fa 100644 --- a/test/registered/8-gpu-models/test_deepseek_v32.py +++ b/test/registered/8-gpu-models/test_deepseek_v32.py @@ -53,7 +53,7 @@ class TestDeepseekV32(unittest.TestCase): "--speculative-num-steps=3", "--speculative-eagle-topk=1", "--speculative-num-draft-tokens=4", - "--mem-frac=0.7", + "--mem-frac=0.85", ] variants = [ # Variant: "dp" - Standard TP=8 + DP=8 with dp-attention diff --git a/test/registered/cp/test_deepseek_v32_cp_single_node.py b/test/registered/cp/test_deepseek_v32_cp_single_node.py index 5b582e827..5b7105d83 100644 --- a/test/registered/cp/test_deepseek_v32_cp_single_node.py +++ b/test/registered/cp/test_deepseek_v32_cp_single_node.py @@ -43,7 +43,7 @@ class TestDeepseekV32CPInSeqSplit(CustomTestCase): "--speculative-num-draft-tokens", "4", "--mem-frac", - "0.7", + "0.85", "--cuda-graph-max-bs", "32", "--max-running-requests", @@ -109,7 +109,7 @@ class TestDeepseekV32CPRoundRobinSplit(CustomTestCase): "--speculative-num-draft-tokens", "4", "--mem-frac", - "0.7", + "0.85", "--cuda-graph-max-bs", "32", "--max-running-requests", diff --git a/test/registered/models_e2e/test_deepseek_v3_mtp.py b/test/registered/models_e2e/test_deepseek_v3_mtp.py index fb77d5bae..595413d52 100644 --- a/test/registered/models_e2e/test_deepseek_v3_mtp.py +++ b/test/registered/models_e2e/test_deepseek_v3_mtp.py @@ -31,7 +31,7 @@ _OTHER_ARGS = [ '{"enable_multithread_load": true, "num_threads": 64}', ] if not is_in_amd_ci(): - _OTHER_ARGS += ["--mem-frac", "0.7"] + _OTHER_ARGS += ["--mem-frac", "0.85"] class TestDeepseekV3MTP(GSM8KMixin, DefaultServerBase): diff --git a/test/registered/models_e2e/test_dsa_dsv32_dp_mtp.py b/test/registered/models_e2e/test_dsa_dsv32_dp_mtp.py index 6b9314528..785390fca 100644 --- a/test/registered/models_e2e/test_dsa_dsv32_dp_mtp.py +++ b/test/registered/models_e2e/test_dsa_dsv32_dp_mtp.py @@ -19,7 +19,7 @@ class TestDeepseekV32DPMTP( DsaMtpServerBase, DsaMtpEvalConfigDefaults, GSM8KMixin, SpecDecodingMixin ): model = "deepseek-ai/DeepSeek-V3.2" - mem_fraction_static = 0.7 + mem_fraction_static = 0.85 enable_dp_attention = True bs_1_speed_thres = 90 diff --git a/test/registered/models_e2e/test_dsa_dsv32_tp_mtp.py b/test/registered/models_e2e/test_dsa_dsv32_tp_mtp.py index 4a0d1150f..9bcd2a7bb 100644 --- a/test/registered/models_e2e/test_dsa_dsv32_tp_mtp.py +++ b/test/registered/models_e2e/test_dsa_dsv32_tp_mtp.py @@ -19,7 +19,7 @@ class TestDeepseekV32TPMTP( DsaMtpServerBase, DsaMtpEvalConfigDefaults, GSM8KMixin, SpecDecodingMixin ): model = "deepseek-ai/DeepSeek-V3.2" - mem_fraction_static = 0.7 + mem_fraction_static = 0.85 bs_1_speed_thres = 180 diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py index 14e60dd6c..ec9090f5e 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py +++ b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_dsv4.py @@ -200,7 +200,7 @@ class TestUnifiedDeepSeekV4FlashEagleHiCacheL3(AccuracyTwoPassMixin, CustomTestC "--chunked-prefill-size", "8192", "--mem-fraction-static", - "0.9", + "0.95", "--disable-shared-experts-fusion", "--enable-hierarchical-cache", "--hicache-ratio",