[Feature] Spec V2 DFlash Support (#23000)

test/registered/cuda_graph/piecewise/test_pcg_with_speculative_decoding_dflash.py

@@ -36,6 +36,10 @@ class TestPCGWithDFlash(PCGSpecBase, CustomTestCase):
         "1",
         "--max-running-requests",
         "64",
+        # Keep headroom for the draft KV pool + piecewise cuda graph
+        # private pools on 32GB CI cards.
+        "--mem-fraction-static",
+        "0.7",
         "--cuda-graph-bs-decode",
         *[str(i) for i in range(1, 65)],
     ]

test/registered/spec/dflash/test_dflash.py

@@ -1,4 +1,3 @@
-import os
 import unittest
 
 import openai
@@ -8,7 +7,10 @@ from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.kits.matched_stop_kit import MatchedStopMixin
-from sglang.test.kits.radix_cache_server_kit import gen_radix_tree
+from sglang.test.kits.radix_cache_server_kit import (
+    gen_radix_tree,
+    run_radix_attention_test,
+)
 from sglang.test.test_utils import (
     DEFAULT_DRAFT_MODEL_DFLASH,
     DEFAULT_TARGET_MODEL_DFLASH,
@@ -26,6 +28,8 @@ class TestDFlashServerBase(CustomTestCase, MatchedStopMixin, GSM8KMixin):
     attention_backend = "flashinfer"
     page_size = 1
     other_launch_args = []
+    spec_v2 = False
+    overlap_plan_stream = False
     model = DEFAULT_TARGET_MODEL_DFLASH
     draft_model = DEFAULT_DRAFT_MODEL_DFLASH
     gsm8k_accuracy_thres = 0.75
@@ -46,29 +50,32 @@ class TestDFlashServerBase(CustomTestCase, MatchedStopMixin, GSM8KMixin):
             str(cls.page_size),
             "--max-running-requests",
             str(cls.max_running_requests),
+            # Keep headroom for the draft KV pool + piecewise cuda graph
+            # private pools on 32GB CI cards.
+            "--mem-fraction-static",
+            "0.7",
             "--cuda-graph-bs",
             *[str(i) for i in range(1, cls.max_running_requests + 1)],
         ]
         launch_args.extend(cls.other_launch_args)
-        old_value = os.environ.get("SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN")
-        os.environ["SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN"] = "1"
-        try:
-            with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1):
-                cls.process = popen_launch_server(
-                    cls.model,
-                    cls.base_url,
-                    timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-                    other_args=launch_args,
-                )
-        finally:
-            if old_value is None:
-                del os.environ["SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN"]
-            else:
-                os.environ["SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN"] = old_value
+        with (
+            envs.SGLANG_ENABLE_SPEC_V2.override(cls.spec_v2),
+            envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.override(cls.overlap_plan_stream),
+            envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1),
+            envs.SGLANG_ENABLE_ASYNC_ASSERT.override(True),
+            envs.SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.override(True),
+        ):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+                other_args=launch_args,
+            )
 
     @classmethod
     def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
+        if hasattr(cls, "process") and cls.process:
+            kill_process_tree(cls.process.pid)
 
     def test_early_stop(self):
         client = openai.Client(base_url=self.base_url + "/v1", api_key="EMPTY")
@@ -142,5 +149,17 @@ class TestDFlashServerNoCudaGraph(TestDFlashServerBase):
     other_launch_args = ["--disable-cuda-graph"]
 
 
+class TestDFlashServerSpecV2(TestDFlashServerBase):
+    spec_v2 = True
+
+    def test_radix_attention(self):
+        run_radix_attention_test(self.base_url)
+        assert self.process.poll() is None
+
+
+class TestDFlashServerSpecV2PlanStream(TestDFlashServerSpecV2):
+    overlap_plan_stream = True
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/spec/test_decode_bookkeeping_ownership.py

@@ -60,8 +60,15 @@ _OWNER_SITES = {
     (*_MIXIN, "evict"): 1,
     (*_MIXIN, "kv_committed_len"): 1,
     (*_MIXIN, "kv_allocated_len"): 1,
-    (*_RESOLVE, "kv_committed_len"): 2,
+    # 3rd resolve mutation: DFLASH settles its full commit_lens here (no
+    # pre-claim in prepare_for_decode, unlike the EAGLE mixin).
+    (*_RESOLVE, "kv_committed_len"): 3,
     (*_RESOLVE, "spec_verify_ct"): 1,
+    (
+        "speculative/dflash_info_v2.py",
+        "DFlashDraftInputV2.prepare_for_decode",
+        "kv_allocated_len",
+    ): 1,
     # spec v1: each verify path owns its own settlement
     ("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_committed_len"): 1,
     ("speculative/eagle_info.py", "EagleVerifyInput.verify", "kv_allocated_len"): 1,