[Spec] DFlash2: local convolution + candidate selector (#35371)

Co-authored-by: Jian Chen <jianchen0311@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com> Co-authored-by: hnyls2002 <lsyincs@gmail.com>
2026-08-18 17:07:28 -07:00
parent 3a8f522f65
commit c14312a664
7 changed files with 929 additions and 61 deletions
@@ -0,0 +1,38 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+
+from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import (
+    _map_muse_target_layer_ids,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+@pytest.mark.parametrize(
+    ("target_model_type", "draft_architecture", "expected"),
+    [
+        ("muse_glimmer", "MuseGlimmerAssistantModel", [2, 14, 26, 38, 50]),
+        ("muse_glimmer", "DFlash2DraftModel", [2, 14, 26, 38, 50]),
+        ("muse_glimmer", "DFlashDraftModel", [1, 13, 25, 37, 49]),
+        ("qwen3", "DFlash2DraftModel", [1, 13, 25, 37, 49]),
+        ("qwen3", "MuseGlimmerAssistantModel", [1, 13, 25, 37, 49]),
+    ],
+)
+def test_muse_target_layer_id_mapping(target_model_type, draft_architecture, expected):
+    """The +1 belongs to Muse targets, which report layer outputs where the rest
+    report layer inputs. The draft architecture alone does not earn it."""
+    assert (
+        _map_muse_target_layer_ids(
+            target_hf_config=SimpleNamespace(model_type=target_model_type),
+            draft_hf_config=SimpleNamespace(architectures=[draft_architecture]),
+            layer_ids=[1, 13, 25, 37, 49],
+        )
+        == expected
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))