minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Enable breakable prefill CUDA graph for DP attention (#30898)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-07-12 17:10:04 -07:00
committed by GitHub
parent c616d5a55e
commit b94ac87e0c
9 changed files with 478 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler_components/dp_attn.py
+13 -5
View File
@@ -14,7 +14,12 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.cuda_graph_config import (
Backend,
Phase,
check_cuda_graph_backend,
cuda_graph_fully_disabled,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.observability.metrics_collector import DPCooperationInfo
from sglang.srt.server_args import ServerArgs
@@ -184,11 +189,14 @@ def prepare_mlp_sync_batch_raw(
or local_batch.forward_mode.is_decode_or_idle()
or local_batch.forward_mode.is_prebuilt()
) and not disable_cuda_graph
# Idle/None ranks are permissive (like can_cuda_graph): the all-gather
# min()-reduces this across DP ranks, so a prefill batch with idle ranks
# still resolves to True (idle ranks become a padded dummy extend).
can_run_breakable_cuda_graph = (
local_batch is not None
and local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
and not disable_cuda_graph
)
local_batch is None
or local_batch.forward_mode.is_idle()
or local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED)
) and check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE)
is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False
if local_batch is not None:
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+10
View File
@@ -788,6 +788,7 @@ def build_prefill_registry(
hidden_size: int = 0,
embed_dtype: Optional[torch.dtype] = None,
enable_mamba_track: bool = False,
enable_num_token_non_padded: bool = False,
register_input_embeds: bool = True,
share_pool: bool = True,
source: Optional[Any] = None,
@@ -876,6 +877,15 @@ def build_prefill_registry(
slots.append(GraphSlot("mamba_track_indices", _bs, torch.int64, axis="bs"))
slots.append(GraphSlot("mamba_track_mask", _bs, torch.bool, axis="bs"))
slots.append(GraphSlot("mamba_track_seqlens", _bs, torch.int32, axis="bs"))
if enable_num_token_non_padded:
slots.append(
GraphSlot(
"num_token_non_padded",
lambda _bs2, _mt: (1,),
torch.int32,
axis="none",
)
)
for slot in slots:
bind = None
python/sglang/srt/model_executor/forward_batch_info.py
+39 -1
View File
@@ -1177,6 +1177,26 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
self.is_extend_in_batch, global_num_tokens
)
# Prefill breakable CUDA graph requires every DP rank to run the SAME
# captured shape. Under SUM_LEN each rank pads to its own local token
# count and can select a different capture bucket, so the in-graph DP
# collectives (all_gather / reduce_scatter) mismatch across ranks and
# corrupt the output. Force MAX_LEN so every rank pads to the global
# max and picks the same bucket (mirrors the decode cuda graph
# contract, which always runs MAX_LEN).
#
# Only force MAX_LEN when the batch fits a captured breakable prefill
# graph; larger prefills fall back to eager and keep the
# memory-efficient SUM_LEN. global_num_tokens is identical across ranks
# (all-gathered), so the decision is consistent cluster-wide.
prefill_cg = model_runner.server_args.cuda_graph_config.prefill
if (
self.can_run_dp_breakable_cuda_graph
and self.is_extend_in_batch
and prefill_cg.bs
and max(global_num_tokens) <= max(prefill_cg.bs)
):
dp_padding_mode = DpPaddingMode.MAX_LEN
self.dp_padding_mode = dp_padding_mode
if dp_padding_mode.is_max_len():
@@ -1233,7 +1253,13 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
elif self.is_extend_in_batch and dp_padding_mode.is_max_len():
self._original_forward_mode = self.forward_mode
self.forward_mode = ForwardMode.EXTEND
if hybrid_ssm:
# Fabricate a single dummy request covering num_tokens for an
# empty (idle) rank. Hybrid-SSM families always take this path;
# non-hybrid ranks reach it once MAX_LEN is forced for the
# prefill breakable CUDA graph (idle + prefill), which needs
# every DP rank to run the same captured shape. The `else`
# branch handles decode rows padded to a 1-token extend.
if hybrid_ssm or self.seq_lens.shape[0] == 0:
dev = self.seq_lens.device
assert (
self.seq_lens.shape[0] == 0
@@ -1251,6 +1277,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
self.seq_lens = torch.tensor(
[num_tokens], dtype=self.seq_lens.dtype, device=dev
)
# orig_seq_lens is not padded by _pad_inputs_to_size, so
# fabricate it to match the dummy request (the breakable
# prefill CUDA graph runner reads it).
self.orig_seq_lens = torch.tensor(
[num_tokens], dtype=self.orig_seq_lens.dtype, device=dev
)
self.seq_lens_sum = int(num_tokens)
if self.seq_lens_cpu is not None:
self.seq_lens_cpu = torch.tensor(
@@ -1260,6 +1292,12 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
self.extend_seq_lens_cpu = [int(num_tokens)]
self.extend_logprob_start_lens_cpu = [0]
bs = self.batch_size = 1
# Count the dummy tokens as real, else MoE topk/all-to-all
# treats this rank as empty and starves later layers.
# (num_token_non_padded is None unless moe_ep_size > 1.)
if self.num_token_non_padded is not None:
self.num_token_non_padded.fill_(num_tokens)
self.num_token_non_padded_cpu = num_tokens
else:
self.extend_num_tokens = bs
self.extend_seq_lens = torch.full_like(self.seq_lens, 1)
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+70 -8
View File
@@ -61,6 +61,8 @@ from sglang.srt.model_executor.forward_batch_info import (
ForwardBatch,
ForwardMode,
PPProxyTensors,
compute_local_num_token_non_padded,
enable_num_token_non_padded,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
@@ -98,6 +100,7 @@ from sglang.srt.utils import (
is_hip,
is_npu,
require_attn_tp_gather,
require_gathered_buffer,
require_mlp_tp_gather,
)
@@ -229,6 +232,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
hidden_size=self.model_runner.model_config.hidden_size,
embed_dtype=self.model_runner.dtype,
enable_mamba_track=self.mamba_track_enabled,
enable_num_token_non_padded=enable_num_token_non_padded(),
source=self.buffers,
)
@@ -381,6 +385,39 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.model_runner.model_config.vocab_size, rows=rows
)
def _prefill_logits_buffer_rows(self, forward_batch: ForwardBatch) -> int:
if not forward_batch.return_logprob:
return forward_batch.batch_size
if not isinstance(self.backend, BreakableCudaGraphBackend):
return forward_batch.batch_size
global_num_tokens = forward_batch.global_num_tokens_for_logprob_cpu
if global_num_tokens is not None:
dp_rank = get_parallel().attn_dp_rank
return int(global_num_tokens[dp_rank if len(global_num_tokens) > 1 else 0])
return sum(
max(int(seq_len) - int(start_len), 1)
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
)
)
def _capture_num_token_non_padded(self, num_tokens: int) -> Optional[torch.Tensor]:
if not self.buffer_registry.has_slot("num_token_non_padded"):
return None
buf = self.buffer_registry.get_slot("num_token_non_padded").buffer
buf.fill_(num_tokens)
if require_gathered_buffer(self.model_runner.server_args):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=buf,
num_tokens_per_dp=num_tokens,
)
buf.copy_(local)
return buf
_aiter_chip_info_cached = False
@classmethod
@@ -592,7 +629,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
):
return False
num_tokens = len(forward_batch.input_ids)
if forward_batch.return_logprob:
if forward_batch.return_logprob and not isinstance(
self.backend, BreakableCudaGraphBackend
):
for start_len, seq_len in zip(
forward_batch.extend_logprob_start_lens_cpu,
forward_batch.extend_seq_lens_cpu,
@@ -631,7 +670,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
Returns ``(forward_batch, attn_backend)`` to mirror decode's
capture_prepare signature.
"""
buffers = self.buffers
bs = self._capture_req_slots
# Slot 0 carries num_tokens; slots 1..bs-1 are zero-length sentinels.
lens_cpu = [num_tokens] + [0] * (bs - 1)
@@ -748,7 +786,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# FULL aux hidden states) captures with the right mode.
# Ported from main #27468.
capture_hidden_mode=self.capture_hidden_mode,
num_token_non_padded=None,
num_token_non_padded=self._capture_num_token_non_padded(num_tokens),
num_token_non_padded_cpu=num_tokens,
global_forward_mode=ForwardMode.EXTEND,
lora_ids=None,
@@ -829,7 +867,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
"""Pad, populate static buffers, and build the static_forward_batch
the model code reads during replay.
"""
buffers = self.buffers
num_tokens = len(forward_batch.input_ids)
static_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens)
self.raw_num_tokens = num_tokens
@@ -881,6 +918,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
and forward_batch.mrope_positions is not None
else None
)
num_token_non_padded = (
_slot("num_token_non_padded")
if registry.has_slot("num_token_non_padded")
else forward_batch.num_token_non_padded
)
# Normalize MIXED→EXTEND so dynamo's guard (captured with EXTEND=1)
# doesn't fail on MIXED=3.
@@ -902,7 +944,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
input_embeds=input_embeds,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
next_token_logits_buffer=self._next_token_logits_buffer(bs),
next_token_logits_buffer=self._next_token_logits_buffer(
self._prefill_logits_buffer_rows(forward_batch)
),
orig_seq_lens=forward_batch.orig_seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
out_cache_loc=out_cache_loc,
@@ -911,25 +955,34 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
encoder_lens=forward_batch.encoder_lens,
return_logprob=False,
return_logprob=(
forward_batch.return_logprob
if isinstance(self.backend, BreakableCudaGraphBackend)
else False
),
is_prefill_only=forward_batch.is_prefill_only,
extend_seq_lens=forward_batch.extend_seq_lens,
extend_prefix_lens=forward_batch.extend_prefix_lens,
extend_start_loc=forward_batch.extend_start_loc,
extend_prefix_lens_cpu=forward_batch.extend_prefix_lens_cpu,
extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu,
extend_logprob_start_lens_cpu=forward_batch.extend_logprob_start_lens_cpu,
top_logprobs_nums=forward_batch.top_logprobs_nums,
token_ids_logprobs=forward_batch.token_ids_logprobs,
multi_item_delimiter_indices=forward_batch.multi_item_delimiter_indices,
extend_num_tokens=forward_batch.extend_num_tokens,
extend_input_logprob_token_ids_gpu=forward_batch.extend_input_logprob_token_ids_gpu,
positions=positions,
global_num_tokens_gpu=forward_batch.global_num_tokens_gpu,
global_num_tokens_for_logprob_gpu=forward_batch.global_num_tokens_for_logprob_gpu,
global_num_tokens_for_logprob_cpu=forward_batch.global_num_tokens_for_logprob_cpu,
dp_padding_mode=forward_batch.dp_padding_mode,
global_dp_buffer_len=forward_batch.global_dp_buffer_len,
mrope_positions=mrope_positions,
spec_algorithm=forward_batch.spec_algorithm,
spec_info=forward_batch.spec_info,
capture_hidden_mode=forward_batch.capture_hidden_mode,
num_token_non_padded=forward_batch.num_token_non_padded,
num_token_non_padded=num_token_non_padded,
num_token_non_padded_cpu=forward_batch.num_token_non_padded_cpu,
global_forward_mode=pcg_global_forward_mode,
lora_ids=forward_batch.lora_ids,
@@ -1096,12 +1149,21 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.raw_bs if self._is_full_backend else self.raw_num_tokens
)
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[:logits_rows],
next_token_logits=(
output.next_token_logits[:logits_rows]
if output.next_token_logits is not None
else None
),
hidden_states=(
output.hidden_states[: self.raw_num_tokens]
if output.hidden_states is not None
else None
),
input_token_logprobs=output.input_token_logprobs,
input_top_logprobs_val=output.input_top_logprobs_val,
input_top_logprobs_idx=output.input_top_logprobs_idx,
input_token_ids_logprobs_val=output.input_token_ids_logprobs_val,
input_token_ids_logprobs_idx=output.input_token_ids_logprobs_idx,
mm_input_embeds=mm_input_embeds,
)
elif isinstance(output, EmbeddingPoolerOutput):
python/sglang/srt/model_executor/runner_utils/buffers.py
+3 -1
View File
@@ -62,7 +62,6 @@ def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -
@dataclass
class DecodeInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
input_embeds: torch.Tensor
req_pool_indices: torch.Tensor
@@ -328,6 +327,7 @@ class DecodeInputBuffers(ForwardInputBuffers):
class PrefillInputBuffers(ForwardInputBuffers):
input_ids: torch.Tensor
out_cache_loc: torch.Tensor
num_token_non_padded: torch.Tensor
mamba_track_indices: Optional[torch.Tensor]
mamba_track_mask: Optional[torch.Tensor]
mamba_track_seqlens: Optional[torch.Tensor]
@@ -351,6 +351,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
with torch.device(device):
input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
mamba_track_indices = (
torch.zeros((max_bs,), dtype=torch.int64)
if enable_mamba_track
@@ -376,6 +377,7 @@ class PrefillInputBuffers(ForwardInputBuffers):
return cls(
input_ids=input_ids,
out_cache_loc=out_cache_loc,
num_token_non_padded=num_token_non_padded,
mamba_track_indices=mamba_track_indices,
mamba_track_mask=mamba_track_mask,
mamba_track_seqlens=mamba_track_seqlens,
python/sglang/srt/server_args.py
+33 -5
View File
@@ -219,9 +219,15 @@ ATTENTION_BACKEND_CHOICES = [
"intel_xpu",
]
DETERMINISTIC_ATTENTION_BACKEND_CHOICES = ["flashinfer", "fa3", "triton", "ascend"]
DETERMINISTIC_ATTENTION_BACKEND_CHOICES = [
"ascend",
"fa3",
"fa4",
"flashinfer",
"triton",
]
RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["fa3", "triton", "ascend"]
RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND = ["ascend", "fa3", "fa4", "triton"]
DISAGG_TRANSFER_BACKEND_CHOICES = [
"mooncake",
@@ -5256,11 +5262,32 @@ class ServerArgs:
if self._resolved().enable_dp_attention:
self.schedule_conservativeness = self.schedule_conservativeness * 0.3
assert self.tp_size % self.dp_size == 0
original_chunked_prefill_size = self.chunked_prefill_size
self.chunked_prefill_size = self.chunked_prefill_size // self.dp_size
logger.warning(
f"DP attention is enabled. The chunked prefill size is adjusted to {self.chunked_prefill_size} to avoid MoE kernel issues. "
f"DP attention is enabled. chunked prefill size is adjusted "
f"from {original_chunked_prefill_size} to {self.chunked_prefill_size}."
)
# The prefill CUDA graph max_bs was derived from the pre-DP-division
# chunked_prefill_size in _handle_gpu_memory_settings (which runs
# before this handler). Re-clamp it (and the captured shape list) to
# the per-DP-rank chunked_prefill_size so breakable CUDA graph
# capture never exceeds the MoE all-to-all's max_num_tokens budget,
# which is also sized from the DP-adjusted chunked_prefill_size.
prefill_cfg = self.cuda_graph_config.prefill
if (
prefill_cfg.backend != Backend.DISABLED
and prefill_cfg.max_bs is not None
and prefill_cfg.max_bs > self.chunked_prefill_size
and (Phase.PREFILL, "max_bs") not in self._cuda_graph_config_locked
):
prefill_cfg.max_bs = self.chunked_prefill_size
if (Phase.PREFILL, "bs") not in self._cuda_graph_config_locked:
prefill_cfg.bs = self._generate_prefill_cuda_graph_batch_sizes(
prefill_cfg.max_bs
)
# The dp-lm-head validation moved to the resolution pipeline
# (arg_groups/overrides.py: _dp_lm_head_validation), invoked here at
# its legacy slot.
@@ -6190,9 +6217,10 @@ class ServerArgs:
attention_backend = resolved_view(self).attention_backend
if is_deepseek_model:
if attention_backend not in ["fa3", "triton"]:
deepseek_deterministic_attention_backends = ["fa3", "triton"]
if attention_backend not in deepseek_deterministic_attention_backends:
raise ValueError(
f"Currently only {RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
f"Currently only {deepseek_deterministic_attention_backends} attention backends are supported for deterministic inference with DeepSeek models. But you're using {attention_backend}."
)
if attention_backend not in RADIX_SUPPORTED_DETERMINISTIC_ATTENTION_BACKEND:
test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py
+2
View File
@@ -80,6 +80,8 @@ class TestDSV4FlashFP4B200Balanced_CP(
"round-robin-split",
"--deepep-config",
DEEPEP_CONFIG,
"--mem-fraction-static",
"0.80",
],
env=_DEEPEP_ENV,
)
test/registered/dp_attn/test_dp_attention_bcg_kl.py
+264
View File
@@ -0,0 +1,264 @@
from __future__ import annotations
import random
import unittest
import numpy as np
import requests
from sglang.srt.utils import get_device_capability, is_blackwell, kill_process_tree
from sglang.test.ci.ci_register import register_cuda_ci
from sglang.test.kl_test_utils import (
_extract_output_logprobs,
_flush_cache,
_generate,
_get_input_logprobs,
get_input_ids,
)
from sglang.test.test_utils import (
DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN,
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
DEFAULT_URL_FOR_TEST,
CustomTestCase,
popen_launch_server,
)
register_cuda_ci(est_time=160, stage="base-b", runner_config="2-gpu-large")
# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------
def _load_input_ids(tokenizer_path, num_samples, max_prompt_tokens):
try:
return get_input_ids(
tokenizer_path,
max_prompt_tokens=max_prompt_tokens,
num_samples=num_samples,
)
except (ValueError, OSError) as e:
print(
f"WARNING: Could not load LongBench inputs with tokenizer "
f"'{tokenizer_path}': {e}"
)
print("Falling back to random token IDs")
return [
[
random.randint(1, 32000 - 1)
for _ in range(int(max_prompt_tokens * random.uniform(0.5, 1.5)))
]
for _ in range(num_samples)
]
def _compute_kl(input_logprobs, output_logprobs):
kl_divs = []
for idx, (inp_lp, out_lp) in enumerate(zip(input_logprobs, output_logprobs)):
inp_none = any(v is None for v in inp_lp)
out_none = any(v is None for v in out_lp)
if inp_none or out_none:
src = "input" if inp_none else "output"
if inp_none and out_none:
src = "input and output"
print(f" WARNING: sample {idx}: skipping due to None in {src} logprobs")
continue
logr = np.array(inp_lp) - np.array(out_lp)
kl_divs.append(float(np.mean((np.exp(logr) - 1) - logr)))
avg = sum(kl_divs) / len(kl_divs)
print(f" per-sample KL: {kl_divs}")
print(f" avg KL: {avg:.6f}")
return avg
def _device_only_hit(meta_info):
details = meta_info.get("cached_tokens_details") or {}
if (details.get("host", 0) or 0) > 0:
return 0
return details.get("device", 0) or 0
# ---------------------------------------------------------------------------
# Hit detection helpers
# ---------------------------------------------------------------------------
def _is_prefill_hit(result, is_hicache):
if is_hicache:
return _device_only_hit(result["meta_info"]) > 0
return result["meta_info"]["cached_tokens"] > 0
def _is_decode_hit(result, first_turn_len, is_hicache):
if result["meta_info"]["cached_tokens"] <= first_turn_len + 1:
return False
if is_hicache:
return _device_only_hit(result["meta_info"]) > 0
return True
def _hit_info(result, is_hicache):
if is_hicache:
return f"device_only={_device_only_hit(result['meta_info'])}"
return f"cached_tokens={result['meta_info']['cached_tokens']}"
# ---------------------------------------------------------------------------
# Prefill / decode cache hit tests
# ---------------------------------------------------------------------------
def test_prefill_cache_hit(base_url, input_ids, max_new_tokens, is_hicache=False):
label = "device (L1) " if is_hicache else ""
print(f"--- Prefill {label}cache hit KL test ---")
_flush_cache(base_url)
_generate(base_url, input_ids, max_new_tokens=0)
results = _generate(base_url, input_ids, max_new_tokens, return_logprob=True)
new_input_ids, output_logprobs = [], []
for i, r in enumerate(results):
hit = _is_prefill_hit(r, is_hicache)
info = _hit_info(r, is_hicache)
print(f" [{i}] prefix_len={len(input_ids[i])} {info} hit={hit}")
if not hit:
continue
new_input_ids.append(input_ids[i] + r["output_ids"])
output_logprobs.append(_extract_output_logprobs(r))
hit_label = "L1 hits" if is_hicache else "cache hits"
print(f" {hit_label}: {len(new_input_ids)}/{len(input_ids)}")
assert (
len(new_input_ids) > len(input_ids) // 2
), f"too few {hit_label}: {len(new_input_ids)}/{len(input_ids)}"
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
return _compute_kl(input_logprobs, output_logprobs)
def test_decode_cache_hit(base_url, input_ids, max_new_tokens, is_hicache=False):
label = "device (L1) " if is_hicache else ""
print(f"--- Decode {label}cache hit KL test ---")
suffix_token = [1]
_flush_cache(base_url)
first = _generate(base_url, input_ids, max_new_tokens, return_logprob=True)
turn2_ids = [
input_ids[i] + r["output_ids"] + suffix_token for i, r in enumerate(first)
]
results = _generate(base_url, turn2_ids, max_new_tokens, return_logprob=True)
new_input_ids, output_logprobs = [], []
for i, r in enumerate(results):
hit = _is_decode_hit(r, len(input_ids[i]), is_hicache)
info = _hit_info(r, is_hicache)
print(f" [{i}] prefix_len={len(turn2_ids[i])} {info} hit={hit}")
if not hit:
continue
new_input_ids.append(turn2_ids[i] + r["output_ids"])
output_logprobs.append(_extract_output_logprobs(r))
hit_label = "L1 decode hits" if is_hicache else "cache hits"
print(f" {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}")
assert (
len(new_input_ids) > len(turn2_ids) // 2
), f"too few {hit_label}: {len(new_input_ids)}/{len(turn2_ids)}"
input_logprobs = _get_input_logprobs(base_url, new_input_ids, output_logprobs)
return _compute_kl(input_logprobs, output_logprobs)
# ---------------------------------------------------------------------------
# Server test
# ---------------------------------------------------------------------------
def _select_attention_backend():
major, minor = get_device_capability()
if major == 9:
return "fa3"
if is_blackwell():
return "fa4"
raise NotImplementedError(
f"DP attention BCG KL test only supports Hopper (fa3) and "
f"Blackwell (fa4); got compute capability {major}.{minor}"
)
class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase):
num_samples = 48
max_prompt_tokens = 1024
max_new_tokens = 256
@classmethod
def setUpClass(cls):
random.seed(42)
cls.model = DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN
cls.base_url = DEFAULT_URL_FOR_TEST
cls.attention_backend = _select_attention_backend()
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--trust-remote-code",
"--tp",
"2",
"--dp",
"2",
"--enable-dp-attention",
"--enable-deterministic-inference",
"--attention-backend",
cls.attention_backend,
"--moe-runner-backend",
"triton",
"--cuda-graph-backend-prefill=breakable",
"--chunked-prefill-size",
"2048",
"--mem-fraction-static",
"0.70",
],
)
server_info = requests.get(f"{cls.base_url}/server_info", timeout=30).json()
tokenizer_path = (
server_info.get("tokenizer_path")
or server_info.get("model_path")
or cls.model
)
cls.input_ids = _load_input_ids(
tokenizer_path, cls.num_samples, cls.max_prompt_tokens
)
print(f"Built {len(cls.input_ids)} prompts\n")
@classmethod
def tearDownClass(cls):
if hasattr(cls, "process") and cls.process:
kill_process_tree(cls.process.pid)
def test_prefill_and_decode_cache_hit_kl_is_zero(self):
server_info = requests.get(self.base_url + "/server_info", timeout=30).json()
self.assertFalse(server_info["disable_radix_cache"])
self.assertTrue(server_info["enable_dp_attention"])
self.assertTrue(server_info["enable_deterministic_inference"])
self.assertEqual(server_info["attention_backend"], self.attention_backend)
self.assertEqual(
server_info["cuda_graph_config"]["prefill"]["backend"], "breakable"
)
print("=== Radix Cache KL Divergence Eval ===")
print(f"Server: {self.base_url} Samples: {self.num_samples}\n")
prefill_kl = test_prefill_cache_hit(
self.base_url, self.input_ids, self.max_new_tokens
)
decode_kl = test_decode_cache_hit(
self.base_url, self.input_ids, self.max_new_tokens
)
self.assertEqual(prefill_kl, 0.0)
self.assertEqual(decode_kl, 0.0)
if __name__ == "__main__":
unittest.main()
test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py
+44
View File
@@ -1058,6 +1058,50 @@ class TestBuildPrefillRegistry(unittest.TestCase):
self.assertTrue(torch.all(ids[3:8] == 0)) # padded tail reset
self.assertTrue(torch.all(ids[8:] == 7)) # beyond the bucket: untouched
def test_num_token_non_padded_scalar_copy(self):
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
build_prefill_registry,
)
src = self._src(num_token_non_padded=torch.zeros((1,), dtype=torch.int32))
reg = build_prefill_registry(
device=torch.device("cpu"),
max_bs=1,
max_num_token=16,
cache_loc_dtype=torch.int64,
enable_num_token_non_padded=True,
source=src,
)
self.assertTrue(reg.has_slot("num_token_non_padded"))
self.assertEqual(
reg.get_slot("num_token_non_padded").buffer.data_ptr(),
src.num_token_non_padded.data_ptr(),
)
fb = _MiniForwardBatch(
input_ids=torch.tensor([1, 2, 3], dtype=torch.int64),
positions=torch.tensor([4, 5, 6], dtype=torch.int64),
out_cache_loc=torch.tensor([8, 9, 10], dtype=torch.int64),
num_token_non_padded=torch.tensor([3], dtype=torch.int32),
)
reg.fill_from(fb, raw_bs=1, padded_bs=1, raw_num_tokens=3, padded_num_tokens=8)
self.assertTrue(
torch.equal(
reg.get_slot("num_token_non_padded").buffer,
torch.tensor([3], dtype=torch.int32),
)
)
static_fb = reg.extract_buffer(
padded_bs=1,
padded_num_tokens=8,
forward_batch_template=fb,
)
self.assertEqual(
static_fb.num_token_non_padded.data_ptr(),
src.num_token_non_padded.data_ptr(),
)
def test_multimodal_input_embeds_reset_only(self):
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
build_prefill_registry,