From b94ac87e0c411e3c1b120b3d13bba9092b198fe0 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Sun, 12 Jul 2026 17:10:04 -0700 Subject: [PATCH] Enable breakable prefill CUDA graph for DP attention (#30898) --- .../managers/scheduler_components/dp_attn.py | 18 +- .../cuda_graph_buffer_registry.py | 10 + .../srt/model_executor/forward_batch_info.py | 40 ++- .../runner/prefill_cuda_graph_runner.py | 78 +++++- .../model_executor/runner_utils/buffers.py | 4 +- python/sglang/srt/server_args.py | 38 ++- .../cp/test_deepseek_v4_flash_fp4_b200_cp.py | 2 + .../dp_attn/test_dp_attention_bcg_kl.py | 264 ++++++++++++++++++ .../test_cuda_graph_buffer_registry.py | 44 +++ 9 files changed, 478 insertions(+), 20 deletions(-) create mode 100644 test/registered/dp_attn/test_dp_attention_bcg_kl.py diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index b473a11ee..a4efe8218 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -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: diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index a90b9bef9..b62b822ca 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py @@ -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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 1df55a02f..51a562361 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -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) diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index d658b5c59..8fe294763 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -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): diff --git a/python/sglang/srt/model_executor/runner_utils/buffers.py b/python/sglang/srt/model_executor/runner_utils/buffers.py index 00d32a249..6dc963082 100644 --- a/python/sglang/srt/model_executor/runner_utils/buffers.py +++ b/python/sglang/srt/model_executor/runner_utils/buffers.py @@ -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, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 437073a64..a79fd28af 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -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: diff --git a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py index 815f36aa4..6f79cdb4b 100644 --- a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py +++ b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py @@ -80,6 +80,8 @@ class TestDSV4FlashFP4B200Balanced_CP( "round-robin-split", "--deepep-config", DEEPEP_CONFIG, + "--mem-fraction-static", + "0.80", ], env=_DEEPEP_ENV, ) diff --git a/test/registered/dp_attn/test_dp_attention_bcg_kl.py b/test/registered/dp_attn/test_dp_attention_bcg_kl.py new file mode 100644 index 000000000..5445343ba --- /dev/null +++ b/test/registered/dp_attn/test_dp_attention_bcg_kl.py @@ -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() diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index c8e35ca19..18d964f41 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py @@ -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,