diff --git a/python/sglang/srt/layers/attention/triton_backend.py b/python/sglang/srt/layers/attention/triton_backend.py index 333df0827..d71871e5b 100644 --- a/python/sglang/srt/layers/attention/triton_backend.py +++ b/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( diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 51e4d75f6..c4fff2fbd 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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() diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index ad65a9556..fb0c91b68 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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) diff --git a/python/sglang/test/kits/fwd_occupancy_kit.py b/python/sglang/test/kits/fwd_occupancy_kit.py index 0f66eaed5..4a34e7f63 100644 --- a/python/sglang/test/kits/fwd_occupancy_kit.py +++ b/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", ) diff --git a/test/registered/core/test_basic_sanity_eagle3.py b/test/registered/core/test_basic_sanity_eagle3.py index 72ef8ffff..ce3c4d192 100644 --- a/test/registered/core/test_basic_sanity_eagle3.py +++ b/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"}, )