Fix multi_layer_eagle_worker_v2 draft extend selection, add chain style multi layer mtp test (#22340)

Co-authored-by: 0xNullPath <luyan@nvidia.com>

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -27,6 +27,7 @@ from sglang.srt.managers.tp_worker import TpModelWorker
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardBatch
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.speculative.base_spec_worker import BaseDraftWorker, BaseSpecWorker
+from sglang.srt.speculative.draft_utils import DraftBackendFactory
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.eagle_info_v2 import fill_new_verified_id
 from sglang.srt.speculative.eagle_utils import TreeMaskMode, build_tree_kernel_efficient
@@ -183,16 +184,14 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker):
         # Create attn backends
         self.draft_extend_attn_backend_list = []
         for step in range(self.speculative_num_steps):
-            from sglang.srt.layers.attention.flashattention_backend import (
+            draft_backend_factory = DraftBackendFactory(
-                FlashAttentionBackend,
+                self.server_args,
+                self.draft_runner_list[step],
+                self.topk,
+                self.speculative_num_steps,
             )
-
             self.draft_extend_attn_backend_list.append(
-                FlashAttentionBackend(
+                draft_backend_factory.create_draft_extend_backend()
-                    model_runner=self.draft_runner_list[step],
-                    skip_prefill=False,
-                    speculative_step_id=step,
-                )
             )
             self.draft_runner_list[step].attn_backend = (
                 self.draft_extend_attn_backend_list[-1]

test/registered/8-gpu-models/test_step3p5_flash_chain_mtp.py

@@ -0,0 +1,105 @@
+import unittest
+from types import SimpleNamespace
+
+import requests
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.run_eval import run_eval
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    is_in_ci,
+    popen_launch_server,
+    write_github_step_summary,
+)
+
+register_cuda_ci(est_time=500, suite="stage-c-test-8-gpu-h200")
+
+STEP3P5_FLASH_MODEL_PATH = "stepfun-ai/Step-3.5-Flash"
+
+
+class TestStep3p5FlashChainMTP(CustomTestCase):
+    """Chain-style multi-layer EAGLE speculative decoding on Step-3.5-Flash.
+
+    Step3p5ForCausalLM auto-enables multi-layer EAGLE and spec v2 when
+    --speculative-algorithm=EAGLE is set.  The chain MTP propagation
+    (each MTP layer feeds its hidden states to the next) is activated
+    automatically for the Step3p5MTP draft architecture.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = STEP3P5_FLASH_MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = [
+            "--tp",
+            "8",
+            "--trust-remote-code",
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-num-steps",
+            "3",
+            "--speculative-eagle-topk",
+            "1",
+            "--speculative-num-draft-tokens",
+            "4",
+            "--attention-backend",
+            "fa3",
+            "--enable-multi-layer-eagle",
+            "--mem-fraction-static",
+            "0.75",
+            "--chunked-prefill-size",
+            "4096",
+            "--model-loader-extra-config",
+            '{"enable_multithread_load": true, "num_threads": 64}',
+        ]
+        with envs.SGLANG_ENABLE_SPEC_V2.override(True):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH * 3,
+                other_args=other_args,
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
+
+    def test_gsm8k(self):
+        requests.get(self.base_url + "/flush_cache")
+
+        args = SimpleNamespace(
+            base_url=self.base_url,
+            model=self.model,
+            eval_name="gsm8k",
+            api="completion",
+            max_tokens=512,
+            num_examples=200,
+            num_threads=128,
+        )
+        metrics = run_eval(args)
+        print(f"{metrics=}")
+
+        server_info = requests.get(self.base_url + "/server_info")
+        avg_spec_accept_length = server_info.json()["internal_states"][0][
+            "avg_spec_accept_length"
+        ]
+        print(f"{avg_spec_accept_length=}")
+        print(f"{metrics=}")
+
+        if is_in_ci():
+            write_github_step_summary(
+                f"### test_gsm8k (step-3.5-flash chain mtp)\n"
+                f'{metrics["score"]=:.3f}\n'
+                f"{avg_spec_accept_length=:.2f}\n"
+            )
+            self.assertGreater(metrics["score"], 0.84)
+            self.assertGreater(avg_spec_accept_length, 2.6)
+
+
+if __name__ == "__main__":
+    unittest.main()