[core] Fix crashes on the gpu_only spec_v2 path (#26738)

python/sglang/srt/layers/attention/triton_backend.py

@@ -537,9 +537,13 @@ class TritonAttnBackend(AttentionBackend):
                 dtype=torch.int32,
                 device=self.device,
             )
-            # Different with flashinfer kv_indptr and kv_indices construction
+            # Different with flashinfer kv_indptr and kv_indices construction.
+            # gpu_only: seq_lens_sum may be None; ub-allocate is safe (ragged write).
+            seq_lens_sum = forward_batch.seq_lens_sum
+            if seq_lens_sum is None:
+                seq_lens_sum = bs * self.max_context_len
             kv_indices = torch.empty(
-                forward_batch.seq_lens_sum, dtype=torch.int64, device=self.device
+                seq_lens_sum, dtype=torch.int64, device=self.device
             )
             kv_indptr = self._fill_kv_indptr_and_indices(
                 bs,
@@ -605,8 +609,14 @@ class TritonAttnBackend(AttentionBackend):
             attn_logits = None
             attn_lse = None
         else:
+            # gpu_only leaves _cpu unset; ub-allocate is safe (ragged write
+            # from GPU tensor, extra tail unused).
+            if forward_batch.extend_prefix_lens_cpu is not None:
+                kv_indices_len = sum(forward_batch.extend_prefix_lens_cpu)
+            else:
+                kv_indices_len = bs * self.max_context_len
             kv_indices = torch.empty(
-                sum(forward_batch.extend_prefix_lens_cpu),
+                kv_indices_len,
                 dtype=torch.int64,
                 device=self.device,
             )
@@ -641,7 +651,12 @@ class TritonAttnBackend(AttentionBackend):
             mask_indptr = None
             attn_logits = None
             attn_lse = None
-            max_extend_len = max(forward_batch.extend_seq_lens_cpu)
+            # Caller usually supplies extend_seq_lens_cpu (eagle_info gpu_only
+            # sets host-constant mirror); defensive GPU-max fallback if not.
+            if forward_batch.extend_seq_lens_cpu is not None:
+                max_extend_len = max(forward_batch.extend_seq_lens_cpu)
+            else:
+                max_extend_len = int(forward_batch.extend_seq_lens.max())
             num_kv_splits = None
 
         self.forward_metadata = ForwardMetadata(

python/sglang/srt/managers/overlap_utils.py

@@ -32,7 +32,11 @@ def decide_needs_cpu_seq_lens(
     if not server_args.disable_piecewise_cuda_graph:
         # FIXME: support PCG without seq lens cpu value
         return True
-    return any(b.needs_cpu_seq_lens for b in attn_backends)
+    # Skip unset slots (e.g. draft_extend_attn_backend on some spec configs);
+    # missing flag -> True so undeclared backends stay on the legacy path.
+    return any(
+        getattr(b, "needs_cpu_seq_lens", True) for b in attn_backends if b is not None
+    )
 
 
 _is_cuda = is_cuda()

python/sglang/srt/speculative/eagle_info_v2.py

@@ -266,6 +266,10 @@ class EagleDraftInputV2Mixin:
         if not gpu_only:
             forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens
             forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum())
+        else:
+            # Supply CPU mirror (extend_seq_lens are all num_draft_tokens) so
+            # backend max() reads from list without a per-iter D2H sync.
+            forward_batch.extend_seq_lens_cpu = [num_draft_tokens] * bs
         can_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run(forward_batch)
         if not batch.forward_mode.is_idle() and not can_cuda_graph:
             draft_model_runner.attn_backend.init_forward_metadata(forward_batch)

python/sglang/test/kits/fwd_occupancy_kit.py

@@ -41,7 +41,7 @@ class FwdOccupancyMixin:
     # in /server_info); silently skipped otherwise. EAGLE3 3/1/4 on
     # 5090 + Llama-3.1-8B measured ~2.0 in CI; 1.8 leaves a small
     # buffer while still catching silent fallback to vanilla (~1.0).
-    spec_accept_length_threshold: float = 1.8
+    fwd_occupancy_acc_length_threshold: float = 1.8
 
     # Warmup: one short request to fill cuda graphs + get the
     # device-timer past its first NaN window.
@@ -51,7 +51,9 @@ class FwdOccupancyMixin:
     # Measurement: one long single-batch request -- max_new_tokens must
     # span several decode_log_interval windows for enough samples.
     fwd_occupancy_max_new_tokens: int = 2048
-    fwd_occupancy_prompt: str = "Write a long, detailed, multi-paragraph story about "
+    fwd_occupancy_prompt: str = (
+        "Human: Give me a fully functional FastAPI server. Show the python code.\n\nAssistant:"
+    )
 
     def _scrape_fwd_occupancy(self):
         """Max non-NaN gauge value across exposed labels (e.g. dp ranks);
@@ -202,8 +204,8 @@ class FwdOccupancyMixin:
             print(f"avg_spec_accept_length = {avg_accept:.3f}")
             self.assertGreater(
                 avg_accept,
-                self.spec_accept_length_threshold,
+                self.fwd_occupancy_acc_length_threshold,
                 f"avg_spec_accept_length={avg_accept:.3f} did not exceed "
-                f"threshold {self.spec_accept_length_threshold} -- spec "
+                f"threshold {self.fwd_occupancy_acc_length_threshold} -- spec "
                 "barely accepted, possibly degraded to vanilla decode",
             )

test/registered/core/test_basic_sanity_eagle3.py

@@ -8,8 +8,8 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kits.basic_api_contract_kit import BasicAPIContractMixin
 from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
 from sglang.test.kits.basic_scheduler_stress_kit import BasicSchedulerStressMixin
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.kits.fwd_occupancy_kit import FwdOccupancyMixin
-from sglang.test.kits.hellaswag_kit import HellaswagMixin
 from sglang.test.test_utils import (
     DEFAULT_DRAFT_MODEL_EAGLE3,
     DEFAULT_TARGET_MODEL_EAGLE3,
@@ -29,7 +29,7 @@ class TestBasicSanityEagle3(
     BasicDecodeCorrectnessMixin,
     BasicSchedulerStressMixin,
     FwdOccupancyMixin,
-    HellaswagMixin,
+    GSM8KMixin,
     CustomTestCase,
 ):
     served_model_name = DEFAULT_TARGET_MODEL_EAGLE3
@@ -38,6 +38,11 @@ class TestBasicSanityEagle3(
     # measurement window to avoid too few non-NaN samples.
     fwd_occupancy_threshold = 80.0 if is_in_amd_ci() else 97.0
     fwd_occupancy_max_new_tokens = 4096 if is_in_amd_ci() else 2048
+    fwd_occupancy_acc_length_threshold: float = 1.6
+
+    model = DEFAULT_TARGET_MODEL_EAGLE3
+    gsm8k_num_questions = 1400
+    gsm8k_accuracy_thres = 0.74
 
     @classmethod
     def setUpClass(cls):
@@ -59,16 +64,17 @@ class TestBasicSanityEagle3(
                 "--speculative-draft-model-path",
                 DEFAULT_DRAFT_MODEL_EAGLE3,
                 "--speculative-num-steps",
-                "3",
+                "1",
                 "--speculative-eagle-topk",
                 "1",
                 "--speculative-num-draft-tokens",
-                "4",
+                "2",
                 "--cuda-graph-max-bs",
                 "4",
                 "--mem-fraction-static",
                 "0.7",
                 "--enable-metrics",
+                "--disable-piecewise-cuda-graph",
             ],
             env={"SGLANG_ENABLE_METRICS_DEVICE_TIMER": "1"},
         )