[core] Fix crashes on the gpu_only spec_v2 path (#26738)
This commit is contained in:
@@ -537,9 +537,13 @@ class TritonAttnBackend(AttentionBackend):
|
|||||||
dtype=torch.int32,
|
dtype=torch.int32,
|
||||||
device=self.device,
|
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(
|
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(
|
kv_indptr = self._fill_kv_indptr_and_indices(
|
||||||
bs,
|
bs,
|
||||||
@@ -605,8 +609,14 @@ class TritonAttnBackend(AttentionBackend):
|
|||||||
attn_logits = None
|
attn_logits = None
|
||||||
attn_lse = None
|
attn_lse = None
|
||||||
else:
|
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(
|
kv_indices = torch.empty(
|
||||||
sum(forward_batch.extend_prefix_lens_cpu),
|
kv_indices_len,
|
||||||
dtype=torch.int64,
|
dtype=torch.int64,
|
||||||
device=self.device,
|
device=self.device,
|
||||||
)
|
)
|
||||||
@@ -641,7 +651,12 @@ class TritonAttnBackend(AttentionBackend):
|
|||||||
mask_indptr = None
|
mask_indptr = None
|
||||||
attn_logits = None
|
attn_logits = None
|
||||||
attn_lse = None
|
attn_lse = None
|
||||||
|
# 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)
|
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
|
num_kv_splits = None
|
||||||
|
|
||||||
self.forward_metadata = ForwardMetadata(
|
self.forward_metadata = ForwardMetadata(
|
||||||
|
|||||||
@@ -32,7 +32,11 @@ def decide_needs_cpu_seq_lens(
|
|||||||
if not server_args.disable_piecewise_cuda_graph:
|
if not server_args.disable_piecewise_cuda_graph:
|
||||||
# FIXME: support PCG without seq lens cpu value
|
# FIXME: support PCG without seq lens cpu value
|
||||||
return True
|
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()
|
_is_cuda = is_cuda()
|
||||||
|
|||||||
@@ -266,6 +266,10 @@ class EagleDraftInputV2Mixin:
|
|||||||
if not gpu_only:
|
if not gpu_only:
|
||||||
forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens
|
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())
|
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)
|
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:
|
if not batch.forward_mode.is_idle() and not can_cuda_graph:
|
||||||
draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
|
draft_model_runner.attn_backend.init_forward_metadata(forward_batch)
|
||||||
|
|||||||
@@ -41,7 +41,7 @@ class FwdOccupancyMixin:
|
|||||||
# in /server_info); silently skipped otherwise. EAGLE3 3/1/4 on
|
# 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
|
# 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).
|
# 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
|
# Warmup: one short request to fill cuda graphs + get the
|
||||||
# device-timer past its first NaN window.
|
# device-timer past its first NaN window.
|
||||||
@@ -51,7 +51,9 @@ class FwdOccupancyMixin:
|
|||||||
# Measurement: one long single-batch request -- max_new_tokens must
|
# Measurement: one long single-batch request -- max_new_tokens must
|
||||||
# span several decode_log_interval windows for enough samples.
|
# span several decode_log_interval windows for enough samples.
|
||||||
fwd_occupancy_max_new_tokens: int = 2048
|
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):
|
def _scrape_fwd_occupancy(self):
|
||||||
"""Max non-NaN gauge value across exposed labels (e.g. dp ranks);
|
"""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}")
|
print(f"avg_spec_accept_length = {avg_accept:.3f}")
|
||||||
self.assertGreater(
|
self.assertGreater(
|
||||||
avg_accept,
|
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"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",
|
"barely accepted, possibly degraded to vanilla decode",
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -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_api_contract_kit import BasicAPIContractMixin
|
||||||
from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
|
from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectnessMixin
|
||||||
from sglang.test.kits.basic_scheduler_stress_kit import BasicSchedulerStressMixin
|
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.fwd_occupancy_kit import FwdOccupancyMixin
|
||||||
from sglang.test.kits.hellaswag_kit import HellaswagMixin
|
|
||||||
from sglang.test.test_utils import (
|
from sglang.test.test_utils import (
|
||||||
DEFAULT_DRAFT_MODEL_EAGLE3,
|
DEFAULT_DRAFT_MODEL_EAGLE3,
|
||||||
DEFAULT_TARGET_MODEL_EAGLE3,
|
DEFAULT_TARGET_MODEL_EAGLE3,
|
||||||
@@ -29,7 +29,7 @@ class TestBasicSanityEagle3(
|
|||||||
BasicDecodeCorrectnessMixin,
|
BasicDecodeCorrectnessMixin,
|
||||||
BasicSchedulerStressMixin,
|
BasicSchedulerStressMixin,
|
||||||
FwdOccupancyMixin,
|
FwdOccupancyMixin,
|
||||||
HellaswagMixin,
|
GSM8KMixin,
|
||||||
CustomTestCase,
|
CustomTestCase,
|
||||||
):
|
):
|
||||||
served_model_name = DEFAULT_TARGET_MODEL_EAGLE3
|
served_model_name = DEFAULT_TARGET_MODEL_EAGLE3
|
||||||
@@ -38,6 +38,11 @@ class TestBasicSanityEagle3(
|
|||||||
# measurement window to avoid too few non-NaN samples.
|
# measurement window to avoid too few non-NaN samples.
|
||||||
fwd_occupancy_threshold = 80.0 if is_in_amd_ci() else 97.0
|
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_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
|
@classmethod
|
||||||
def setUpClass(cls):
|
def setUpClass(cls):
|
||||||
@@ -59,16 +64,17 @@ class TestBasicSanityEagle3(
|
|||||||
"--speculative-draft-model-path",
|
"--speculative-draft-model-path",
|
||||||
DEFAULT_DRAFT_MODEL_EAGLE3,
|
DEFAULT_DRAFT_MODEL_EAGLE3,
|
||||||
"--speculative-num-steps",
|
"--speculative-num-steps",
|
||||||
"3",
|
"1",
|
||||||
"--speculative-eagle-topk",
|
"--speculative-eagle-topk",
|
||||||
"1",
|
"1",
|
||||||
"--speculative-num-draft-tokens",
|
"--speculative-num-draft-tokens",
|
||||||
"4",
|
"2",
|
||||||
"--cuda-graph-max-bs",
|
"--cuda-graph-max-bs",
|
||||||
"4",
|
"4",
|
||||||
"--mem-fraction-static",
|
"--mem-fraction-static",
|
||||||
"0.7",
|
"0.7",
|
||||||
"--enable-metrics",
|
"--enable-metrics",
|
||||||
|
"--disable-piecewise-cuda-graph",
|
||||||
],
|
],
|
||||||
env={"SGLANG_ENABLE_METRICS_DEVICE_TIMER": "1"},
|
env={"SGLANG_ENABLE_METRICS_DEVICE_TIMER": "1"},
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user