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()