From 4ea227fa91fd569082947fae22dcf1dd908174e2 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Wed, 5 Aug 2026 19:32:24 -0700 Subject: [PATCH] config: the draft runner carries its own attention backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `build_draft_tp_worker` built a `ServerArgs` variant whose only job was to make four config reads answer with the draft's backend instead of the target's, and published it for the duration of the build so the bags agreed. The backend is a per-runner fact — target and draft coexist in one process — so it moves onto the runner, and the variant and the construction-time publish both go away. `ModelRunner` takes `draft_attention_backend` and resolves the runner's effective value once (`resolve_draft_attention_backend`: the algorithm's resolved backend, else `--speculative-draft-attention-backend`, else None for a target runner); `TpModelWorker` threads it to both runner constructions. `resolve_attention_backend_strs` reads it off the runner, and `ModelRunner` stamps the resolved pair *before* building backends so a backend can read it while it constructs — which is what the FlashInfer KV-access check needs now that it no longer asks the config. `configure_kv_cache_dtype` and the draft backend factory read the runner too. One latent bug falls out: the non-hybrid branch of the backend build ignored the resolved pair and re-read `server_args.attention_backend`, which is why the variant had to set that field as well as the split pair. It now uses the value that was resolved for the runner. `draft_server_args_overrides` and the `preserve_config()` publish switch are deleted; with them goes the last production `ServerArgs.derive` outside pre-publish config building, and the last construction-time publish. The chunked-prefix gate the target resolved simply stays in the bags, since nothing re-projects them. --- .../layers/attention/attention_registry.py | 23 ++---- .../layers/attention/flashinfer_backend.py | 5 +- python/sglang/srt/managers/tp_worker.py | 5 ++ .../srt/model_executor/forward_batch_info.py | 2 +- .../sglang/srt/model_executor/model_runner.py | 40 +++++++-- .../attention_backend_setup.py | 31 ++++--- .../srt/speculative/dflash_worker_v2.py | 2 +- python/sglang/srt/speculative/draft_utils.py | 3 +- .../srt/speculative/draft_worker_common.py | 45 ++-------- .../dspark_components/dspark_worker_v2.py | 2 +- .../attention_methods/dense_attention.py | 5 ++ .../attention_methods/dsa_attention.py | 5 ++ .../attention_methods/dsv4_attention.py | 5 ++ .../attention_methods/dual_chunk_attention.py | 5 ++ .../attention_methods/gdn_attention.py | 5 ++ .../attention_methods/kda_attention.py | 5 ++ .../attention_methods/lightning_attention.py | 5 ++ .../attention_methods/mamba2_attention.py | 5 ++ .../attention_methods/mla_attention.py | 5 ++ .../test_fp4_kv_cache_quant_method.py | 1 + .../test_chunked_prefix_cache_gate.py | 16 ++-- .../unit/spec/test_draft_per_runner_config.py | 82 ++++++++++++++----- 22 files changed, 199 insertions(+), 103 deletions(-) diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 9af985d73..bdc15d51c 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -398,23 +398,16 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac allowed = {"triton", "trtllm_mha", "flashinfer"} else: allowed = {"triton", "trtllm_mha", "fa4"} - attn_be = runner.server_args.attention_backend - prefill_be = runner.server_args.prefill_attention_backend - decode_be = runner.server_args.decode_attention_backend - # When using split prefill/decode backends, check each individually - if prefill_be and decode_be: - assert prefill_be in allowed and decode_be in allowed, ( - f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. " - f"Got prefill={prefill_be}, decode={decode_be}." - ) - else: - assert attn_be in allowed, ( - f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. " - f"Got attention_backend={attn_be}." - ) + prefill_be = runner.prefill_attention_backend_str + decode_be = runner.decode_attention_backend_str + assert prefill_be in allowed and decode_be in allowed, ( + f"Only {allowed} backends are supported on Blackwell GPUs for hybrid GDN models. " + f"Got prefill={prefill_be}, decode={decode_be}." + ) elif is_npu(): assert ( - runner.server_args.attention_backend == "ascend" + runner.prefill_attention_backend_str == "ascend" + and runner.decode_attention_backend_str == "ascend" ), "ascend backend is the only supported backend on NPU for hybrid GDN models, use --attention-backend ascend to specify the backend." logger.info(f"Using hybrid linear attention backend for hybrid GDN models.") linear_attn_backend = GDNAttnBackend(runner) diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 02bc8e6de..b524a9afc 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -319,9 +319,8 @@ class FlashInferAttnBackend(AttentionBackend): self.decode_kv_access = self.kv_cache_quant_method.resolve_attention_access( "decode", "flashinfer" ) - prefill_backend, decode_backend = ( - model_runner.server_args.get_attention_backends() - ) + prefill_backend = model_runner.prefill_attention_backend_str + decode_backend = model_runner.decode_attention_backend_str if self.__class__ is FlashInferAttnBackend: if prefill_backend == "flashinfer": self._check_kv_attention_access("prefill", self.prefill_kv_access) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 0fdf5b966..f6eec4b37 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -310,6 +310,7 @@ class TpModelWorker(BaseTpWorker): memory_pool_config: Optional[MemoryPoolConfig] = None, is_multi_layer_eagle: bool = False, context_length: Optional[int] = None, + draft_attention_backend: Optional[str] = None, ): # Parse args self.server_args = server_args @@ -325,6 +326,8 @@ class TpModelWorker(BaseTpWorker): # Draft worker: target's effective context length; the draft runs at # absolute target positions. None keeps server_args.context_length. self.context_length = context_length + # Draft worker: the attention backend the algorithm resolved for it. + self.draft_attention_backend = draft_attention_backend # MTP model runners self.model_runner_list: List[ModelRunner] = [] @@ -459,6 +462,7 @@ class TpModelWorker(BaseTpWorker): req_to_token_pool=self.req_to_token_pool, token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, memory_pool_config=self.memory_pool_config, + draft_attention_backend=self.draft_attention_backend, draft_model_idx=0 if self.is_multi_layer_eagle else None, ) @@ -479,6 +483,7 @@ class TpModelWorker(BaseTpWorker): req_to_token_pool=self.req_to_token_pool, token_to_kv_pool_allocator=self.token_to_kv_pool_allocator, memory_pool_config=self.memory_pool_config, + draft_attention_backend=self.draft_attention_backend, draft_model_idx=i, ) ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 05c71d00d..14ebe2e83 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -924,7 +924,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ret.extend_prefix_lens = extend_prefix_lens ret.extend_num_tokens = batch.extend_num_tokens positions, ret.extend_start_loc = compute_position( - model_runner.server_args.attention_backend, + model_runner.prefill_attention_backend_str, ret.extend_prefix_lens, ret.extend_seq_lens, ret.extend_num_tokens, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index cc029ae5b..12bbd4fba 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -112,6 +112,7 @@ from sglang.srt.model_executor.model_runner_components.attention_backend_setup i build_attention_backends, configure_aux_hidden_state_capture, get_attention_backend, + resolve_attention_backend_strs, ) from sglang.srt.model_executor.model_runner_components.cuda_graph_setup import ( capture_cuda_graphs, @@ -261,6 +262,24 @@ class ModelRunnerOutput: indexer_topk_output: Optional[TopkCaptureOutput] = None +def resolve_draft_attention_backend( + *, + draft_attention_backend: Optional[str], + server_args: ServerArgs, + is_draft_worker: bool, +) -> Optional[str]: + """The attention backend a runner uses because it is a draft runner. + + ``None`` for a target runner. For a draft: the backend the algorithm that + built it resolved (the supported-backend fallback in + ``build_draft_tp_worker``), else ``--speculative-draft-attention-backend``. + It belongs to the runner, not the process: target and draft coexist. + """ + if not is_draft_worker: + return None + return draft_attention_backend or server_args.speculative_draft_attention_backend + + class ModelRunner: """ModelRunner runs the forward passes of the models.""" @@ -277,6 +296,7 @@ class ModelRunner: token_to_kv_pool_allocator: Optional[BaseTokenToKVPoolAllocator] = None, memory_pool_config: Optional[MemoryPoolConfig] = None, draft_model_idx: Optional[int] = None, + draft_attention_backend: Optional[str] = None, ): # Parse args self.mem_fraction_static = mem_fraction_static @@ -293,6 +313,11 @@ class ModelRunner: self.dist_port = nccl_port self.server_args = server_args self.is_draft_worker = is_draft_worker + self.draft_attention_backend = resolve_draft_attention_backend( + draft_attention_backend=draft_attention_backend, + server_args=server_args, + is_draft_worker=is_draft_worker, + ) # This runner's own load format, resolved before anything keys off it: # the remote-instance transfer engine is initialized at the top of # initialize(), long before the weights are loaded. @@ -902,12 +927,15 @@ class ModelRunner: dflash_target_layer_ids=self.spec_aux_config.dflash_target_layer_ids, is_dspark=self.spec_algorithm.is_dspark(), ) + # Resolve before building: backends read the pair off the runner while + # they construct (the FlashInfer KV-access check). + resolved = resolve_attention_backend_strs(model_runner=self) + self.prefill_attention_backend_str = resolved.prefill + self.decode_attention_backend_str = resolved.decode backends = build_attention_backends(model_runner=self) self.attn_backend = backends.attn_backend self.decode_attn_backend = backends.decode_attn_backend self.decode_attn_backend_group = backends.decode_attn_backend_group - self.prefill_attention_backend_str = backends.prefill_attention_backend_str - self.decode_attention_backend_str = backends.decode_attention_backend_str if self.server_args.dcp_size > 1 and get_parallel().dcp_replicate_q_proj: self._prepare_replicated_q_proj() @@ -1061,7 +1089,9 @@ class ModelRunner: ) maybe_trigger_remote_instance_nccl_send_group( - server_args=self.server_args, tp_rank=self.ps.tp_rank + server_args=self.server_args, + tp_rank=self.ps.tp_rank, + load_format=draft_load_format, ) with self._load_format_scope(draft_load_format): @@ -1248,9 +1278,7 @@ class ModelRunner: if spec_algorithm is not None else False ), - speculative_draft_attention_backend=getattr( - self.server_args, "speculative_draft_attention_backend", None - ), + speculative_draft_attention_backend=self.draft_attention_backend, ) ) # This runner's OWN resolved dtype string (target or draft). Attention diff --git a/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py b/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py index e9a863690..6d0a8af30 100644 --- a/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py +++ b/python/sglang/srt/model_executor/model_runner_components/attention_backend_setup.py @@ -17,7 +17,6 @@ from sglang.srt.utils import init_cublas if TYPE_CHECKING: from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.model_executor.model_runner import ModelRunner - from sglang.srt.server_args import ServerArgs logger = logging.getLogger(__name__) @@ -73,8 +72,13 @@ def build_attention_backends(*, model_runner: ModelRunner) -> AttentionBackends: if model_runner.device in ("cuda", "musa"): init_cublas() - resolved = _resolve_attention_backend_strs( - server_args=server_args, is_draft_worker=model_runner.is_draft_worker + # Already resolved and stamped on the runner before this call. + resolved = ResolvedAttentionBackendStr( + prefill=model_runner.prefill_attention_backend_str, + decode=model_runner.decode_attention_backend_str, + is_draft_override=bool( + model_runner.is_draft_worker and model_runner.draft_attention_backend + ), ) if server_args.enable_pdmux: @@ -140,10 +144,7 @@ def get_attention_backend( *, model_runner: ModelRunner, init_new_workspace: bool = False ) -> AttentionBackend: """Init attention kernel backend.""" - resolved = _resolve_attention_backend_strs( - server_args=model_runner.server_args, - is_draft_worker=model_runner.is_draft_worker, - ) + resolved = resolve_attention_backend_strs(model_runner=model_runner) return _build_resolved_backend( model_runner=model_runner, resolved=resolved, @@ -151,10 +152,18 @@ def get_attention_backend( ) -def _resolve_attention_backend_strs( - *, server_args: ServerArgs, is_draft_worker: bool +def resolve_attention_backend_strs( + *, model_runner: ModelRunner ) -> ResolvedAttentionBackendStr: - draft_attn_backend = server_args.speculative_draft_attention_backend + """The (prefill, decode) backends this runner runs. + + A draft runner's backend is its own (``ModelRunner.draft_attention_backend``): + target and draft coexist in one process, so it cannot come from the + process-wide config. + """ + server_args = model_runner.server_args + is_draft_worker = model_runner.is_draft_worker + draft_attn_backend = model_runner.draft_attention_backend if is_draft_worker and draft_attn_backend: logger.warning(f"Overriding draft attention backend to {draft_attn_backend}.") # Single backend for all draft modes (no prefill/decode split). @@ -218,7 +227,7 @@ def _build_resolved_backend( else: attn_backend = _build_backend_from_str( model_runner=model_runner, - backend_str=model_runner.server_args.attention_backend, + backend_str=resolved.prefill, init_new_workspace=init_new_workspace, ) return attn_backend diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 23fccf90d..4f3c100b8 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -1425,7 +1425,7 @@ class DFlashWorkerV2(BaseSpecWorker): "DFLASH prefill expected out_cache_loc, but got None." ) positions, _ = compute_position( - self.model_runner.server_args.attention_backend, + self.model_runner.prefill_attention_backend_str, draft_seq_lens, ctx_lens, int(sum(batch.extend_lens)), diff --git a/python/sglang/srt/speculative/draft_utils.py b/python/sglang/srt/speculative/draft_utils.py index c05a57709..bf3e92f29 100644 --- a/python/sglang/srt/speculative/draft_utils.py +++ b/python/sglang/srt/speculative/draft_utils.py @@ -39,7 +39,8 @@ class DraftBackendFactory: self.topk = topk self.speculative_num_steps = speculative_num_steps self.seed_dsa_topk_from_draft_extend = seed_dsa_topk_from_draft_extend - self.draft_attn_backend = server_args.speculative_draft_attention_backend + # The draft runner's own backend, not the process-wide config. + self.draft_attn_backend = draft_model_runner.draft_attention_backend def _create_backend( self, backend_name: str, backend_map: dict, error_template: str diff --git a/python/sglang/srt/speculative/draft_worker_common.py b/python/sglang/srt/speculative/draft_worker_common.py index 857178a6a..ca0be1060 100644 --- a/python/sglang/srt/speculative/draft_worker_common.py +++ b/python/sglang/srt/speculative/draft_worker_common.py @@ -9,7 +9,6 @@ import torch from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode -from sglang.srt.runtime_context import get_context, get_schedule from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 @@ -63,26 +62,6 @@ def _resolve_draft_attention_backend_fallback( return draft_backend -def draft_server_args_overrides(draft_backend) -> dict: - """The fields a draft variant must carry: its attention backend. - - Backend selection reads them off the config object the draft runner holds -- - ``speculative_draft_attention_backend`` in ``_resolve_attention_backend_strs`` - and ``configure_kv_cache_dtype``, ``attention_backend`` in the non-hybrid - branch of the backend build, and the split pair must not shadow either with - the target's. ``disable_chunked_prefix_cache`` is the target's resolved gate, - which lives in the bags only: publishing the variant re-projects the bags - from it, so the value has to travel on the variant. - """ - return dict( - speculative_draft_attention_backend=draft_backend, - prefill_attention_backend=None, - decode_attention_backend=None, - attention_backend=draft_backend, - disable_chunked_prefix_cache=get_schedule().disable_chunked_prefix_cache, - ) - - def build_draft_tp_worker( *, server_args: ServerArgs, @@ -101,23 +80,17 @@ def build_draft_tp_worker( server_args=server_args, algo_label=algo_label ) ) - draft_server_args = server_args.derive( - "draft_worker.build", **draft_server_args_overrides(draft_backend) + draft_worker = TpModelWorker( + server_args=server_args, + gpu_id=gpu_id, + ps=ps, + nccl_port=nccl_port, + is_draft_worker=True, + # The draft runs at absolute target positions. + context_length=target_model_config.context_len, + draft_attention_backend=draft_backend, ) - # The draft's layers must resolve config from the draft's own bags. - with get_context().preserve_config(): - get_context().set_server_args(draft_server_args) - draft_worker = TpModelWorker( - server_args=draft_server_args, - gpu_id=gpu_id, - ps=ps, - nccl_port=nccl_port, - is_draft_worker=True, - # The draft runs at absolute target positions. - context_length=target_model_config.context_len, - ) - draft_model_runner = draft_worker.model_runner draft_worker.draft_runner = draft_model_runner return DraftWorkerBundle( diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 0182c8082..ed64e21f3 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -444,7 +444,7 @@ class DSparkWorkerV2(BaseSpecWorker): batch.prefix_lens, dtype=torch.int32, device=device ) positions, _ = compute_position( - self.model_runner.server_args.attention_backend, + self.model_runner.prefill_attention_backend_str, draft_seq_lens, ctx_lens, int(sum(batch.extend_lens)), diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py index ddb5465d7..3ed889b2f 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py @@ -322,6 +322,11 @@ class MockModelRunner(ModelRunner): self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.canary_manager = None self.page_size = case.page_size diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index c097add8f..72a37f5e6 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -293,6 +293,11 @@ class DSAMockModelRunner(ModelRunner): # `set_mla_kv_buffer` does the quantize on the way in. self.kv_cache_dtype = torch.float8_e4m3fn if fp8_kv_cache else dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None # For TARGET_VERIFY / DRAFT_EXTEND, the DSA backend uses # `self.speculative_num_draft_tokens` to size `seqlens_expanded` # (`dsa_backend.py:482-486,510-515`). When zero, deep_gemm's diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py index 2ce09fa28..1e9550d77 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py @@ -332,6 +332,11 @@ class MockDSV4ModelRunner: self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.canary_manager = None self.page_size = case.page_size diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py index a6e44d24d..3eabff6ef 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py @@ -322,6 +322,11 @@ class DualChunkMockModelRunner(ModelRunner): self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.canary_manager = None self.page_size = case.page_size diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py index cabc12937..db551e0db 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py @@ -217,6 +217,11 @@ class MockGDNModelRunner(ModelRunner): self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.ps = ParallelState.trivial() self.canary_manager = None diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py index 219e67cb6..dd53c74cb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py @@ -222,6 +222,11 @@ class MockKDAModelRunner(ModelRunner): self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.ps = ParallelState.trivial() self.canary_manager = None diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py index 43e4cdf8b..a1fcf89ef 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py @@ -230,6 +230,11 @@ class MockLightningModelRunner(ModelRunner): self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.ps = ParallelState.trivial() self.canary_manager = None diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py index f23c3cd4e..78436ba14 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py @@ -315,6 +315,11 @@ class MockMamba2ModelRunner(ModelRunner): self.dtype = dtype self.kv_cache_dtype = dtype self.kv_cache_dtype_str = "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.ps = ParallelState.trivial() self.canary_manager = None diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py index e7debefa5..d2e2f04fb 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py @@ -233,6 +233,11 @@ class MockMLAModelRunner(ModelRunner): # does the BF16->FP8 cast on the way in. self.kv_cache_dtype = torch.float8_e4m3fn if fp8_kv_cache else dtype self.kv_cache_dtype_str = "fp8_e4m3" if fp8_kv_cache else "auto" + # This runner's own resolved backends (production stamps these in + # ModelRunner.initialize); a draft runner would carry its own. + self.prefill_attention_backend_str = case.backend + self.decode_attention_backend_str = case.backend + self.draft_attention_backend = None self.gpu_id = 0 self.canary_manager = None self.page_size = case.page_size diff --git a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py index 316456a5e..118de4a62 100644 --- a/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py +++ b/test/registered/unit/layers/quantization/test_fp4_kv_cache_quant_method.py @@ -75,6 +75,7 @@ class TestKVCacheQuantRegistry(CustomTestCase): runner = object.__new__(ModelRunner) runner.server_args = SimpleNamespace(kv_cache_dtype="fp4_e2m1") + runner.draft_attention_backend = None with self.assertRaisesRegex(ValueError, "fp4_mx_block16"): runner.configure_kv_cache_dtype() diff --git a/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py b/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py index f20223985..8b0c7f9f0 100644 --- a/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py +++ b/test/registered/unit/model_executor/test_chunked_prefix_cache_gate.py @@ -51,17 +51,15 @@ class TestChunkedPrefixCacheGate(CustomTestCase): get_context().set_server_args(sa) # what a later republish would do self.assertFalse(get_schedule().disable_chunked_prefix_cache) - def test_draft_variant_fields_carry_the_gate(self): - # Publishing the draft variant re-projects the bags from it, so the - # gate — which lives in the bags only — has to travel on the variant. - from sglang.srt.speculative.draft_worker_common import ( - draft_server_args_overrides, - ) - + def test_the_gate_survives_a_draft_build(self): + # The draft build no longer publishes a config of its own, so the gate + # the target resolved stays in the bags for the rest of the process. self._seed(attention_backend="triton") maybe_disable_chunked_prefix_cache(use_mla_backend=True, is_draft_worker=False) - fields = draft_server_args_overrides(draft_backend="fa3") - self.assertTrue(fields["disable_chunked_prefix_cache"]) + self.assertTrue(get_schedule().disable_chunked_prefix_cache) + + maybe_disable_chunked_prefix_cache(use_mla_backend=False, is_draft_worker=True) + self.assertTrue(get_schedule().disable_chunked_prefix_cache) if __name__ == "__main__": diff --git a/test/registered/unit/spec/test_draft_per_runner_config.py b/test/registered/unit/spec/test_draft_per_runner_config.py index 325ab1d79..274a71a51 100644 --- a/test/registered/unit/spec/test_draft_per_runner_config.py +++ b/test/registered/unit/spec/test_draft_per_runner_config.py @@ -12,12 +12,17 @@ import unittest from types import SimpleNamespace from sglang.srt.managers.scheduler import Scheduler -from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.model_executor.model_runner import ( + ModelRunner, + resolve_draft_attention_backend, +) +from sglang.srt.model_executor.model_runner_components.attention_backend_setup import ( + resolve_attention_backend_strs, +) from sglang.srt.model_executor.model_runner_components.load_model_utils import ( build_load_config, ) from sglang.srt.runtime_context import get_context, get_model -from sglang.srt.speculative.draft_worker_common import draft_server_args_overrides from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -102,27 +107,66 @@ class TestDraftPerRunnerConfig(CustomTestCase): build_load_config(load_format="dummy", **common).load_format, "dummy" ) - # -- the variant left in the dflash / dspark path carries backends only ---- + # -- the attention backend is per-runner, not a config variant ------------- - def test_the_draft_variant_carries_the_backend_family_only(self): - self._seed(disable_chunked_prefix_cache=False) - fields = draft_server_args_overrides("triton") + def _runner(self, *, is_draft_worker, draft_attention_backend=None): + runner = ModelRunner.__new__(ModelRunner) + runner.server_args = get_context().server_args + runner.is_draft_worker = is_draft_worker + runner.draft_attention_backend = draft_attention_backend + return runner - self.assertEqual(fields["attention_backend"], "triton") - self.assertEqual(fields["speculative_draft_attention_backend"], "triton") - self.assertIsNone(fields["prefill_attention_backend"]) - self.assertIsNone(fields["decode_attention_backend"]) - self.assertNotIn("context_length", fields) - self.assertNotIn("load_format", fields) - self.assertNotIn("skip_tokenizer_init", fields) + def test_the_draft_backend_applies_to_the_draft_runner_only(self): + self._seed(attention_backend="fa3") - def test_the_variant_carries_the_targets_resolved_gate(self): - """Publishing the variant re-projects the bags, so the gate travels.""" - self._seed(disable_chunked_prefix_cache=False) - get_context().override("test.gate", disable_chunked_prefix_cache=True) - self.assertTrue( - draft_server_args_overrides("triton")["disable_chunked_prefix_cache"] + draft = resolve_attention_backend_strs( + model_runner=self._runner( + is_draft_worker=True, draft_attention_backend="triton" + ) ) + self.assertEqual((draft.prefill, draft.decode), ("triton", "triton")) + self.assertTrue(draft.is_draft_override) + + target = resolve_attention_backend_strs( + model_runner=self._runner(is_draft_worker=False) + ) + self.assertEqual((target.prefill, target.decode), ("fa3", "fa3")) + + def test_an_unresolved_draft_falls_back_to_the_config_field(self): + """The v2 workers pass no backend: --speculative-draft-attention-backend.""" + server_args = self._seed( + attention_backend="fa3", speculative_draft_attention_backend="triton" + ) + + def effective(*, is_draft_worker, passed=None): + return resolve_draft_attention_backend( + draft_attention_backend=passed, + server_args=server_args, + is_draft_worker=is_draft_worker, + ) + + self.assertEqual(effective(is_draft_worker=True), "triton") + self.assertEqual(effective(is_draft_worker=True, passed="fa3"), "fa3") + self.assertIsNone(effective(is_draft_worker=False)) + + draft = resolve_attention_backend_strs( + model_runner=self._runner( + is_draft_worker=True, + draft_attention_backend=effective(is_draft_worker=True), + ) + ) + self.assertEqual((draft.prefill, draft.decode), ("triton", "triton")) + + def test_the_target_keeps_its_split_pair(self): + self._seed( + attention_backend="fa3", + prefill_attention_backend="flashinfer", + decode_attention_backend="fa3", + ) + target = resolve_attention_backend_strs( + model_runner=self._runner(is_draft_worker=False) + ) + self.assertEqual((target.prefill, target.decode), ("flashinfer", "fa3")) # -- the scheduler hands over the process's own config ---------------------