From ab810e40524c8489a9eb9c3cbcdb7ade1fa6b947 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Sat, 15 Aug 2026 00:36:04 -0700 Subject: [PATCH] config: each runner carries its own linear-attn kernel choice Two problems in the same family as #33312 (a per-runner decision that one participant answered differently), one fixed here and one guarded. **The linear-attn kernel backends were process-wide.** `attn_backend_wrapper` rebuilt a module-level dict once per runner, from the handed record plus a local `prefill_default`. Two things follow, and both are wrong: - **A draft could not hold a different choice than its target.** Only the runner whose model is GDN gets the SM100 FlashInfer prefill default; the operator's explicit flag belongs to the launch. The full-attention backends already model this correctly -- the runner stamps `prefill_attention_backend_str` / `decode_attention_backend_str` and its backend objects are built from the stamp. Linear attn had no stamp at all. - **The second rebuild replaced the first one's choice.** The default was also recorded into the process-wide config, which the record does not see, so a runner rebuilding without a default of its own resolved `prefill` back to the base backend -- silently swapping the kernel the earlier runner selected. Demonstrated in-process before this change: table `FLASHINFER`, then `TRITON`. `resolve_linear_attn_backends(prefill_default=None)` returns a frozen `LinearAttnBackends(decode, prefill, verify)` from the published `exec.mamba` leaves; the wrapper stamps it as `runner.linear_attn_backends` before building the backends that read it; and the three consumers (GDN, KDA, Ascend GDN) read it off the runner they are built for. Each already took `model_runner` and cached the result on itself, so the value now simply comes from the right place. A backend built outside that path has no stamp and raises on the attribute, the way the full-attention strings do -- no silent fallback to hide the wiring mistake. The recording goes away with it. A per-runner choice in the process-wide config has no meaning the second runner can read correctly: the leaf is how the gate asks "did the operator name a backend", so a recorded default reads back as an operator flag and the next runner declines its own. The leaf now keeps meaning what was asked for at launch, and the effective choice lives in the stamp (and in the log line the gate already emits). Precedence is unchanged: the resolver takes the default as an argument and an explicit `--linear-attn-prefill-backend` wins over it, with the gate declining early so it neither probes the device nor logs. **A draft entry class must answer the loader exactly when its target does.** The loader asks the entry class it instantiates for the shared-experts-fusion decision, and a draft is its own entry class. When the target family carries auto-disable conditions and the draft's class does not expose them, the loader installs one decision for each and the draft's weights are laid out for the wrong one. That shipped: the DSV4 DSpark draft skipped its bundled shared-expert tensors until #33312 gave it the gate, costing accept length 5.60 -> 2.05. `test_fusion_gate_coverage.py` walks the same registry but asks whether an entry class *touches* the decision -- reads the flag, names a gated class. That catches a class once it already consumes the decision; it could not catch one that should consume it and does not, which is what the DSpark class looked like (it built the family's *layer* classes, so the flag reader lived in another module and its own source named no gated class). `test_draft_entry_hook_parity.py` asks the invariant directly: presence parity between a draft entry class and the target it is named after. Identity is deliberately not required -- the Qwen3.5 MTP delegates with adapted arguments (unwrapping `text_config`, using the MTP quantization config), which is right -- and weight-name maps are out of scope, since a draft's checkpoint has its own names. Reverse-verified against the original defect: with the DSpark gate removed the case names the pair and the side that is missing it; with #33312 in place it passes. --- .../npu/attention/ascend_gdn_backend.py | 9 +- .../layers/attention/attention_registry.py | 13 +- .../layers/attention/linear/gdn_backend.py | 9 +- .../layers/attention/linear/kda_backend.py | 10 +- .../srt/layers/attention/linear/utils.py | 88 ++++---- .../attention_methods/gdn_attention.py | 6 +- .../attention_methods/kda_attention.py | 6 +- .../attention_methods/lightning_attention.py | 6 +- .../attention/test_linear_attn_config.py | 190 +++++++++++++----- .../models/test_draft_entry_hook_parity.py | 95 +++++++++ 10 files changed, 295 insertions(+), 137 deletions(-) create mode 100644 test/registered/unit/models/test_draft_entry_hook_parity.py diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py index cc9510ff9..892066a80 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_gdn_backend.py @@ -10,10 +10,6 @@ from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend AscendMambaAttnBackendBase, ) from sglang.srt.layers.attention.linear.gdn_backend import GDNKernelDispatcher -from sglang.srt.layers.attention.linear.utils import ( - get_linear_attn_decode_backend, - get_linear_attn_prefill_backend, -) from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import MambaPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -36,8 +32,9 @@ class AscendGDNAttnBackend(AscendMambaAttnBackendBase): model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape[-2], ) ) - decode_backend = get_linear_attn_decode_backend() - prefill_backend = get_linear_attn_prefill_backend() + backends = model_runner.linear_attn_backends + decode_backend = backends.decode + prefill_backend = backends.prefill self.kernel_dispatcher = GDNKernelDispatcher(decode_backend, prefill_backend) def _prepare_mamba_track_metadata(self, forward_batch: ForwardBatch): diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 84237b424..1d176ce7b 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -13,7 +13,7 @@ from sglang.srt.configs.linear_attn_model_registry import ( get_linear_attn_config, import_backend_class, ) -from sglang.srt.runtime_context import get_context, get_parallel +from sglang.srt.runtime_context import get_parallel from sglang.srt.utils import get_device_capability, is_hip, is_musa, is_npu _is_musa = is_musa() @@ -351,7 +351,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac LightningAttentionBackend, ) from sglang.srt.layers.attention.linear.utils import ( - initialize_linear_attn_config, + resolve_linear_attn_backends, ) from sglang.srt.utils import ( is_blackwell, @@ -383,13 +383,8 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac prefill_default = None if hybrid_gdn_config(runner.model_config) is not None and not is_npu(): prefill_default = flashinfer_gdn_prefill_default(runner) - if prefill_default is not None: - get_context().override( - "gdn_backend.sm100_flashinfer_default", - linear_attn_prefill_backend=prefill_default, - ) - initialize_linear_attn_config( - runner.server_args, prefill_default=prefill_default + runner.linear_attn_backends = resolve_linear_attn_backends( + prefill_default=prefill_default ) hybrid_backend_cls = HybridLinearAttnBackend if hybrid_gdn_config(runner.model_config) is not None: diff --git a/python/sglang/srt/layers/attention/linear/gdn_backend.py b/python/sglang/srt/layers/attention/linear/gdn_backend.py index 89f014703..18634c283 100644 --- a/python/sglang/srt/layers/attention/linear/gdn_backend.py +++ b/python/sglang/srt/layers/attention/linear/gdn_backend.py @@ -13,9 +13,6 @@ from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKerne from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, build_verify_intermediate_state_indices, - get_linear_attn_decode_backend, - get_linear_attn_prefill_backend, - get_linear_attn_verify_backend, ) from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import MambaPool @@ -359,11 +356,9 @@ class GDNAttnBackend(MambaAttnBackendBase): self.conv_states_shape[-1] < FLA_CHUNK_SIZE ), f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}" - decode_backend = get_linear_attn_decode_backend() - prefill_backend = get_linear_attn_prefill_backend() - verify_backend = get_linear_attn_verify_backend() + backends = model_runner.linear_attn_backends self.kernel_dispatcher = GDNKernelDispatcher( - decode_backend, prefill_backend, verify_backend + backends.decode, backends.prefill, backends.verify ) # Sized past the pool for attn_tp-padded warmup/MLP-sync batches (see helper). self.verify_intermediate_state_indices = ( diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index e09984bec..68008fb06 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -13,9 +13,6 @@ from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKerne from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, build_verify_intermediate_state_indices, - get_linear_attn_decode_backend, - get_linear_attn_prefill_backend, - get_linear_attn_verify_backend, ) from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.utils import is_cpu, is_cuda, is_npu @@ -379,9 +376,10 @@ class KDAAttnBackend(MambaAttnBackendBase): .transpose(-1, -2) .shape ) - decode_backend = get_linear_attn_decode_backend() - prefill_backend = get_linear_attn_prefill_backend() - verify_backend = get_linear_attn_verify_backend() + backends = model_runner.linear_attn_backends + decode_backend = backends.decode + prefill_backend = backends.prefill + verify_backend = backends.verify # KDA FlashInfer target_verify (recurrent_kda) is chain-only (no tree-ancestor # traversal). Reject EAGLE tree verify (topk > 1) early at setup, keyed on the # verify backend (not decode). The kernel keeps a per-call diff --git a/python/sglang/srt/layers/attention/linear/utils.py b/python/sglang/srt/layers/attention/linear/utils.py index 681f589c6..74323c216 100644 --- a/python/sglang/srt/layers/attention/linear/utils.py +++ b/python/sglang/srt/layers/attention/linear/utils.py @@ -1,16 +1,16 @@ from __future__ import annotations -import logging from enum import Enum -from typing import TYPE_CHECKING, Dict, Optional +from typing import TYPE_CHECKING, Optional +import msgspec + +from sglang.srt.runtime_context import get_exec from sglang.srt.utils.common import rank0_log if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs -logger = logging.getLogger(__name__) - class LinearAttnKernelBackend(Enum): TRITON = "triton" @@ -55,56 +55,48 @@ class LinearAttnKernelBackend(Enum): return self == LinearAttnKernelBackend.CUSTOM -_BACKENDS: Dict[str, Optional[LinearAttnKernelBackend]] = { - "decode": None, - "prefill": None, - "verify": None, -} +class LinearAttnBackends(msgspec.Struct, frozen=True): + """One runner's linear-attn kernel choice, per phase. + + Per runner, not per process: a target and its draft coexist and can want + different kernels (only the runner whose model is GDN gets the SM100 + FlashInfer prefill default, and an explicit flag applies to whichever runner + was launched with it). + """ + + decode: LinearAttnKernelBackend + prefill: LinearAttnKernelBackend + verify: LinearAttnKernelBackend -def initialize_linear_attn_config( - server_args: ServerArgs, prefill_default: Optional[str] = None -): - base = server_args.linear_attn_backend - decode = server_args.linear_attn_decode_backend or base - prefill = server_args.linear_attn_prefill_backend or prefill_default or base +def resolve_linear_attn_backends( + prefill_default: Optional[str] = None, +) -> LinearAttnBackends: + """This runner's kernel choice from the published leaves. - _BACKENDS["decode"] = LinearAttnKernelBackend(decode) - _BACKENDS["prefill"] = LinearAttnKernelBackend(prefill) - - # Verify backend. Unset -> follow decode (flashinfer -> its recurrent kernel, - # else triton), preserving historical behavior. - verify = server_args.linear_attn_verify_backend - if verify is None: - verify = decode if _BACKENDS["decode"].is_flashinfer() else "triton" - _BACKENDS["verify"] = LinearAttnKernelBackend(verify) - - rank0_log( - f"Linear attention kernel backend: decode={decode}, prefill={prefill}, " - f"verify={verify}" + ``prefill_default`` is the caller's own auto-default (the SM100 GDN + domain); an explicitly configured ``--linear-attn-prefill-backend`` wins. + """ + mamba = get_exec().mamba + base = mamba.linear_attn_backend + decode = LinearAttnKernelBackend(mamba.linear_attn_decode_backend or base) + prefill = LinearAttnKernelBackend( + mamba.linear_attn_prefill_backend or prefill_default or base ) + # Unset verify follows decode (flashinfer -> its recurrent kernel, else triton). + verify = mamba.linear_attn_verify_backend + if verify is None: + verify = decode.value if decode.is_flashinfer() else "triton" -def _get_backend(phase: str) -> LinearAttnKernelBackend: - backend = _BACKENDS[phase] - if backend is None: - logger.warning( - "linear-attn %s backend is not initialized, using triton backend", phase - ) - backend = _BACKENDS[phase] = LinearAttnKernelBackend.TRITON - return backend - - -def get_linear_attn_decode_backend() -> LinearAttnKernelBackend: - return _get_backend("decode") - - -def get_linear_attn_prefill_backend() -> LinearAttnKernelBackend: - return _get_backend("prefill") - - -def get_linear_attn_verify_backend() -> LinearAttnKernelBackend: - return _get_backend("verify") + backends = LinearAttnBackends( + decode=decode, prefill=prefill, verify=LinearAttnKernelBackend(verify) + ) + rank0_log( + f"Linear attention kernel backend: decode={backends.decode.value}, " + f"prefill={backends.prefill.value}, verify={backends.verify.value}" + ) + return backends def build_verify_intermediate_state_indices( 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 50354ff4a..9ef502e94 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 @@ -16,7 +16,7 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( HybridLinearAttnBackend, ) from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend -from sglang.srt.layers.attention.linear.utils import initialize_linear_attn_config +from sglang.srt.layers.attention.linear.utils import resolve_linear_attn_backends from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, @@ -605,7 +605,9 @@ def build_gdn_attention_fixture( except (AssertionError, ImportError, ModuleNotFoundError) as exc: testcase.skipTest(f"{case.backend} backend is not available: {exc}") - initialize_linear_attn_config(runner.server_args) + # Standing in for `attn_backend_wrapper`, which is what stamps this on a + # runner before building the backend that reads it. + runner.linear_attn_backends = resolve_linear_attn_backends() linear_backend = GDNAttnBackend(runner) if case.linear_attn_prefill_backend == "flashinfer": from sglang.srt.layers.attention.linear.kernels.gdn_flashinfer import ( 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 0315c5352..899b2dbce 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 @@ -16,7 +16,7 @@ from sglang.srt.layers.attention.hybrid_linear_attn_backend import ( HybridLinearAttnBackend, ) from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend -from sglang.srt.layers.attention.linear.utils import initialize_linear_attn_config +from sglang.srt.layers.attention.linear.utils import resolve_linear_attn_backends from sglang.srt.layers.radix_linear_attention import RadixLinearAttention from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, @@ -609,7 +609,9 @@ def build_kda_attention_fixture( except (AssertionError, ImportError, ModuleNotFoundError) as exc: testcase.skipTest(f"{case.backend} backend is not available: {exc}") - initialize_linear_attn_config(runner.server_args) + # Standing in for `attn_backend_wrapper`, which is what stamps this on a + # runner before building the backend that reads it. + runner.linear_attn_backends = resolve_linear_attn_backends() linear_backend = KDAAttnBackend(runner) backend = HybridLinearAttnBackend(full_backend, linear_backend, full_attn_layers=[]) actual_module = ProjectedKDAAttention( 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 1b85b55b3..89b104f67 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 @@ -15,7 +15,7 @@ from sglang.srt.layers.attention.attention_registry import ATTENTION_BACKENDS from sglang.srt.layers.attention.linear.lightning_backend import ( LightningAttentionBackend, ) -from sglang.srt.layers.attention.linear.utils import initialize_linear_attn_config +from sglang.srt.layers.attention.linear.utils import resolve_linear_attn_backends from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.mem_cache.memory_pool import ( HybridReqToTokenPool, @@ -581,7 +581,9 @@ def build_lightning_attention_fixture( except (AssertionError, ImportError, ModuleNotFoundError) as exc: testcase.skipTest(f"{case.backend} backend is not available: {exc}") - initialize_linear_attn_config(runner.server_args) + # Standing in for `attn_backend_wrapper`, which is what stamps this on a + # runner before building the backend that reads it. + runner.linear_attn_backends = resolve_linear_attn_backends() backend = LightningAttentionBackend(runner) actual_module = ProjectedLightningAttention( num_heads=case.num_heads, diff --git a/test/registered/unit/layers/attention/test_linear_attn_config.py b/test/registered/unit/layers/attention/test_linear_attn_config.py index d0af273b4..a69371c15 100644 --- a/test/registered/unit/layers/attention/test_linear_attn_config.py +++ b/test/registered/unit/layers/attention/test_linear_attn_config.py @@ -1,80 +1,160 @@ -"""Backend selection in initialize_linear_attn_config. +"""Each runner carries its own linear-attn kernel choice. -The SM100 GDN default reaches the module state as an argument rather than as a -ServerArgs mutation, so the precedence between an explicit flag, that default, -and the shared base backend is pinned here. +A target and its draft coexist in one process and can want different kernels: +only the runner whose model is GDN gets the SM100 FlashInfer prefill default, +and the operator's explicit flag applies to the launch. So the choice is a +per-runner stamp read from the runner, the way the full-attention pair already +works (`prefill_attention_backend_str` / `decode_attention_backend_str`), rather +than one process-wide table. + +It used to be that table: `attn_backend_wrapper` rebuilt a module-level dict +once per runner, from the handed record plus a local default. Two consequences, +both pinned below -- a second runner could not hold a different choice, and its +rebuild replaced the first runner's (the record never carries the recorded +default, so a runner with no default of its own resolved back to the base +backend). """ import unittest +from types import SimpleNamespace -from sglang.srt.layers.attention.linear import utils as linear_utils from sglang.srt.layers.attention.linear.utils import ( LinearAttnKernelBackend, - initialize_linear_attn_config, + resolve_linear_attn_backends, ) -from sglang.srt.server_args import ServerArgs 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") -class TestLinearAttnConfig(CustomTestCase): - def setUp(self): - self.addCleanup(linear_utils._BACKENDS.update, linear_utils._BACKENDS.copy()) +class _Runner: + """The only thing the readers need from a runner: the stamp.""" - def _init(self, prefill_default=None, **fields): - args = ServerArgs(model_path="dummy") - for key, value in fields.items(): - setattr(args, key, value) - initialize_linear_attn_config(args, prefill_default) - return ( - linear_utils.get_linear_attn_prefill_backend(), - linear_utils.get_linear_attn_decode_backend(), - ) - def test_default_applies_when_the_flag_is_unset(self): - prefill, _ = self._init( - prefill_default="flashinfer", linear_attn_backend="triton" - ) - self.assertEqual(prefill, LinearAttnKernelBackend.FLASHINFER) +class TestLinearAttnBackends(CustomTestCase): + def _publish(self, **fields): + from sglang.srt.runtime_context import get_context - def test_explicit_flag_wins_over_the_default(self): - prefill, _ = self._init( - prefill_default="flashinfer", - linear_attn_backend="triton", - linear_attn_prefill_backend="cutedsl", - ) - self.assertEqual(prefill, LinearAttnKernelBackend.CUTEDSL) - - def test_base_backend_applies_without_a_default(self): - prefill, decode = self._init(linear_attn_backend="triton") - self.assertEqual(prefill, LinearAttnKernelBackend.TRITON) - self.assertEqual(decode, LinearAttnKernelBackend.TRITON) - - def test_a_recorded_default_shows_in_the_resolved_config(self): - from sglang.srt.runtime_context import get_context, get_exec - - override = get_context().override_server_args(linear_attn_backend="triton") - server_args = override.install() + override = get_context().override_server_args(**fields) + override.install() self.addCleanup(override.restore) - get_context().override( - "gdn_backend.sm100_flashinfer_default", - linear_attn_prefill_backend="flashinfer", - ) - self.assertEqual(get_exec().mamba.linear_attn_prefill_backend, "flashinfer") - self.assertEqual( - get_context().resolved_server_args_dict()["linear_attn_prefill_backend"], - "flashinfer", - ) - self.assertIsNone(server_args.linear_attn_prefill_backend) + def test_the_default_applies_when_the_flag_is_unset(self): + self._publish(linear_attn_backend="triton") + backends = resolve_linear_attn_backends(prefill_default="flashinfer") + self.assertEqual(backends.prefill, LinearAttnKernelBackend.FLASHINFER) + + def test_the_base_backend_applies_without_a_default(self): + self._publish(linear_attn_backend="triton") + backends = resolve_linear_attn_backends() + self.assertEqual(backends.prefill, LinearAttnKernelBackend.TRITON) + self.assertEqual(backends.decode, LinearAttnKernelBackend.TRITON) def test_the_default_does_not_reach_the_decode_backend(self): - _, decode = self._init( - prefill_default="flashinfer", linear_attn_backend="triton" + self._publish(linear_attn_backend="triton") + backends = resolve_linear_attn_backends(prefill_default="flashinfer") + self.assertEqual(backends.decode, LinearAttnKernelBackend.TRITON) + + def test_explicit_flag_wins_over_the_default(self): + """Precedence lives in the gate, which declines once the flag is set. + + The resolver takes the default as an argument, so the flag has to win + upstream: `flashinfer_gdn_prefill_default` returns None the moment the + leaf is set, and that condition is checked before anything touches the + device -- which is what lets this run anywhere. + """ + from types import SimpleNamespace + + from sglang.srt.layers.attention.linear.gdn_backend import ( + flashinfer_gdn_prefill_default, + ) + from sglang.srt.runtime_context import get_server_args + + self._publish( + linear_attn_backend="triton", linear_attn_prefill_backend="cutedsl" + ) + runner = SimpleNamespace(server_args=get_server_args()) + self.assertIsNone(flashinfer_gdn_prefill_default(runner)) + self.assertEqual( + resolve_linear_attn_backends().prefill, LinearAttnKernelBackend.CUTEDSL + ) + + def test_two_runners_hold_different_choices(self): + """The property the process-wide table could not express. + + The GDN target gets the SM100 default; the draft that is not GDN has no + default of its own. Both stamps stand, and reading one does not disturb + the other -- under the old table the draft's rebuild replaced the + target's choice with the base backend. + """ + self._publish(linear_attn_backend="triton") + + target, draft = _Runner(), _Runner() + target.linear_attn_backends = resolve_linear_attn_backends( + prefill_default="flashinfer" + ) + draft.linear_attn_backends = resolve_linear_attn_backends() + + self.assertEqual( + target.linear_attn_backends.prefill, LinearAttnKernelBackend.FLASHINFER + ) + self.assertEqual( + draft.linear_attn_backends.prefill, LinearAttnKernelBackend.TRITON + ) + + def test_an_unstamped_runner_raises_rather_than_guessing(self): + """No default for "nobody stamped this", on the production read path. + + `attn_backend_wrapper` stamps before it builds the backends that read + the stamp, so a missing one means a backend was built outside that + path. The runner double below satisfies everything `GDNAttnBackend` + touches *before* the stamp read, so the `AttributeError` this asserts + comes from `model_runner.linear_attn_backends` itself -- a default + stamp on the runner or a restored module-level fallback would turn + this red-to-green, which is the regression it guards. Silent triton + fallback would hide the wiring mistake behind a working-but-wrong + kernel. + """ + import torch + + from sglang.srt.layers.attention.linear.gdn_backend import GDNAttnBackend + + runner = SimpleNamespace( + device="cpu", + server_args=SimpleNamespace( + speculative_eagle_topk=0, enable_unified_memory=False + ), + is_draft_worker=False, + req_to_token_pool=SimpleNamespace( + mamba_pool=SimpleNamespace( + mamba_cache=SimpleNamespace(conv=[torch.zeros(1, 1, 4, 7)]) + ) + ), + token_to_kv_pool=None, + ) + with self.assertRaises(AttributeError) as caught: + GDNAttnBackend(runner) + self.assertIn("linear_attn_backends", str(caught.exception)) + + def test_the_per_runner_default_stays_out_of_the_process_config(self): + """The process-wide config cannot represent one runner's choice. + + Recording the auto-default there is how a second runner used to inherit + it: the leaf then reads as "the operator named a backend", which is the + one question the gate asks. So the default lives in the runner's stamp + and the leaf keeps meaning what was asked for at launch. + """ + from sglang.srt.runtime_context import get_context, get_exec + + self._publish(linear_attn_backend="triton") + backends = resolve_linear_attn_backends(prefill_default="flashinfer") + + self.assertEqual(backends.prefill, LinearAttnKernelBackend.FLASHINFER) + self.assertIsNone(get_exec().mamba.linear_attn_prefill_backend) + self.assertIsNone( + get_context().resolved_server_args_dict()["linear_attn_prefill_backend"] ) - self.assertEqual(decode, LinearAttnKernelBackend.TRITON) if __name__ == "__main__": diff --git a/test/registered/unit/models/test_draft_entry_hook_parity.py b/test/registered/unit/models/test_draft_entry_hook_parity.py new file mode 100644 index 000000000..b96ca4308 --- /dev/null +++ b/test/registered/unit/models/test_draft_entry_hook_parity.py @@ -0,0 +1,95 @@ +"""A draft entry class answers the loader's questions exactly when its target does. + +The loader asks the *entry class* it instantiates for the shared-experts-fusion +decision (`install_shared_experts_fusion_decision`). A draft is its own entry +class -- `...NextN`, `...MTP`, `...DSpark`, `...Eagle3` -- so when the target +family carries auto-disable conditions and the draft's class does not expose +them, the draft resolves a *different* decision than the target it drafts for, +and its weights are laid out for the other layout. That shipped once: the DSV4 +DSpark draft skipped its bundled shared-expert tensors until #33312 gave it the +gate. + +`test_fusion_gate_coverage.py` walks the same registry but asks a different +question: does this entry class *touch* the decision (read the flag, name a gated +class)? That catches a class after it starts consuming the decision. It cannot +catch one that should consume it and does not, which is what the DSpark class +looked like: it built the family's *layer* classes, so the flag reader lived in +another module and its own source named no gated class. + +This case asks the invariant directly instead: **presence parity between a draft +entry class and its target**. Identity is deliberately not required -- a draft +that delegates with adapted arguments (the Qwen3.5 MTP unwraps `text_config` and +substitutes the MTP quantization config) is exactly right. +""" + +import re +import unittest + +from sglang.srt.models.registry import ModelRegistry +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=60, suite="base-a-test-cpu") + +# The hooks the loader resolves on the entry class where a draft/target +# disagreement is a defect rather than a difference. Weight-name maps +# (`packed_modules_mapping`, `hf_to_sglang_mapper`) are deliberately absent: a +# draft's checkpoint has its own names, so those differ by design. +PARITY_HOOKS = ("shared_experts_fusion_disable_reason",) + +# `...Eagle`/`...Eagle3` drafts are standalone checkpoints with their own +# architecture; the rest mirror a stage of their target. +DRAFT_SUFFIX = re.compile(r"(NextN|MTP|DSpark|DFlash|Standalone)$") + + +def _target_of(arch: str, archs: dict): + """The arch a draft entry class drafts for, if it is named after it.""" + match = DRAFT_SUFFIX.search(arch) + if not match: + return None + base = arch[: match.start()] + for candidate in (base, f"{base}ForCausalLM"): + if candidate in archs and candidate != arch: + return candidate + return None + + +class TestDraftEntryHookParity(CustomTestCase): + def test_a_draft_resolves_a_gate_exactly_when_its_target_does(self): + archs = {} + for arch in sorted(ModelRegistry.get_supported_archs()): + try: + archs[arch] = ModelRegistry.resolve_model_cls(arch)[0] + except Exception: + continue # an arch this environment cannot import + + checked, offenders = 0, [] + for arch, cls in archs.items(): + target = _target_of(arch, archs) + if target is None: + continue + for hook in PARITY_HOOKS: + checked += 1 + draft_has = hasattr(cls, hook) + target_has = hasattr(archs[target], hook) + if draft_has == target_has: + continue + which = "the draft" if target_has else "the target" + offenders.append( + f"{arch} vs {target}: {hook} missing on {which} " + f"({cls.__module__} / {archs[target].__module__})" + ) + + self.assertGreater(checked, 10, "the draft/target pairing found nothing") + self.assertEqual( + [], + offenders, + "a draft entry class and the target it drafts for must resolve the " + "same loader hooks; otherwise the loader installs one decision for " + "the target and another for the draft, and the draft's weights are " + "laid out for the wrong one:\n " + "\n ".join(offenders), + ) + + +if __name__ == "__main__": + unittest.main()