From 58ecbba0bd8c2aec9adc91dbc66e125e8c211028 Mon Sep 17 00:00:00 2001 From: Aurick Qiao Date: Wed, 26 Aug 2026 02:16:02 -0700 Subject: [PATCH] [Feature] Coordinate FullCG prefill across DP-attention ranks (#35640) Co-authored-by: Yuwei An --- .../srt/batch_overlap/two_batch_overlap.py | 4 +- python/sglang/srt/managers/schedule_batch.py | 8 +- .../managers/scheduler_components/dp_attn.py | 15 +-- .../srt/model_executor/forward_batch_info.py | 23 +++-- .../runner/decode_cuda_graph_runner.py | 6 +- .../runner/prefill_cuda_graph_runner.py | 36 ++++---- .../dspark_components/dspark_draft.py | 2 +- .../eagle_draft_cuda_graph_runner.py | 4 +- .../eagle_draft_extend_cuda_graph_runner.py | 4 +- .../frozen_kv_mtp_cuda_graph_runner.py | 4 +- .../speculative/frozen_kv_mtp_worker_v2.py | 2 +- ...er_eagle_draft_extend_cuda_graph_runner.py | 4 +- .../dp_attn/test_dp_attention_bcg_kl.py | 92 ++++++++++++++++++- .../spec/dspark/test_dspark_dp_tier.py | 4 +- .../model_executor/test_mlp_sync_pad_unpad.py | 4 +- 15 files changed, 155 insertions(+), 57 deletions(-) diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index f26955de4..7ca3d0dbd 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -736,8 +736,8 @@ class TboForwardBatchPreparer: "forward_mode", "is_extend_in_batch", "return_logprob", - "can_run_dp_cuda_graph", - "can_run_dp_breakable_cuda_graph", + "can_run_decode_cuda_graph", + "can_run_dp_prefill_cuda_graph", "dp_padding_mode", "global_forward_mode", "is_prefill_only", diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 22b8e1dde..41513a1c4 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2161,8 +2161,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): # For DP attention is_extend_in_batch: bool = False - can_run_dp_cuda_graph: bool = False - can_run_dp_breakable_cuda_graph: bool = False + can_run_decode_cuda_graph: bool = False + can_run_dp_prefill_cuda_graph: bool = False tbo_split_seq_index: Optional[int] = None # Rank-consistent forward mode for the recv skipper, derived from the MLP # sync all-gather (the TBO-only `global_forward_mode` is None without TBO). @@ -3355,8 +3355,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): spec_info=self.spec_info, global_num_tokens=self.global_num_tokens, global_num_tokens_for_logprob=self.global_num_tokens_for_logprob, - can_run_dp_cuda_graph=self.can_run_dp_cuda_graph, - can_run_dp_breakable_cuda_graph=self.can_run_dp_breakable_cuda_graph, + can_run_decode_cuda_graph=self.can_run_decode_cuda_graph, + can_run_dp_prefill_cuda_graph=self.can_run_dp_prefill_cuda_graph, is_extend_in_batch=self.is_extend_in_batch, is_prefill_only=self.is_prefill_only, seq_lens_cpu=self.seq_lens_cpu, diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 56274517e..2b1b0e28a 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -221,8 +221,8 @@ def _update_gather_batch( batch.global_forward_mode = mlp_sync_info.global_forward_mode # Check forward mode for cuda graph - batch.can_run_dp_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph - batch.can_run_dp_breakable_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph + batch.can_run_decode_cuda_graph = mlp_sync_info.can_run_decode_cuda_graph + batch.can_run_dp_prefill_cuda_graph = mlp_sync_info.can_run_prefill_cuda_graph def prepare_mlp_sync_batch_raw( @@ -271,14 +271,15 @@ 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 - breakable_prefill = check_cuda_graph_backend(Phase.PREFILL, Backend.BREAKABLE) + coordinated_prefill = check_cuda_graph_backend( + Phase.PREFILL, Backend.BREAKABLE + ) or check_cuda_graph_backend(Phase.PREFILL, Backend.FULL) prefill_graph_runner = ( - model_runner.prefill_cuda_graph_runner if breakable_prefill else None + model_runner.prefill_cuda_graph_runner if coordinated_prefill else None ) can_run_prefill_cuda_graph = ( local_batch is None or local_batch.forward_mode.is_idle() - # Breakable Cuda Graph Backend Check. or ( local_batch.forward_mode in (ForwardMode.EXTEND, ForwardMode.MIXED) and ( @@ -287,7 +288,7 @@ def prepare_mlp_sync_batch_raw( batch_size=local_batch.batch_size(), num_tokens=local_batch.extend_num_tokens, input_embeds=local_batch.input_embeds, - replace_embeds=None, + replace_embeds=local_batch.replace_embeds, prefix_lens=local_batch.prefix_lens, is_target_verify=local_batch.forward_mode.is_target_verify(), capture_hidden_mode=None, @@ -295,7 +296,7 @@ def prepare_mlp_sync_batch_raw( lora_ineligible=prefill_graph_runner.enable_lora, ) ) - and breakable_prefill + and coordinated_prefill ) ) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 79c75d196..b92d5ce07 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -446,8 +446,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # For DP attention is_extend_in_batch: bool = False - can_run_dp_cuda_graph: bool = False - can_run_dp_breakable_cuda_graph: bool = False + can_run_decode_cuda_graph: bool = False + can_run_dp_prefill_cuda_graph: bool = False global_forward_mode: Optional[ForwardMode] = None # For two-batch overlap @@ -703,7 +703,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): self.global_num_tokens_for_logprob_gpu = torch.tensor( global_num_tokens_for_logprob, dtype=torch.int64 ).to(device, non_blocking=True) - self.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph + self.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph @classmethod def init_new( @@ -788,8 +788,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): # Scalar config / flags return_logprob=batch.return_logprob, is_extend_in_batch=batch.is_extend_in_batch, - can_run_dp_cuda_graph=batch.can_run_dp_cuda_graph, - can_run_dp_breakable_cuda_graph=batch.can_run_dp_breakable_cuda_graph, + can_run_decode_cuda_graph=batch.can_run_decode_cuda_graph, + can_run_dp_prefill_cuda_graph=batch.can_run_dp_prefill_cuda_graph, global_forward_mode=batch.global_forward_mode, is_prefill_only=batch.is_prefill_only, spec_algorithm=batch.spec_algorithm, @@ -1312,21 +1312,20 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): ): # Joined ranks require real token counts instead of MAX_LEN padding. dp_padding_mode = DpPaddingMode.SUM_LEN - # 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 + # Prefill CUDA graphs require 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. This mismatches the # rank-coupled communication geometry: DP gather/combine uses # all_gather_into_tensor / reduce_scatter_tensor, while MoE backends may # use A2A dispatch/combine. Force MAX_LEN so every rank pads to the global # max and picks the same bucket. # - # 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. + # 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 = get_exec().graph.cuda_graph_config.prefill if ( - self.can_run_dp_breakable_cuda_graph + self.can_run_dp_prefill_cuda_graph and self.is_extend_in_batch and prefill_cg.bs and max(global_num_tokens) <= max(prefill_cg.bs) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 204b17691..934be8494 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -694,7 +694,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): ) if self.require_mlp_sync: - is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph + is_bs_supported = ( + is_bs_supported and forward_batch.can_run_decode_cuda_graph + ) # NOTE: cuda graph cannot handle mixed batch (encoder_len = 0) # If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph @@ -735,7 +737,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): ] and forward_batch.batch_size <= self._ragged_capture_slots(admission_tokens) is_dp_supported = ( - forward_batch.can_run_dp_cuda_graph if self.require_mlp_sync else True + forward_batch.can_run_decode_cuda_graph if self.require_mlp_sync else True ) is_encoder_lens_supported = ( 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 6e38bd5ad..e703d4e1e 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 @@ -40,6 +40,7 @@ from __future__ import annotations import copy import inspect import logging +from collections.abc import Sequence from contextlib import contextmanager from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, Optional, Union @@ -826,19 +827,29 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): return prefix_chunk_len, prefix_chunk_len * capture_req_slots def _select_prefix_capture_chunks( - self, forward_batch: ForwardBatch + self, prefix_lens: Sequence[int] ) -> Optional[int]: """Smallest captured variant covering the batch's max prefix, or None.""" - max_prefix_len = max( - int(length) for length in forward_batch.extend_prefix_lens_cpu - ) + max_prefix_len = max(int(length) for length in prefix_lens) real_n = _ceil_div(max_prefix_len, self._prefix_chunk_len) return next((n for n in self._prefix_capture_variants if n >= real_n), None) + def _has_uncapturable_chunked_prefix( + self, prefix_lens: Sequence[int] | None + ) -> bool: + return ( + self._capture_chunked_prefix + and prefix_lens is not None + and any(int(length) > 0 for length in prefix_lens) + and self._select_prefix_capture_chunks(prefix_lens) is None + ) + def _shape_key(self, num_tokens: int, forward_batch: ForwardBatch) -> ShapeKey: variant = None if self._capture_chunked_prefix and self._has_prefix_hit(forward_batch): - captured_n = self._select_prefix_capture_chunks(forward_batch) + captured_n = self._select_prefix_capture_chunks( + forward_batch.extend_prefix_lens_cpu + ) assert captured_n is not None, "prefix batch has no captured FullCG variant" variant = _chunked_prefix_variant(captured_n) return ShapeKey(size=num_tokens, variant_label=variant) @@ -1055,7 +1066,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): capture_hidden_mode, return_logprob: bool, lora_ineligible: bool = False, - chunked_prefix_uncapturable: bool = False, ) -> bool: """Rank-local replay eligibility: the single source of truth for ``can_run_graph`` (ForwardBatch, forward time) and the dp mlp-sync @@ -1084,10 +1094,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): and any(prefix_lens) ): return False - # FullCG's chunked-prefix topology covers a bounded prefix. The flag - # gating it is FULL-backend-only, so this is inert for the breakable - # vote path. - if chunked_prefix_uncapturable: + # FullCG's chunked-prefix topology covers a bounded prefix. Its capture + # flag is FullCG-only, so this is inert for the BreakableCG vote path. + if self._has_uncapturable_chunked_prefix(prefix_lens): return False # tc_piecewise captures with ForwardMode.EXTEND and spec_info=None. if is_target_verify: @@ -1115,7 +1124,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): # (min-reduced votes; also requires every rank to hold tokens). if ( forward_batch.global_num_tokens_cpu is not None - and not forward_batch.can_run_dp_breakable_cuda_graph + and not forward_batch.can_run_dp_prefill_cuda_graph ): return False @@ -1141,11 +1150,6 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): forward_batch ) ), - chunked_prefix_uncapturable=( - self._capture_chunked_prefix - and self._has_prefix_hit(forward_batch) - and self._select_prefix_capture_chunks(forward_batch) is None - ), ): return False if getattr(self, "enable_cp_v2_bcg_capture", False) and is_cp_v2_active( diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index 4c1983a1b..f5d5b3335 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -444,7 +444,7 @@ class DraftBlockProposer: ) -> None: # The dense DSpark draft still reuses the target batch's graph tier. # Set graph eligibility before the DP-MoE-only metadata early return. - forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph + forward_batch.can_run_decode_cuda_graph = batch.can_run_decode_cuda_graph if not self._dp_moe_sync or batch.global_num_tokens is None: return # Graph bucket selection uses the raw per-rank request counts. Keep diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index 87a7011fb..5ac5189ed 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -318,7 +318,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner): ) if self.require_mlp_sync: - is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph + is_bs_supported = ( + is_bs_supported and forward_batch.can_run_decode_cuda_graph + ) return is_bs_supported diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 642c0954f..deede532a 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -309,7 +309,9 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ) if self.require_mlp_sync: - is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph + is_bs_supported = ( + is_bs_supported and forward_batch.can_run_decode_cuda_graph + ) return is_bs_supported diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py index 4317e5629..b86b8737d 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_cuda_graph_runner.py @@ -237,7 +237,9 @@ class FrozenKVMTPCudaGraphRunner(DecodeCudaGraphRunner): else cuda_graph_bs <= self.max_bs ) if self.require_mlp_sync: - is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph + is_bs_supported = ( + is_bs_supported and forward_batch.can_run_decode_cuda_graph + ) return is_bs_supported def capture_one_shape( diff --git a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py index e36995737..443ed9ed6 100644 --- a/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py +++ b/python/sglang/srt/speculative/frozen_kv_mtp_worker_v2.py @@ -485,7 +485,7 @@ class FrozenKVMTPDraftWorker(EagleDraftWorkerBase, TpModelWorker): self.cuda_graph_runner.execute(forward_batch) ) else: - forward_batch.can_run_dp_cuda_graph = False + forward_batch.can_run_decode_cuda_graph = False parent_list, top_scores_index, draft_tokens = self.draft_forward( forward_batch ) diff --git a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py index 63d8853ec..359bdfa4f 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py @@ -255,7 +255,9 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): ) if self.require_mlp_sync: - is_bs_supported = is_bs_supported and forward_batch.can_run_dp_cuda_graph + is_bs_supported = ( + is_bs_supported and forward_batch.can_run_decode_cuda_graph + ) return is_bs_supported diff --git a/test/registered/dp_attn/test_dp_attention_bcg_kl.py b/test/registered/dp_attn/test_dp_attention_bcg_kl.py index 5445343ba..272783279 100644 --- a/test/registered/dp_attn/test_dp_attention_bcg_kl.py +++ b/test/registered/dp_attn/test_dp_attention_bcg_kl.py @@ -1,6 +1,11 @@ from __future__ import annotations +import os import random +import re +import shutil +import tempfile +import time import unittest import numpy as np @@ -23,7 +28,12 @@ from sglang.test.test_utils import ( popen_launch_server, ) -register_cuda_ci(est_time=160, stage="base-b", runner_config="2-gpu-large") +register_cuda_ci(est_time=320, stage="base-b", runner_config="2-gpu-large") + +PREFILL_GRAPH_REPLAY_PATTERN = re.compile(r"Prefill batch.*cuda graph: True") +CACHED_PREFIX_GRAPH_REPLAY_PATTERN = re.compile( + r"Prefill batch.*#cached-token: [1-9][0-9]*.*cuda graph: True" +) # --------------------------------------------------------------------------- @@ -185,10 +195,11 @@ def _select_attention_backend(): ) -class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase): +class _DPAttentionPrefillCudaGraphKLMixin: num_samples = 48 max_prompt_tokens = 1024 max_new_tokens = 256 + prefill_backend: str @classmethod def setUpClass(cls): @@ -196,6 +207,11 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase): cls.model = DEFAULT_TARGET_MODEL_EAGLE_DP_ATTN cls.base_url = DEFAULT_URL_FOR_TEST cls.attention_backend = _select_attention_backend() + cls.log_dir = tempfile.mkdtemp(prefix=f"dp_attn_{cls.prefill_backend}_") + cls.stdout_path = os.path.join(cls.log_dir, "server.out") + cls.stderr_path = os.path.join(cls.log_dir, "server.err") + cls.stdout = open(cls.stdout_path, "w") + cls.stderr = open(cls.stderr_path, "w") cls.process = popen_launch_server( cls.model, cls.base_url, @@ -212,12 +228,15 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase): cls.attention_backend, "--moe-runner-backend", "triton", - "--cuda-graph-backend-prefill=breakable", + f"--cuda-graph-backend-prefill={cls.prefill_backend}", "--chunked-prefill-size", "2048", + "--prefill-max-requests", + "2", "--mem-fraction-static", "0.70", ], + return_stdout_stderr=(cls.stdout, cls.stderr), ) server_info = requests.get(f"{cls.base_url}/server_info", timeout=30).json() @@ -235,6 +254,36 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase): def tearDownClass(cls): if hasattr(cls, "process") and cls.process: kill_process_tree(cls.process.pid) + for attr in ("stdout", "stderr"): + output = getattr(cls, attr, None) + if output is not None and not output.closed: + output.close() + if hasattr(cls, "log_dir"): + shutil.rmtree(cls.log_dir, ignore_errors=True) + + def _wait_for_prefill_graph_replay( + self, + offsets, + pattern=PREFILL_GRAPH_REPLAY_PATTERN, + case="a lone request", + ): + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + chunks = [] + for path, offset in zip( + (self.stdout_path, self.stderr_path), offsets, strict=True + ): + with open(path, "rb") as log: + log.seek(offset) + chunks.append(log.read().decode(errors="replace")) + logs = "\n".join(chunks) + if pattern.search(logs): + return + time.sleep(0.5) + self.fail( + f"No {self.prefill_backend} prefill CUDA graph replay was logged " + f"for {case}" + ) def test_prefill_and_decode_cache_hit_kl_is_zero(self): server_info = requests.get(self.base_url + "/server_info", timeout=30).json() @@ -243,15 +292,24 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase): 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" + server_info["cuda_graph_config"]["prefill"]["backend"], + self.prefill_backend, ) print("=== Radix Cache KL Divergence Eval ===") print(f"Server: {self.base_url} Samples: {self.num_samples}\n") + offsets = [ + os.path.getsize(path) for path in (self.stdout_path, self.stderr_path) + ] prefill_kl = test_prefill_cache_hit( self.base_url, self.input_ids, self.max_new_tokens ) + self._wait_for_prefill_graph_replay( + offsets, + CACHED_PREFIX_GRAPH_REPLAY_PATTERN, + "a cached-prefix request", + ) decode_kl = test_decode_cache_hit( self.base_url, self.input_ids, self.max_new_tokens ) @@ -259,6 +317,32 @@ class TestDPAttentionBreakablePrefillCudaGraphKL(CustomTestCase): self.assertEqual(prefill_kl, 0.0) self.assertEqual(decode_kl, 0.0) + def test_lone_request_replays_prefill_cuda_graph(self): + _flush_cache(self.base_url) + offsets = [ + os.path.getsize(path) for path in (self.stdout_path, self.stderr_path) + ] + result = _generate( + self.base_url, + self.input_ids[0], + max_new_tokens=1, + return_logprob=True, + ) + self.assertNotIn("error", result) + self._wait_for_prefill_graph_replay(offsets) + + +class TestDPAttentionBreakablePrefillCudaGraphKL( + _DPAttentionPrefillCudaGraphKLMixin, CustomTestCase +): + prefill_backend = "breakable" + + +class TestDPAttentionFullPrefillCudaGraphKL( + _DPAttentionPrefillCudaGraphKLMixin, CustomTestCase +): + prefill_backend = "full" + if __name__ == "__main__": unittest.main() diff --git a/test/registered/spec/dspark/test_dspark_dp_tier.py b/test/registered/spec/dspark/test_dspark_dp_tier.py index 169531224..b3278a15b 100644 --- a/test/registered/spec/dspark/test_dspark_dp_tier.py +++ b/test/registered/spec/dspark/test_dspark_dp_tier.py @@ -67,7 +67,7 @@ class TestDraftDpSyncMetadata(CustomTestCase): batch = SimpleNamespace( global_num_tokens=[1, 3, 0, 2], global_num_tokens_for_logprob=[1, 3, 0, 2], - can_run_dp_cuda_graph=True, + can_run_decode_cuda_graph=True, ) with patch( @@ -84,7 +84,7 @@ class TestDraftDpSyncMetadata(CustomTestCase): self.assertEqual(forward_batch.num_token_non_padded.item(), 6) self.assertEqual(forward_batch.num_token_non_padded.dtype, torch.int32) self.assertEqual(forward_batch.num_token_non_padded_cpu, 6) - self.assertTrue(forward_batch.can_run_dp_cuda_graph) + self.assertTrue(forward_batch.can_run_decode_cuda_graph) class TestBusyIdleGraphKeyIdentity(CustomTestCase): diff --git a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py index 1c97b2e9d..e51c5c7bc 100644 --- a/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py +++ b/test/registered/unit/model_executor/test_mlp_sync_pad_unpad.py @@ -59,7 +59,7 @@ class TestMlpSyncPadUnpad(CustomTestCase): batch = SimpleNamespace( global_num_tokens=[2, 0, 3], global_num_tokens_for_logprob=[2, 0, 3], - can_run_dp_cuda_graph=True, + can_run_decode_cuda_graph=True, ) fb.init_mlp_sync_metadata(batch, torch.device("cpu")) @@ -71,7 +71,7 @@ class TestMlpSyncPadUnpad(CustomTestCase): torch.testing.assert_close( fb.global_num_tokens_for_logprob_gpu, torch.tensor([4, 0, 6]) ) - self.assertTrue(fb.can_run_dp_cuda_graph) + self.assertTrue(fb.can_run_decode_cuda_graph) def test_draft_input_without_hidden_states_can_be_padded(self): spec_info = SimpleNamespace(