diff --git a/docs_new/docs/advanced_features/speculative_decoding.mdx b/docs_new/docs/advanced_features/speculative_decoding.mdx index 52c89027bb4e..789950544a46 100644 --- a/docs_new/docs/advanced_features/speculative_decoding.mdx +++ b/docs_new/docs/advanced_features/speculative_decoding.mdx @@ -1,9 +1,9 @@ --- title: "Speculative Decoding" metatags: - description: "SGLang speculative decoding: EAGLE-2/EAGLE-3, MTP, DFLASH, draft model configuration, and overlap-scheduler guidance." + description: "SGLang speculative decoding: EAGLE-2/EAGLE-3, MTP, DFLASH, DSPARK, draft model configuration, and overlap-scheduler guidance." --- -SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, MTP, DFLASH, classic draft-model decoding, and an NGRAM-based variant. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines. +SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, MTP, DFLASH, DSPARK, classic draft-model decoding, and an NGRAM-based variant. ## Summary @@ -16,6 +16,7 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, - [EAGLE-3 Decoding](#eagle-3-decoding) - [Multi Token Prediction](#multi-token-prediction) - [DFlash Decoding](#dflash-decoding) +- [DSpark Decoding](#dspark-decoding) - [Standalone Speculative Decoding (Small Draft Model)](#standalone-speculative-decoding-small-draft-model) - [Speculative Decoding V2 (Overlap Scheduler)](#speculative-decoding-v2-overlap-scheduler) - [Ngram Speculative Decoding](#ngram-speculative-decoding) @@ -31,6 +32,7 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, - **Lower `lm_head` overhead for EAGLE-2**: Enable **FR-Spec** with `--speculative-token-map`. - **Model is MTP-enabled**: Use **MTP via speculative decoding** (often with small `speculative_num_steps/topk/num_draft_tokens`, see the example section). - **You have a DFlash draft checkpoint**: Use **DFLASH** with `--speculative-algorithm DFLASH` and `--speculative-draft-model-path ...`. +- **You have a DSpark checkpoint**: Use **DSPARK** with `--speculative-algorithm DSPARK`. A bundled DSpark checkpoint does not need a separate draft path. - **You have a smaller draft LLM**: Use **STANDALONE** (`--speculative-algorithm STANDALONE`). - **No extra model available**: Use **NGRAM** (`--speculative-algorithm NGRAM`, CUDA-only). @@ -93,6 +95,13 @@ SGLang provides several speculative decoding options, including EAGLE-2/EAGLE-3, --speculative-algorithm DFLASH + --speculative-draft-model-path ... No --enable-dp-attention; pp_size == 1; disables overlap scheduler & mixed chunked prefill + + DSPARK + Confidence-driven variable-length block draft + Checkpoint-dependent + --speculative-algorithm DSPARK + The target checkpoint can bundle the DSpark draft head; structured output requires the batched XGrammar API + STANDALONE Smaller draft LLM (token-level) @@ -498,6 +507,39 @@ print(response.choices[0].message.content) --- +## DSpark Decoding + +DSpark proposes a fixed-size token block and uses confidence-driven, variable-length target verification. DeepSeek V4 DSpark checkpoints can bundle the draft head into the target checkpoint, so you omit `--speculative-draft-model-path`. + +### Prerequisites + +- NVIDIA CUDA +- A DSpark-compatible checkpoint such as `deepseek-ai/DeepSeek-V4-Flash-DSpark` +- `--speculative-eagle-topk 1` +- For constrained decoding, an XGrammar build that provides `BatchGrammarMatcher.batch_traverse_draft_tree` + +The draft checkpoint records its block size (`gamma`). `--speculative-num-draft-tokens` is the verify-window size and must equal `gamma + 1`. Omit it to read the value from the checkpoint, or set it explicitly to make checkpoint changes fail during startup. + +```bash Command +python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-V4-Flash-DSpark \ + --served-model-name deepseek-ai/DeepSeek-V4-Flash \ + --tp-size 2 \ + --dp-size 2 \ + --enable-dp-attention \ + --speculative-algorithm DSPARK \ + --speculative-num-steps 1 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 6 \ + --enable-metrics +``` + +For a co-located DP2 deployment (`dp_size=2`, attention `tp_size=1`, `cp_size=1`), you can enable the native shared-memory scheduler exchange described in the [environment-variable reference](/docs/references/environment_variables#scheduler--batching). The path is fail-closed: startup rejects incompatible geometry, missing native support, invalid settings, or a synchronization timeout. + +With `--grammar-backend xgrammar`, DSpark batches grammar draft-tree traversal on CPU, overlaps the resulting mask transfer with target verification, and records cache and pipeline metrics. The server rejects this configuration at startup when the required XGrammar API is unavailable. + +--- + ## Standalone Speculative Decoding (Small Draft Model) Besides EAGLE/MTP, SGLang also supports **token-level speculative decoding** using a smaller **draft model**. Enable it with `--speculative-algorithm STANDALONE` and provide a draft model via `--speculative-draft-model-path`. @@ -746,7 +788,7 @@ Below is a comprehensive list of all speculative decoding parameters available i --speculative-algorithm str None - Algorithm to use: DFLASH, EAGLE, EAGLE3, STANDALONE, NGRAM, NEXTN (alias of EAGLE) + Algorithm to use: DFLASH, DSPARK, EAGLE, EAGLE3, STANDALONE, NGRAM, NEXTN (alias of EAGLE) --speculative-draft-model-path diff --git a/docs_new/docs/basic_usage/native_api.mdx b/docs_new/docs/basic_usage/native_api.mdx index 42c5bd228319..5b6e3dc96b26 100644 --- a/docs_new/docs/basic_usage/native_api.mdx +++ b/docs_new/docs/basic_usage/native_api.mdx @@ -9,6 +9,7 @@ Apart from the OpenAI compatible APIs, the SGLang Runtime also provides its nati - `/get_model_info` - `/server_info` - `/health` +- `/health_scheduler` - `/health_generate` - `/flush_cache` - `/update_weights` @@ -104,6 +105,7 @@ print_highlight(response.text) ## Health Check - `/health`: Check the health of the server. +- `/health_scheduler`: Check that every DP scheduler is present and publishing a fresh load snapshot without running inference. It returns `503` during startup or shutdown, when a rank is missing or duplicated, when a snapshot is stale, or when scheduler state cannot be read. - `/health_generate`: Check the health of the server by generating one token. ```python Example @@ -113,6 +115,13 @@ response = requests.get(url) print_highlight(response.text) ``` +```python Example +url = f"http://localhost:{port}/health_scheduler" + +response = requests.get(url) +print_highlight(response.text) +``` + ```python Example url = f"http://localhost:{port}/health" diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 67750563c55b..1025b3e9bfec 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -1127,6 +1127,11 @@ SGLang supports various environment variables that can be used to configure its Decode-side incremental KV cache offload stride. Rounded down to a multiple of --page-size (min is --page-size). If unset/invalid/<=0, it falls back to --page-size. Not set (uses --page-size) + + SGLANG_HICACHE_GATE_CONTROL + Arm a one-shot, idle-only HiCache device eviction through /set_internal_state. Accepted values are disabled and armed; an armed request disarms after one attempt. + disabled + SGLANG_HICACHE_NIXL_USE_DIRECT_IO Enable O_DIRECT for any file-based NIXL backend (POSIX, GDS, GDS_MT, 3FS) when opening cache files (bypasses the OS page cache, reducing memory pressure and improving throughput on NVMe). Can also be disabled via {'{"use_direct_io": false}'} in --hicache-storage-backend-extra-config. Falls back to buffered I/O with a warning when O_DIRECT is unavailable on the current OS. @@ -1280,6 +1285,61 @@ SGLang supports various environment variables that can be used to configure its Maximum poll iterations before grammar compilation is treated as stuck. 10000 + + SGLANG_GRAMMAR_COMPILATION_WORKERS + Thread-pool workers used for asynchronous grammar compilation. Must be positive. + 8 + + + SGLANG_GRAMMAR_CACHE_BYTES + Hard byte limit for the in-process compiled-grammar LRU. + 10737418240 (10 GiB) + + + SGLANG_GRAMMAR_CACHE_ENTRIES + Hard entry limit for the in-process compiled-grammar LRU. + 100000 + + + SGLANG_XGRAMMAR_COMPILER_THREADS + Native XGrammar compiler threads. Must be positive. + 12 + + + SGLANG_XGRAMMAR_CACHE_DIR + Root directory for the tokenizer- and compiler-scoped persistent XGrammar cache. + ~/.cache/sglang/xgrammar + + + SGLANG_XGRAMMAR_DISK_CACHE_BYTES + Hard byte limit for the persistent XGrammar cache. + 42949672960 (40 GiB) + + + SGLANG_XGRAMMAR_DESERIALIZE_BYTES_PER_SECOND + Estimated deserialization throughput used to choose between reading a cached artifact and compiling locally. + 134217728 (128 MiB/s) + + + SGLANG_XGRAMMAR_LOCAL_COMPILE_SPEEDUP + Safety factor required before adaptive cache policy selects local compilation over deserialization. + 2 + + + SGLANG_XGRAMMAR_CACHE_SESSION_ID + Optional ASCII launch identifier. A new value reconciles the persistent-cache size ledger after an unclean exit. + Not set + + + SGLANG_GRAMMAR_TRAVERSAL_THREADS + Native CPU threads for each DSpark batched draft-tree traversal. + 12 + + + SGLANG_DSPARK_GRAMMAR_METRICS_INTERVAL + DSpark grammar steps between aggregated metric emissions. + 16 + @@ -1359,6 +1419,31 @@ SGLang supports various environment variables that can be used to configure its Skip the scheduler all-gather step. false + + SGLANG_DSPARK_DP2_SHM_MLP_SYNC + Enable the native shared-memory/futex scheduler exchange for co-located DSpark DP2. The runtime rejects any geometry other than DP2, attention TP1, and CP1. + 0 + + + SGLANG_DSPARK_DP2_SHM_SESSION_ID + Required shared-memory namespace when native DP2 synchronization is enabled. Use a new value for each server launch. + Not set + + + SGLANG_DSPARK_DP2_SHM_TIMEOUT_MS + Positive peer-wait timeout for native DP2 synchronization. A timeout is fatal; the runtime does not select another transport. + 30000 + + + SGLANG_DSPARK_DP2_SHM_METRICS + Emit native DP2 synchronization metrics. Must be 1 when native synchronization is enabled. + 1 + + + SGLANG_DSPARK_DP2_SHM_LIBRARY + Optional path override for the native DP2 library. When unset, SGLang loads sglang_dp2_sync.so from the installed sgl_kernel package. + Not set + SGLANG_ENABLE_WAR_BARRIER Force-enable the write-after-read barrier for the overlap scheduler even when CUDA is not detected (e.g. AMD/ROCm). On CUDA the barrier is always enabled. @@ -1611,6 +1696,11 @@ SGLang supports various environment variables that can be used to configure its Skip the sgl-kernel version compatibility check. false + + SGLANG_PREFILL_CUDA_GRAPH_EAGER_VALIDATION + During full-prefill CUDA graph capture, compare the captured output with an eager reference and fail startup on a numerical mismatch. + false + diff --git a/docs_new/docs/references/production_metrics.mdx b/docs_new/docs/references/production_metrics.mdx index f15f51c9f67c..5e6e8cbdac30 100644 --- a/docs_new/docs/references/production_metrics.mdx +++ b/docs_new/docs/references/production_metrics.mdx @@ -274,3 +274,25 @@ Average estimated memory bandwidth in GB/s: - These metrics are estimates intended for observability and trend analysis. - Estimated memory bytes reflect modeled traffic and are not a direct hardware counter from GPU profilers. + +## Scheduler, prefill graph, grammar, and HiCache metrics + +The runtime exports additional per-rank metrics when `--enable-metrics` is set: + +| Metric family | What it measures | +| --- | --- | +| `sglang:scheduler_phase_*` | Scheduler-thread time, calls, and one-second maximum by phase | +| `sglang:runtime_gc_frozen` | Whether the warmed scheduler object graph moved to Python's permanent GC generation | +| `sglang:prefill_graph_admissions_total` | Full-prefill CUDA graph admissions and eager-routing reasons | +| `sglang:prefill_graph_shapes_total` | Captured token bucket and request-slot shape used by each graph replay | +| `sglang:prefill_execution_tokens_total` | Scheduled versus executed token-axis elements, including graph padding | +| `sglang:prefill_batch_scheduled_tokens` and `sglang:prefill_batch_requests` | Prefill batch shape by graph or eager path | +| `sglang:grammar_cache_lookup_total` | Grammar resolutions from memory, an in-flight compile, disk, local compile, or a new compile | +| `sglang:grammar_cache_*_time_seconds` | End-to-end cache latency, persistent-cache lock wait, and internal cache phases | +| `sglang:dspark_grammar_*` | DSpark structured-output batch size, active matchers, mask outcomes, and overlapped pipeline phases | +| `sglang:hicache_backup_*` | HiCache GPU-to-host backup tokens and duration | +| `sglang:hicache_scheduler_phase_*` | Scheduler-thread HiCache control time, calls, and one-second maximum | +| `sglang:hicache_pending_operations` | Outstanding HiCache operations by kind | +| `sglang:dp2_mlp_sync_*` and `sglang:dp2_verify_tier_sync_*` | Native co-located DP2 exchange state and timing, including arrival skew and post-arrival transport cost | + +Use the `dp_rank` label when aggregating scheduler metrics. Gauges such as `sglang:gen_throughput` are emitted per rank; sum them for aggregate DP throughput. Counters can also be summed across ranks when you want deployment-wide work, while histograms should retain the rank label when diagnosing imbalance. diff --git a/python/sglang/benchmark/one_batch.py b/python/sglang/benchmark/one_batch.py index 8556becec68b..0b1e5d6f6f56 100644 --- a/python/sglang/benchmark/one_batch.py +++ b/python/sglang/benchmark/one_batch.py @@ -545,6 +545,7 @@ def _maybe_prepare_mlp_sync_batch(batch: ScheduleBatch, model_runner): require_mlp_tp_gather=require_mlp_tp_gather(model_runner.server_args), disable_overlap_schedule=model_runner.server_args.disable_overlap_schedule, offload_tags=set(), + server_args=model_runner.server_args, ) diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh index c86b92135437..269a360a68ec 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/c_plan.cuh @@ -114,6 +114,14 @@ SGL_DEVICE uint32_t warp_inclusive_sum(uint32_t lane_id, uint32_t val) { return val; } +SGL_DEVICE void fail_prefill_plan_bounds() { +#ifndef USE_ROCM + __trap(); +#else + __builtin_trap(); +#endif +} + __global__ __launch_bounds__(1024, 1) // void plan_compress_prefill_kernel0(const Prefill0Params params) { using namespace device; @@ -154,6 +162,9 @@ __global__ __launch_bounds__(1024, 1) // warp_max[tx] = 0; warp_min[tx] = 0xFFFFFFFFu; } + // Warp i may otherwise publish its reduction before warp 0 has initialized + // warp_{min,max}[i], allowing the initializer to clobber the result. + __syncthreads(); // === Stage B: min/max(extend_len) for MTP-uniform detection === // For min, treat threads outside `batch_size` as +inf so they don't pull the min down. @@ -171,16 +182,16 @@ __global__ __launch_bounds__(1024, 1) // const auto num_q = params.num_q_tokens; // MTP-uniform: every batch shares the same small extend_len `E`, so we can decompose // a global token id `k` into (batch_id, j) = (k / E, k % E) and skip the per-batch loop. - const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32); + // The product check is also a fail-safe against a corrupted min/max result: + // the fast path must cover exactly the allocated ragged-token domain. + const bool is_mtp_extend = (s_min_extend == s_max_extend) && (s_max_extend > 0) && (s_max_extend <= 32) && + (static_cast(params.batch_size) * s_max_extend == num_q); // === Stage C: emit valid plans, slot allocation via shared-mem atomicAdd === if (is_mtp_extend) { // Path 1: token-driven. Each global token id maps to exactly one (batch_id, j). const uint32_t E = s_max_extend; - // num_q is the padded buffer size (graph bucket), not the work size: cap the - // loop at the real token count so batch_id = k / E stays < batch_size on an - // underfilled replay; Stage D pads [counter, num_q) with invalid. - const uint32_t num_real_q = params.batch_size * E; + const uint32_t num_real_q = num_q; for (uint32_t k = tx; k < num_real_q; k += block_size) { const uint32_t batch_id = k / E; const uint32_t j = k % E; @@ -192,6 +203,7 @@ __global__ __launch_bounds__(1024, 1) // if ((position + 1) % cr == 0) { const int32_t buffer_len = window_size - min(static_cast(j) + 1, window_size); const uint32_t out_idx = atomicAdd(&counter_c, 1u); + if (out_idx >= num_q) fail_prefill_plan_bounds(); params.plan_c[out_idx] = { .seq_len = static_cast(position + 1), .ragged_id = static_cast(ragged_id), @@ -207,6 +219,7 @@ __global__ __launch_bounds__(1024, 1) // if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr); if (do_write) { const uint32_t out_idx = atomicAdd(&counter_w, 1u); + if (out_idx >= num_q) fail_prefill_plan_bounds(); params.plan_w[out_idx] = pack_w(ragged_id, batch_id, position + 1); } } @@ -227,6 +240,7 @@ __global__ __launch_bounds__(1024, 1) // if ((position + 1) % cr == 0) { const int32_t buffer_len = window_size - min(j + 1, window_size); const uint32_t out_idx = atomicAdd(&counter_c, 1u); + if (out_idx >= num_q) fail_prefill_plan_bounds(); params.plan_c[out_idx] = { .seq_len = static_cast(position + 1), .ragged_id = static_cast(ragged_id), @@ -240,11 +254,15 @@ __global__ __launch_bounds__(1024, 1) // if (!do_write && is_overlap) do_write = (position % sps) >= (sps - cr); if (do_write) { const uint32_t out_idx = atomicAdd(&counter_w, 1u); + if (out_idx >= num_q) fail_prefill_plan_bounds(); params.plan_w[out_idx] = pack_w(ragged_id, static_cast(batch_id), position + 1); } } base_e += static_cast(el); } + // num_q can be a larger CUDA-graph token bucket; Stage D intentionally + // pads [base_e, num_q) with invalid records. Only overflow is illegal. + if (tx == 0 && base_e > num_q) fail_prefill_plan_bounds(); } __syncthreads(); @@ -504,7 +522,11 @@ inline PrefillPlan plan_compress_prefill( const auto batch_size = static_cast(B.unwrap()); constexpr auto kMaxTokens = static_cast(std::numeric_limits::max()); RuntimeCheck(compress_ratio == 4 || compress_ratio == 128); - RuntimeCheck(batch_size <= num_q_tokens && num_q_tokens <= kMaxTokens); + // The GPU planner accepts fixed request axes containing zero-length graph + // sentinels, so batch_size can legitimately exceed the child token bucket. + // kernel0 validates the stronger invariant sum(extend_lens) <= num_q_tokens + // on device before any out-of-bounds plan write is possible. + RuntimeCheck(num_q_tokens <= kMaxTokens); // `swa_page_size` >= `ring_size` >= `compress_ratio` RuntimeCheck(swa_page_size % ring_size == 0 && ring_size % compress_ratio == 0); @@ -557,6 +579,9 @@ inline PrefillPlan plan_compress_prefill( return PrefillPlan{std::move(C), std::move(W)}; } + // The CPU planner still requires one or more query tokens per request. + RuntimeCheck(batch_size <= num_q_tokens); + // CPU input path: only here do we need the pinned scratch buffer. const auto pin_buffer_bytes = static_cast(pin_buffer.numel()) * sizeof(uint8_t); RuntimeCheck(pin_buffer_bytes >= num_q_tokens * (sizeof(PlanC) + sizeof(PlanW))); diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh index 848cb3646302..1e9289a0e45e 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh +++ b/python/sglang/kernels/jit/csrc/deepseek_v4/main_norm_rope.cuh @@ -65,8 +65,12 @@ load_rope_first_cos_sin(const float* __restrict__ cos_sin_cache, int32_t lane_id // ============================================================================ struct FusedQNormRopeParams { - const void* __restrict__ q_input; // (B, num_q_heads, kHeadDim) DType - void* __restrict__ q_output; // (B, num_q_heads, kHeadDim) DType + // q_input and q_output may alias for the allocation-free eager-prefill + // path. Every warp loads its complete input head into registers before its + // first output store, so exact in-place execution is safe. Do not mark this + // pair restrict: doing so would make the intentional alias undefined. + const void* q_input; // (B, num_q_heads, kHeadDim) DType + void* q_output; // (B, num_q_heads, kHeadDim) DType const float* __restrict__ freqs_cis; // (max_pos, kRopeDim) fp32 (re/im interleaved) const void* __restrict__ positions; // (B,) PosT int64_t q_input_stride_batch; diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 326d313df2a0..b20cd2dd6b29 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/python/sglang/srt/batch_overlap/two_batch_overlap.py @@ -41,7 +41,13 @@ from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.runtime_context import get_parallel, get_server_args from sglang.srt.speculative.spec_info import SpecInput -from sglang.srt.utils import BumpAllocator, empty_context, get_bool_env_var, is_hip +from sglang.srt.utils import ( + BumpAllocator, + empty_context, + get_bool_env_var, + get_int_env_var, + is_hip, +) if TYPE_CHECKING: from sglang.srt.batch_overlap.single_batch_overlap import CombineOverlapArgs @@ -51,6 +57,7 @@ _is_hip = is_hip() _tbo_debug = get_bool_env_var("SGLANG_TBO_DEBUG") +_tbo_max_prefill_tokens = get_int_env_var("SGLANG_TBO_MAX_PREFILL_TOKENS", 0) logger = logging.getLogger(__name__) @@ -338,20 +345,31 @@ def __init__(self): (2,), dtype=torch.int32, device=get_server_args().device ) - def capture_one_batch_size(self, batch: ForwardBatch, num_tokens: int): - if not is_tbo_enabled(): + def capture_one_batch_size( + self, + batch: ForwardBatch, + num_tokens: int, + *, + enabled: Optional[bool] = None, + attn_backend=None, + ): + if enabled is None: + enabled = is_tbo_enabled() + if not enabled or not batch.forward_mode.is_extend_without_speculative(): return - token_num_per_seq = get_token_num_per_seq( - forward_mode=batch.forward_mode, spec_info=batch.spec_info - ) + if _tbo_max_prefill_tokens > 0 and num_tokens > _tbo_max_prefill_tokens: + raise RuntimeError( + "prefill CUDA-graph TBO capture exceeds the configured dense " + f"ceiling: tokens={num_tokens}, " + f"limit={_tbo_max_prefill_tokens}" + ) batch.tbo_split_seq_index = compute_split_seq_index( forward_mode=batch.forward_mode, num_tokens=num_tokens, - extend_lens=None, - token_num_per_seq=token_num_per_seq, + extend_lens=batch.extend_seq_lens_cpu, + token_num_per_seq=None, ) - # For simplicity, when two_batch_overlap is enabled, we only capture CUDA Graph for tbo=true assert batch.tbo_split_seq_index is not None, f"{num_tokens=}" self._tbo_children_num_token_non_padded[...] = ( @@ -361,6 +379,46 @@ def capture_one_batch_size(self, batch: ForwardBatch, num_tokens: int): TboForwardBatchPreparer.prepare_raw( batch, tbo_children_num_token_non_padded=self._tbo_children_num_token_non_padded, + attn_backend=attn_backend, + ) + # Full prefill graphs use a fixed request axis. The capture batch has + # one real sequence followed by zero-length sentinels, so a two-chunk + # split naturally gives child A one slot and child B every slot. Pad + # child A to the same fixed geometry: replay can then refresh both + # child attention plans for arbitrary live request layouts without + # changing any captured tensor shape. + for child in batch.tbo_children: + TboForwardBatchPreparer.pad_sequence_axis_for_cuda_graph( + child, + target_batch_size=batch.batch_size, + ) + # Full DP prefill graphs execute a fixed token bucket on every rank. + # Seed each child's fixed DP geometry before capture so model.forward + # never needs a CUDA->host size gather while the stream is recording. + if batch.global_num_tokens_cpu is not None: + world = len(batch.global_num_tokens_cpu) + for child in batch.tbo_children: + rows = int(child.tbo_padded_len) + child.global_num_tokens_cpu = [rows] * world + child.global_num_tokens_gpu = ( + torch.full_like(batch.global_num_tokens_gpu, rows) + if batch.global_num_tokens_gpu is not None + else None + ) + child.global_dp_buffer_len = rows * world + + def replay_prepare_prefill( + self, + *, + split_token_index: int, + num_token_non_padded: int, + ) -> None: + """Refresh graph-captured child logical token counts for prefill.""" + self._tbo_children_num_token_non_padded[...] = ( + TboForwardBatchPreparer.compute_tbo_children_num_token_non_padded_raw( + tbo_split_token_index=split_token_index, + num_token_non_padded=num_token_non_padded, + ) ) def replay_prepare( @@ -410,19 +468,22 @@ def prepare_all_gather( return False, self._compute_local_forward_mode(local_batch) if local_batch is not None: + is_real_prefill = local_batch.forward_mode.is_extend_without_speculative() + within_prefill_limit = ( + _tbo_max_prefill_tokens <= 0 + or local_batch.extend_num_tokens <= _tbo_max_prefill_tokens + ) + if not is_real_prefill or not within_prefill_limit: + self.local_tbo_split_seq_index = None + local_can_run_tbo = False + local_forward_mode = self._compute_local_forward_mode(local_batch) + return local_can_run_tbo, local_forward_mode + token_num_per_seq = get_token_num_per_seq( forward_mode=local_batch.forward_mode, spec_info=local_batch.spec_info ) - if ( - local_batch.forward_mode.is_target_verify() - or local_batch.forward_mode.is_decode() - ): - num_tokens = local_batch.batch_size() * token_num_per_seq - elif local_batch.forward_mode.is_prebuilt(): - num_tokens = 0 - else: - num_tokens = local_batch.extend_num_tokens + num_tokens = local_batch.extend_num_tokens self.local_tbo_split_seq_index = compute_split_seq_index( forward_mode=local_batch.forward_mode, num_tokens=num_tokens, @@ -449,8 +510,14 @@ def prepare_all_gather( def compute_output(self, partial_global_info): # Perform only one Device-to-Host (D2H) memory copy cpu_data = partial_global_info[:, :2].cpu() - local_can_run_tbo_aggregated = min(cpu_data[:, 0].tolist()) - forward_modes = cpu_data[:, 1].tolist() + return self.compute_output_from_values( + local_can_run_tbo=cpu_data[:, 0].tolist(), + forward_modes=cpu_data[:, 1].tolist(), + ) + + def compute_output_from_values(self, local_can_run_tbo, forward_modes): + """Resolve TBO from already-host-resident native DP sync values.""" + local_can_run_tbo_aggregated = min(local_can_run_tbo) global_forward_mode, forward_mode_agree = self._compute_global_forward_mode( forward_modes @@ -514,7 +581,11 @@ def prepare(cls, batch: ForwardBatch, is_draft_worker: bool = False): @classmethod def prepare_raw( - cls, batch: ForwardBatch, tbo_children_num_token_non_padded: torch.Tensor + cls, + batch: ForwardBatch, + tbo_children_num_token_non_padded: torch.Tensor, + *, + attn_backend=None, ): from sglang.srt.layers.attention.tbo_backend import TboAttnBackend @@ -538,7 +609,8 @@ def prepare_raw( # Sanity check: the global attn_backend should be a TboAttnBackend # whose children handle the two halves. - attn_backend = get_attn_backend() + if attn_backend is None: + attn_backend = get_attn_backend() assert isinstance(attn_backend, TboAttnBackend) [out_num_token_non_padded_a, out_num_token_non_padded_b] = ( @@ -819,6 +891,57 @@ def filter_batch( return ForwardBatch(**output_dict) + @staticmethod + def pad_sequence_axis_for_cuda_graph( + batch: ForwardBatch, + *, + target_batch_size: int, + ) -> None: + """Pad a TBO child to a graph-stable request axis with zero sentinels.""" + current_batch_size = batch.batch_size + if current_batch_size == target_batch_size: + return + if current_batch_size > target_batch_size: + raise RuntimeError( + "TBO CUDA-graph child request axis exceeds its parent: " + f"child={current_batch_size}, parent={target_batch_size}" + ) + + pad = target_batch_size - current_batch_size + token_end = batch.input_ids.shape[0] + tensor_fill = { + "req_pool_indices": 0, + "seq_lens": 0, + "seq_lens_cpu": 0, + "orig_seq_lens": 0, + "extend_seq_lens": 0, + "extend_prefix_lens": 0, + "extend_start_loc": token_end, + } + for field, fill_value in tensor_fill.items(): + value = getattr(batch, field, None) + if value is None: + continue + padding = torch.full( + (pad, *value.shape[1:]), + fill_value, + dtype=value.dtype, + device=value.device, + ) + setattr(batch, field, torch.cat((value, padding), dim=0)) + + list_fill = { + "extend_seq_lens_cpu": 0, + "extend_prefix_lens_cpu": 0, + "extend_logprob_start_lens_cpu": 0, + } + for field, fill_value in list_fill.items(): + value = getattr(batch, field, None) + if value is not None: + setattr(batch, field, [*value, *([fill_value] * pad)]) + + batch.batch_size = target_batch_size + @classmethod def compute_tbo_children_num_token_non_padded(cls, batch: ForwardBatch): return cls.compute_tbo_children_num_token_non_padded_raw( diff --git a/python/sglang/srt/constrained/base_grammar_backend.py b/python/sglang/srt/constrained/base_grammar_backend.py index 494a3303cdc8..11f4d575a4dc 100644 --- a/python/sglang/srt/constrained/base_grammar_backend.py +++ b/python/sglang/srt/constrained/base_grammar_backend.py @@ -13,8 +13,14 @@ # ============================================================================== """The base class of a backend for grammar-guided constrained decoding.""" +import dataclasses +import functools +import json import logging +import os +import threading import time +from collections import OrderedDict from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from typing import Dict, List, NamedTuple, Optional, Tuple @@ -40,6 +46,10 @@ class GrammarStats: tree_traversal_time: List[float] = field(default_factory=list) dispatch_type: Optional[str] = None num_timeout: int = 0 + cache_source: str = "compile" + cache_resolution_time: Optional[float] = None + cache_lock_wait_time: Optional[float] = None + cache_phase_seconds: Dict[str, float] = field(default_factory=dict) class GrammarRow(NamedTuple): @@ -103,6 +113,9 @@ def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: def copy(self) -> "BaseGrammarObject": return self + def cache_memory_bytes(self) -> int: + return 0 + @property def finished(self): return self._finished @@ -156,20 +169,76 @@ def apply(self, logits: torch.Tensor) -> None: class InvalidGrammarObject(BaseGrammarObject): """Represents a grammar that failed to compile, carrying the original error message.""" - def __init__(self, error_message: str = "Unknown grammar error"): + def __init__( + self, + error_message: str = "Unknown grammar error", + grammar_stats: Optional[GrammarStats] = None, + ): super().__init__() self.error_message = error_message + self.grammar_stats = grammar_stats def __repr__(self): return f"InvalidGrammarObject(error_message={self.error_message!r})" + def copy(self) -> "InvalidGrammarObject": + stats = ( + None + if self.grammar_stats is None + else dataclasses.replace(self.grammar_stats) + ) + return InvalidGrammarObject(self.error_message, stats) + class BaseGrammarBackend: _enable_strict_thinking: bool = False def __init__(self): - self.executor = ThreadPoolExecutor() - self.cache: Dict[Tuple[str, str], BaseGrammarObject] = {} + executor_threads = int( + os.environ.get("SGLANG_GRAMMAR_COMPILATION_WORKERS", "8") + ) + self.cache_limit_bytes = int( + os.environ.get("SGLANG_GRAMMAR_CACHE_BYTES", str(10 * 1024**3)) + ) + self.cache_limit_entries = int( + os.environ.get("SGLANG_GRAMMAR_CACHE_ENTRIES", "100000") + ) + if ( + executor_threads <= 0 + or self.cache_limit_bytes <= 0 + or self.cache_limit_entries <= 0 + ): + raise ValueError( + "optimized grammar worker and cache limits must all be positive" + ) + self.executor = ThreadPoolExecutor( + max_workers=executor_threads, + thread_name_prefix="sglang-grammar", + ) + self.cache: OrderedDict[Tuple[str, str], BaseGrammarObject] = OrderedDict() + self._cache_entry_bytes: Dict[Tuple[str, str], int] = {} + self._cache_bytes = 0 + self._inflight: Dict[Tuple[str, str], Future[BaseGrammarObject]] = {} + self._cache_lock = threading.Lock() + + @staticmethod + @functools.lru_cache(maxsize=16384) + def _normalize_cache_key(key: Tuple[str, str]) -> Tuple[str, str]: + key_type, key_string = key + if key_type not in {"json", "structural_tag"} or key_string == "$$ANY$$": + return key + try: + parsed = json.loads(key_string) + except (json.JSONDecodeError, UnicodeDecodeError): + return key + return ( + key_type, + json.dumps( + parsed, + ensure_ascii=False, + separators=(",", ":"), + ), + ) def initialize_vocab_mask_buffer( self, @@ -244,19 +313,149 @@ def _init_value_dispatch( def get_cached_or_future_value( self, key: Tuple[str, str], require_reasoning: bool ) -> Tuple[BaseGrammarObject | Future[BaseGrammarObject], bool]: - value = self.cache.get(key) - if value: - copied_value = value.copy() - copied_value.maybe_init_reasoning(require_reasoning) - return copied_value, True - value = self.executor.submit(self._init_value_dispatch, key, require_reasoning) - return value, False + key = self._normalize_cache_key(key) + with self._cache_lock: + value = self.cache.get(key) + if value is not None: + self.cache.move_to_end(key) + return ( + self._copy_for_request( + value, + require_reasoning=require_reasoning, + cache_source="memory", + resolution_time=0.0, + preserve_compilation_time=False, + ), + True, + ) + + shared_future = self._inflight.get(key) + is_owner = shared_future is None + if shared_future is None: + shared_future = self.executor.submit( + self._init_value_dispatch, + key, + False, + ) + self._inflight[key] = shared_future + + # Future.add_done_callback() invokes the callback synchronously when + # the future has already completed. Register outside _cache_lock so a + # fast compile cannot deadlock in _finish_inflight while reacquiring it. + if is_owner: + shared_future.add_done_callback( + lambda future, cache_key=key: self._finish_inflight( + cache_key, + future, + ) + ) + + request_future: Future[BaseGrammarObject] = Future() + wait_started = time.perf_counter() + + def finish_request(future: Future[BaseGrammarObject]) -> None: + if request_future.cancelled(): + return + try: + prototype = future.result() + prototype_stats = prototype.grammar_stats + prototype_source = ( + prototype_stats.cache_source + if prototype_stats is not None + else "compile" + ) + source = prototype_source if is_owner else "inflight" + request_future.set_result( + self._copy_for_request( + prototype, + require_reasoning=require_reasoning, + cache_source=source, + resolution_time=time.perf_counter() - wait_started, + preserve_compilation_time=is_owner, + ) + ) + except BaseException as exc: + request_future.set_exception(exc) + + shared_future.add_done_callback(finish_request) + return request_future, False + + def _finish_inflight( + self, + key: Tuple[str, str], + future: Future[BaseGrammarObject], + ) -> None: + try: + prototype = future.result() + except BaseException: + with self._cache_lock: + self._inflight.pop(key, None) + return + with self._cache_lock: + self._inflight.pop(key, None) + self._set_cache_locked(key, prototype) + + @staticmethod + def _copy_for_request( + prototype: BaseGrammarObject, + *, + require_reasoning: bool, + cache_source: str, + resolution_time: float, + preserve_compilation_time: bool, + ) -> BaseGrammarObject: + prototype_stats = getattr(prototype, "grammar_stats", None) + copied_value = prototype.copy() + copied_value.maybe_init_reasoning(require_reasoning) + copied_stats = getattr(copied_value, "grammar_stats", None) + if copied_stats is not None: + copied_stats.cache_source = cache_source + copied_stats.is_cache_hit = cache_source in { + "memory", + "disk", + "inflight", + } + copied_stats.cache_resolution_time = resolution_time + if not preserve_compilation_time: + copied_stats.compilation_time = None + copied_stats.cache_phase_seconds = {} + elif prototype_stats is not None: + copied_stats.compilation_time = prototype_stats.compilation_time + copied_stats.cache_lock_wait_time = prototype_stats.cache_lock_wait_time + copied_stats.cache_phase_seconds = dict( + prototype_stats.cache_phase_seconds + ) + return copied_value def set_cache(self, key: Tuple[str, str], value: BaseGrammarObject): + key = self._normalize_cache_key(key) + with self._cache_lock: + self._set_cache_locked(key, value) + + def _set_cache_locked( + self, + key: Tuple[str, str], + value: BaseGrammarObject, + ) -> None: + previous_size = self._cache_entry_bytes.pop(key, 0) + if key in self.cache: + self.cache.pop(key) + size = max(0, int(value.cache_memory_bytes())) self.cache[key] = value + self._cache_entry_bytes[key] = size + self._cache_bytes += size - previous_size + while ( + self._cache_bytes > self.cache_limit_bytes + or len(self.cache) > self.cache_limit_entries + ): + evicted_key, _ = self.cache.popitem(last=False) + self._cache_bytes -= self._cache_entry_bytes.pop(evicted_key, 0) def reset(self): - self.cache.clear() + with self._cache_lock: + self.cache.clear() + self._cache_entry_bytes.clear() + self._cache_bytes = 0 def register_vocab_mask_buffer( diff --git a/python/sglang/srt/constrained/grammar_manager.py b/python/sglang/srt/constrained/grammar_manager.py index b039020fdb64..1f1074e28be8 100644 --- a/python/sglang/srt/constrained/grammar_manager.py +++ b/python/sglang/srt/constrained/grammar_manager.py @@ -8,6 +8,7 @@ import torch from sglang.srt.constrained.base_grammar_backend import ( + GrammarStats, InvalidGrammarObject, create_grammar_backend, ) @@ -112,7 +113,18 @@ def abort_requests(self, recv_req: AbortReq): logger.debug(f"Abort grammar queue request. {req.rid=}") if isinstance(req.grammar, futures.Future) and req.grammar: req.grammar.cancel() + dispatch_type = ( + req.grammar_key[0] if req.grammar_key is not None else None + ) + req.grammar = InvalidGrammarObject( + "Aborted by AbortReq.", + GrammarStats( + dispatch_type=dispatch_type, + is_grammar_aborted=True, + ), + ) req.set_finish_with_abort("Aborted by AbortReq.") + req.log_grammar_stats_once(self.scheduler.metrics_collector) def _get_request_thinking_budget(self, req: Req) -> int | None: custom_params = req.sampling_params.custom_params @@ -166,6 +178,7 @@ def process_req_with_grammar(self, req: Req) -> bool: f"Failed to compile {key[0]} grammar: {value.error_message}" ) req.set_finish_with_abort(error_msg) + req.log_grammar_stats_once(self.scheduler.metrics_collector) else: self._apply_request_reasoning_budget(req) elif self._enable_strict_thinking: @@ -282,12 +295,18 @@ def get_ready_grammar_requests(self) -> List[Req]: f"Grammar compilation raised an exception: {e}, " f"grammar_key={req.grammar_key}" ) - req.grammar = InvalidGrammarObject(f"Grammar compilation failed: {e}") - self.grammar_backend.set_cache(req.grammar_key, req.grammar.copy()) + req.grammar = InvalidGrammarObject( + f"Grammar compilation failed: {e}", + GrammarStats( + dispatch_type=req.grammar_key[0], + is_grammar_aborted=True, + ), + ) self._apply_request_reasoning_budget(req) if isinstance(req.grammar, InvalidGrammarObject): error_msg = f"Failed to compile {req.grammar_key[0]} grammar: {req.grammar.error_message}" req.set_finish_with_abort(error_msg) + req.log_grammar_stats_once(self.scheduler.metrics_collector) # Return failed requests for i in synced_failed_req_idxs: @@ -296,11 +315,19 @@ def get_ready_grammar_requests(self) -> List[Req]: assert isinstance(req.grammar, futures.Future) and req.grammar_key req.grammar.cancel() - self.grammar_backend.set_cache( - req.grammar_key, InvalidGrammarObject("Grammar preprocessing timed out") + invalid_grammar = InvalidGrammarObject( + "Grammar preprocessing timed out", + GrammarStats( + dispatch_type=req.grammar_key[0], + is_grammar_aborted=True, + num_timeout=1, + ), ) + req.grammar = invalid_grammar + self.grammar_backend.set_cache(req.grammar_key, invalid_grammar.copy()) error_msg = f"Grammar preprocessing timed out: {req.grammar_key=}" req.set_finish_with_abort(error_msg) + req.log_grammar_stats_once(self.scheduler.metrics_collector) # Remove finished requests from grammar_queue self.grammar_queue = [ diff --git a/python/sglang/srt/constrained/reasoner_grammar_backend.py b/python/sglang/srt/constrained/reasoner_grammar_backend.py index a567bab6acc7..6c3f4d293275 100644 --- a/python/sglang/srt/constrained/reasoner_grammar_backend.py +++ b/python/sglang/srt/constrained/reasoner_grammar_backend.py @@ -71,6 +71,7 @@ def __init__( self.move_vocab_mask_fn = move_vocab_mask_fn self.apply_vocab_mask_fn = apply_vocab_mask_fn self._think_end_id_list = [think_end_id] + self.grammar_stats = grammar.grammar_stats if grammar is not None else None self.tokens_in_think = -1 self.tokens_after_end = -1 @@ -194,6 +195,11 @@ def copy(self): new_obj._finished = self._finished return new_obj + def cache_memory_bytes(self) -> int: + if self.grammar is None: + return 0 + return self.grammar.cache_memory_bytes() + @property def finished(self): if self.grammar is not None: diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index a0b44930c226..b51c4a939107 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -16,6 +16,7 @@ import dataclasses import json import logging +import os from typing import Dict, List, Optional, Tuple, Union import torch @@ -37,6 +38,10 @@ InvalidGrammarObject, ) from sglang.srt.constrained.utils import is_legacy_structural_tag +from sglang.srt.constrained.xgrammar_persistent_cache import ( + CompiledGrammarLookup, + PersistentXGrammarCache, +) from sglang.srt.utils import is_hip _is_hip = is_hip() @@ -78,7 +83,7 @@ def __init__( ctx: CompiledGrammar, override_stop_tokens: Optional[Union[List[int], int]], key_string: Optional[str] = None, - grammar_stats: Optional[GrammarStats] = GrammarStats(), + grammar_stats: Optional[GrammarStats] = None, ) -> None: super().__init__() self.matcher = matcher @@ -143,7 +148,13 @@ def copy(self): ) if grammar_stats := self.grammar_stats: grammar_stats = dataclasses.replace( - grammar_stats, is_cache_hit=True, tree_traversal_time=[] + grammar_stats, + compilation_time=None, + is_cache_hit=True, + tree_traversal_time=[], + cache_source="memory", + cache_resolution_time=0.0, + cache_lock_wait_time=0.0, ) return XGrammarGrammar( matcher, @@ -154,6 +165,9 @@ def copy(self): grammar_stats, ) + def cache_memory_bytes(self) -> int: + return int(self.ctx.memory_size_bytes) + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: s = self.matcher.find_jump_forward_string() if s: @@ -230,10 +244,45 @@ def __init__( f"Failed to create XGrammar TokenizerInfo from tokenizer: {e}" ) - self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info) + compiler_threads = int(os.environ.get("SGLANG_XGRAMMAR_COMPILER_THREADS", "12")) + persistent_cache_bytes = int( + os.environ.get( + "SGLANG_XGRAMMAR_DISK_CACHE_BYTES", + str(40 * 1024**3), + ) + ) + if compiler_threads <= 0 or persistent_cache_bytes <= 0: + raise ValueError( + "optimized XGrammar compiler and disk cache limits must be positive" + ) + self.grammar_compiler = GrammarCompiler( + tokenizer_info=tokenizer_info, + max_threads=compiler_threads, + cache_enabled=False, + ) self.vocab_size = vocab_size self.override_stop_tokens = override_stop_tokens self.any_whitespace = any_whitespace + self.persistent_cache = PersistentXGrammarCache( + tokenizer_info=tokenizer_info, + cache_directory=os.environ.get( + "SGLANG_XGRAMMAR_CACHE_DIR", + "~/.cache/sglang/xgrammar", + ), + max_bytes=persistent_cache_bytes, + compiler_identity={ + "any_whitespace": any_whitespace, + "compiler_threads": compiler_threads, + "override_stop_tokens": ( + list(override_stop_tokens) + if isinstance(override_stop_tokens, list) + else ( + [] if override_stop_tokens is None else [override_stop_tokens] + ) + ), + "vocab_size": vocab_size, + }, + ) @property def is_support_token_filter(self): @@ -327,36 +376,107 @@ def _from_context( grammar_stats, ) + @staticmethod + def _stats_from_lookup( + dispatch_type: str, + lookup: CompiledGrammarLookup, + ) -> GrammarStats: + return GrammarStats( + dispatch_type=dispatch_type, + is_cache_hit=lookup.source == "disk", + cache_source=lookup.source, + cache_resolution_time=lookup.resolution_seconds, + cache_lock_wait_time=lookup.lock_wait_seconds, + cache_phase_seconds=lookup.phase_seconds, + ) + + def _persistent_get_or_compile( + self, + *, + key_type: str, + key_string: str, + compile_fn, + ) -> CompiledGrammarLookup: + return self.persistent_cache.get_or_compile( + key_type=key_type, + key_string=key_string, + compile_fn=compile_fn, + ) + def dispatch_json(self, key_string: str) -> BaseGrammarObject: try: if key_string == "$$ANY$$": # Note: This builtin JSON grammar includes *all* valid JSON (including, for example, arrays at the root) - ctx = self.grammar_compiler.compile_builtin_json_grammar() + compile_fn = self.grammar_compiler.compile_builtin_json_grammar else: - ctx = self.grammar_compiler.compile_json_schema( - schema=key_string, any_whitespace=self.any_whitespace + compile_fn = lambda: self.grammar_compiler.compile_json_schema( + schema=key_string, + any_whitespace=self.any_whitespace, ) + lookup = self._persistent_get_or_compile( + key_type="json", + key_string=key_string, + compile_fn=compile_fn, + ) except (RuntimeError, json.decoder.JSONDecodeError, UnicodeDecodeError) as e: logger.error(f"Hit invalid json_schema: {key_string=}, {e=}") - return InvalidGrammarObject(str(e)) - return self._from_context(ctx, key_string, GrammarStats(dispatch_type="json")) + return InvalidGrammarObject( + str(e), + GrammarStats( + dispatch_type="json", + is_grammar_aborted=True, + ), + ) + return self._from_context( + lookup.grammar, + key_string, + self._stats_from_lookup("json", lookup), + ) def dispatch_ebnf(self, key_string: str) -> BaseGrammarObject: try: - ctx = self.grammar_compiler.compile_grammar(key_string) + lookup = self._persistent_get_or_compile( + key_type="ebnf", + key_string=key_string, + compile_fn=lambda: self.grammar_compiler.compile_grammar(key_string), + ) except RuntimeError as e: logger.error(f"Hit invalid ebnf: {key_string=}, {e=}") - return InvalidGrammarObject(str(e)) - return self._from_context(ctx, key_string, GrammarStats(dispatch_type="ebnf")) + return InvalidGrammarObject( + str(e), + GrammarStats( + dispatch_type="ebnf", + is_grammar_aborted=True, + ), + ) + return self._from_context( + lookup.grammar, + key_string, + self._stats_from_lookup("ebnf", lookup), + ) def dispatch_regex(self, key_string: str) -> BaseGrammarObject: try: - ctx = self.grammar_compiler.compile_regex(key_string) + lookup = self._persistent_get_or_compile( + key_type="regex", + key_string=key_string, + compile_fn=lambda: self.grammar_compiler.compile_regex(key_string), + ) except RuntimeError as e: logger.error(f"Hit invalid regex: {key_string=}, {e=}") - return InvalidGrammarObject(str(e)) - return self._from_context(ctx, key_string, GrammarStats(dispatch_type="regex")) + return InvalidGrammarObject( + str(e), + GrammarStats( + dispatch_type="regex", + is_grammar_aborted=True, + ), + ) + return self._from_context( + lookup.grammar, + key_string, + self._stats_from_lookup("regex", lookup), + ) def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject: try: @@ -376,19 +496,36 @@ def dispatch_structural_tag(self, key_string: str) -> BaseGrammarObject: tags, structural_tag["triggers"] ) new_tag.format.at_least_one = structural_tag.get("at_least_one", False) - ctx = self.grammar_compiler.compile_structural_tag(new_tag) + compile_fn = lambda: self.grammar_compiler.compile_structural_tag( + new_tag + ) else: format_dict = structural_tag.get("format") if isinstance(format_dict, dict): self._sanitize_structural_format(format_dict) structural_tag["format"] = format_dict key_string = json.dumps(structural_tag) - ctx = self.grammar_compiler.compile_structural_tag(key_string) + compile_fn = lambda: self.grammar_compiler.compile_structural_tag( + key_string + ) + lookup = self._persistent_get_or_compile( + key_type="structural_tag", + key_string=key_string, + compile_fn=compile_fn, + ) except (RuntimeError, json.decoder.JSONDecodeError) as e: logger.error(f"Hit invalid structural_tag: {key_string=}, {e=}") - return InvalidGrammarObject(str(e)) + return InvalidGrammarObject( + str(e), + GrammarStats( + dispatch_type="structural_tag", + is_grammar_aborted=True, + ), + ) return self._from_context( - ctx, key_string, GrammarStats(dispatch_type="structural_tag") + lookup.grammar, + key_string, + self._stats_from_lookup("structural_tag", lookup), ) def reset(self): diff --git a/python/sglang/srt/constrained/xgrammar_persistent_cache.py b/python/sglang/srt/constrained/xgrammar_persistent_cache.py new file mode 100644 index 000000000000..b0958c05b51d --- /dev/null +++ b/python/sglang/srt/constrained/xgrammar_persistent_cache.py @@ -0,0 +1,520 @@ +"""Process-safe persistent cache for compatible XGrammar compiled grammars. + +The cache is intentionally fail-closed. An invalid serialized entry or +adaptive local-compile marker is surfaced as an error instead of being hidden +by a recovery path. That makes image/tokenizer/cache incompatibilities +observable immediately. +""" + +from __future__ import annotations + +import fcntl +import hashlib +import importlib.metadata +import json +import os +import struct +import tempfile +import time +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Literal + +from xgrammar import CompiledGrammar, TokenizerInfo + +type CacheSource = Literal["compile", "disk", "local_compile"] + + +@dataclass(frozen=True, slots=True) +class CompiledGrammarLookup: + grammar: CompiledGrammar + source: CacheSource + lock_wait_seconds: float + resolution_seconds: float + phase_seconds: dict[str, float] + + +@dataclass(frozen=True, slots=True) +class _LocalCompileMarker: + native_compile_nanoseconds: int + serialized_bytes: int + + +class PersistentXGrammarCache: + """A tokenizer-scoped, adaptive cross-process compiled-grammar cache.""" + + _COUNTER_FORMAT = " None: + if deserialize_bytes_per_second is None: + deserialize_bytes_per_second = int( + os.environ.get( + "SGLANG_XGRAMMAR_DESERIALIZE_BYTES_PER_SECOND", + str(128 * 1024**2), + ) + ) + if local_compile_speedup is None: + local_compile_speedup = int( + os.environ.get( + "SGLANG_XGRAMMAR_LOCAL_COMPILE_SPEEDUP", + "2", + ) + ) + if ( + max_bytes < self._LOCAL_COMPILE_MARKER_BYTES + or deserialize_bytes_per_second <= 0 + or local_compile_speedup <= 0 + ): + raise ValueError( + "persistent XGrammar cache must fit one local-compile marker, " + "and deserialize throughput/local-compile speedup must be positive" + ) + + xgrammar_version = importlib.metadata.version("xgrammar") + tokenizer_digest = hashlib.sha256( + tokenizer_info.serialize_json().encode("utf-8") + ).hexdigest() + namespace_payload = { + "format": 2, + "xgrammar_version": xgrammar_version, + "tokenizer_sha256": tokenizer_digest, + "compiler": compiler_identity, + "max_bytes": max_bytes, + "deserialize_bytes_per_second": deserialize_bytes_per_second, + "local_compile_speedup": local_compile_speedup, + } + namespace = hashlib.sha256( + json.dumps( + namespace_payload, + ensure_ascii=False, + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + ).hexdigest() + + self.max_bytes = max_bytes + self.deserialize_bytes_per_second = deserialize_bytes_per_second + self.local_compile_speedup = local_compile_speedup + self.root = Path(cache_directory).expanduser().resolve() / namespace + self.entries = self.root / "entries" + self.locks = self.root / "locks" + self.entries.mkdir(parents=True, exist_ok=True, mode=0o700) + self.locks.mkdir(parents=True, exist_ok=True, mode=0o700) + self._tokenizer_info = tokenizer_info + self._size_ledger = self.root / "size-bytes-v2" + self._accounting_session = self.root / "accounting-session-v2" + self._prune_lock = self.root / "prune.lock" + self._initialize_size_ledger() + + def get_or_compile( + self, + *, + key_type: str, + key_string: str, + compile_fn: Callable[[], CompiledGrammar], + ) -> CompiledGrammarLookup: + key_digest = hashlib.sha256( + json.dumps( + [key_type, key_string], + ensure_ascii=False, + separators=(",", ":"), + ).encode("utf-8") + ).hexdigest() + entry_path = self.entries / f"{key_digest}.json" + lock_path = self.locks / f"{key_digest}.lock" + + lock_started = time.perf_counter() + with lock_path.open("a+b") as lock_file: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + lock_wait_seconds = time.perf_counter() - lock_started + resolution_started = time.perf_counter() + + if entry_path.is_file(): + entry_read_started = time.perf_counter() + try: + entry_payload = entry_path.read_bytes() + local_compile_marker = self._decode_local_compile_marker( + entry_payload + ) + if ( + local_compile_marker is not None + and not self._prefer_local_compile( + native_compile_nanoseconds=( + local_compile_marker.native_compile_nanoseconds + ), + serialized_bytes=local_compile_marker.serialized_bytes, + ) + ): + raise ValueError( + "adaptive XGrammar local-compile marker contradicts " + "the active cache policy" + ) + except Exception as exc: + raise RuntimeError( + f"persistent XGrammar cache entry is invalid: {entry_path}" + ) from exc + entry_read_seconds = time.perf_counter() - entry_read_started + + if local_compile_marker is not None: + compile_started_ns = time.perf_counter_ns() + grammar = compile_fn() + compile_nanoseconds = time.perf_counter_ns() - compile_started_ns + os.utime(entry_path, None) + resolved_at = time.perf_counter() + return CompiledGrammarLookup( + grammar=grammar, + source="local_compile", + lock_wait_seconds=lock_wait_seconds, + resolution_seconds=resolved_at - resolution_started, + phase_seconds={ + "adaptive_marker_read": entry_read_seconds, + "native_compile": compile_nanoseconds / 1_000_000_000, + }, + ) + + disk_started = time.perf_counter() + try: + serialized = entry_payload.decode("utf-8") + grammar = CompiledGrammar.deserialize_json( + serialized, + self._tokenizer_info, + ) + os.utime(entry_path, None) + except Exception as exc: + raise RuntimeError( + f"persistent XGrammar cache entry is invalid: {entry_path}" + ) from exc + resolved_at = time.perf_counter() + return CompiledGrammarLookup( + grammar=grammar, + source="disk", + lock_wait_seconds=lock_wait_seconds, + resolution_seconds=resolved_at - resolution_started, + phase_seconds={ + "disk_read": entry_read_seconds, + "disk_deserialize": resolved_at - disk_started, + }, + ) + + compile_started_ns = time.perf_counter_ns() + grammar = compile_fn() + compile_nanoseconds = time.perf_counter_ns() - compile_started_ns + compile_seconds = compile_nanoseconds / 1_000_000_000 + serialize_started = time.perf_counter() + serialized = grammar.serialize_json() + serialized_payload = serialized.encode("utf-8") + serialize_seconds = time.perf_counter() - serialize_started + serialized_bytes = len(serialized_payload) + prefer_local_compile = self._prefer_local_compile( + native_compile_nanoseconds=compile_nanoseconds, + serialized_bytes=serialized_bytes, + ) + if prefer_local_compile: + stored_payload = self._encode_local_compile_marker( + native_compile_nanoseconds=compile_nanoseconds, + serialized_bytes=serialized_bytes, + ) + cache_policy_phase = "policy_local_compile" + lookup_source: CacheSource = "local_compile" + else: + stored_payload = serialized_payload + cache_policy_phase = "policy_serialized" + lookup_source = "compile" + write_started = time.perf_counter() + self._write_atomic_bytes(entry_path, stored_payload) + write_seconds = time.perf_counter() - write_started + accounting_started = time.perf_counter() + try: + self._account_and_prune( + protected=entry_path, + added_bytes=len(stored_payload), + ) + except BaseException: + # The cache is reconstructible. Do not leave an unaccounted + # entry behind when quota accounting fails. + entry_path.unlink(missing_ok=True) + raise + accounting_seconds = time.perf_counter() - accounting_started + resolution_seconds = time.perf_counter() - resolution_started + return CompiledGrammarLookup( + grammar=grammar, + source=lookup_source, + lock_wait_seconds=lock_wait_seconds, + resolution_seconds=resolution_seconds, + phase_seconds={ + "native_compile": compile_seconds, + "serialize": serialize_seconds, + "entry_write": write_seconds, + "account_prune": accounting_seconds, + cache_policy_phase: 0.0, + }, + ) + + def _prefer_local_compile( + self, + *, + native_compile_nanoseconds: int, + serialized_bytes: int, + ) -> bool: + return ( + serialized_bytes > self.max_bytes + or native_compile_nanoseconds + * self.local_compile_speedup + * self.deserialize_bytes_per_second + < serialized_bytes * 1_000_000_000 + ) + + @classmethod + def _encode_local_compile_marker( + cls, + *, + native_compile_nanoseconds: int, + serialized_bytes: int, + ) -> bytes: + if native_compile_nanoseconds < 0 or serialized_bytes <= 0: + raise RuntimeError( + "adaptive XGrammar local-compile marker values are invalid" + ) + body = struct.pack( + cls._LOCAL_COMPILE_BODY_FORMAT, + cls._LOCAL_COMPILE_MAGIC, + native_compile_nanoseconds, + serialized_bytes, + ) + return body + hashlib.sha256(body).digest() + + @classmethod + def _decode_local_compile_marker( + cls, + payload: bytes, + ) -> _LocalCompileMarker | None: + if not payload.startswith(cls._LOCAL_COMPILE_MAGIC): + return None + if len(payload) != cls._LOCAL_COMPILE_MARKER_BYTES: + raise ValueError( + "adaptive XGrammar local-compile marker has the wrong size" + ) + body = payload[: cls._LOCAL_COMPILE_BODY_BYTES] + checksum = payload[cls._LOCAL_COMPILE_BODY_BYTES :] + if hashlib.sha256(body).digest() != checksum: + raise ValueError( + "adaptive XGrammar local-compile marker checksum is invalid" + ) + magic, native_compile_nanoseconds, serialized_bytes = struct.unpack( + cls._LOCAL_COMPILE_BODY_FORMAT, + body, + ) + if ( + magic != cls._LOCAL_COMPILE_MAGIC + or native_compile_nanoseconds < 0 + or serialized_bytes <= 0 + ): + raise ValueError( + "adaptive XGrammar local-compile marker values are invalid" + ) + return _LocalCompileMarker( + native_compile_nanoseconds=native_compile_nanoseconds, + serialized_bytes=serialized_bytes, + ) + + @staticmethod + def _write_atomic_bytes(path: Path, payload: bytes) -> None: + # Cache data is reconstructible. Close + replace gives readers + # process-level atomicity, while the per-launch session scan reconciles + # accounting after a crash. Per-entry fsync (without a directory fsync) + # did not provide complete power-loss durability and reduced burst + # compile throughput by more than an order of magnitude. + temp_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=".xgrammar-", + suffix=".tmp", + delete=False, + ) as temp_file: + temp_path = Path(temp_file.name) + os.chmod(temp_path, 0o600) + if temp_file.write(payload) != len(payload): + raise RuntimeError( + f"short write while creating XGrammar cache entry: {path}" + ) + os.replace(temp_path, path) + finally: + if temp_path is not None and temp_path.exists(): + temp_path.unlink() + + @classmethod + def _write_atomic_text(cls, path: Path, text: str) -> None: + cls._write_atomic_bytes(path, text.encode("utf-8")) + + def _initialize_size_ledger(self) -> None: + with self._prune_lock.open("a+b") as prune_lock: + fcntl.flock(prune_lock.fileno(), fcntl.LOCK_EX) + session_id = os.environ.get("SGLANG_XGRAMMAR_CACHE_SESSION_ID") + if ( + session_id + and self._accounting_session.is_file() + and self._accounting_session.read_text(encoding="ascii") == session_id + ): + if self._size_ledger.is_file(): + self._read_size_ledger() + return + + total_bytes, entries = self._scan_entries() + total_bytes = self._prune_entries( + total_bytes=total_bytes, + entries=entries, + protected=None, + ) + self._reset_size_ledger(total_bytes) + if session_id: + try: + session_id.encode("ascii") + except UnicodeEncodeError as exc: + raise ValueError( + "SGLANG_XGRAMMAR_CACHE_SESSION_ID must be ASCII" + ) from exc + self._write_atomic_text(self._accounting_session, session_id) + + def _account_and_prune( + self, + *, + protected: Path, + added_bytes: int, + ) -> None: + with self._prune_lock.open("a+b") as prune_lock: + fcntl.flock(prune_lock.fileno(), fcntl.LOCK_EX) + total_bytes = self._read_size_ledger() + added_bytes + + if total_bytes <= self.max_bytes: + self._update_size_ledger(total_bytes) + return + total_bytes, entries = self._scan_entries() + total_bytes = self._prune_entries( + total_bytes=total_bytes, + entries=entries, + protected=protected, + ) + if total_bytes > self.max_bytes: + raise RuntimeError( + "persistent XGrammar cache could not prune below its hard " + f"limit: {total_bytes} > {self.max_bytes} bytes" + ) + self._update_size_ledger(total_bytes) + + def _read_size_ledger(self) -> int: + try: + raw = self._size_ledger.read_bytes() + if len(raw) != self._COUNTER_BYTES: + raise ValueError( + f"expected {self._COUNTER_BYTES} bytes, got {len(raw)}" + ) + (value,) = struct.unpack(self._COUNTER_FORMAT, raw) + except Exception as exc: + raise RuntimeError( + f"persistent XGrammar size ledger is invalid: {self._size_ledger}" + ) from exc + return value + + def _reset_size_ledger(self, value: int) -> None: + payload = self._counter_payload(value) + temp_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=self._size_ledger.parent, + prefix=".xgrammar-counter-", + suffix=".tmp", + delete=False, + ) as temp_file: + temp_path = Path(temp_file.name) + os.chmod(temp_path, 0o600) + if temp_file.write(payload) != len(payload): + raise RuntimeError("short write while initializing size ledger") + os.replace(temp_path, self._size_ledger) + finally: + if temp_path is not None and temp_path.exists(): + temp_path.unlink() + + def _update_size_ledger(self, value: int) -> None: + payload = self._counter_payload(value) + try: + with self._size_ledger.open("r+b", buffering=0) as ledger: + written = os.pwrite(ledger.fileno(), payload, 0) + except Exception as exc: + raise RuntimeError( + f"failed to update persistent XGrammar size ledger: {self._size_ledger}" + ) from exc + if written != len(payload): + raise RuntimeError( + "short write while updating persistent XGrammar size ledger: " + f"{written} != {len(payload)}" + ) + + def _counter_payload(self, value: int) -> bytes: + if value < 0 or value >= 1 << (8 * self._COUNTER_BYTES): + raise RuntimeError( + f"persistent XGrammar size ledger value is out of range: {value}" + ) + return struct.pack(self._COUNTER_FORMAT, value) + + def _prune_entries( + self, + *, + total_bytes: int, + entries: list[tuple[int, int, Path]], + protected: Path | None, + ) -> int: + for _, size, path in sorted(entries): + if total_bytes <= self.max_bytes: + break + if path == protected: + continue + if self._unlink_if_idle(path): + total_bytes -= size + return total_bytes + + def _scan_entries(self) -> tuple[int, list[tuple[int, int, Path]]]: + entries: list[tuple[int, int, Path]] = [] + total_bytes = 0 + for item in self.entries.iterdir(): + if not item.is_file() or item.name.startswith(".xgrammar-"): + continue + stat = item.stat() + total_bytes += stat.st_size + entries.append((stat.st_mtime_ns, stat.st_size, item)) + return total_bytes, entries + + def _unlink_if_idle(self, entry: Path) -> bool: + lock_path = self.locks / f"{entry.stem}.lock" + with lock_path.open("a+b") as entry_lock: + try: + fcntl.flock( + entry_lock.fileno(), + fcntl.LOCK_EX | fcntl.LOCK_NB, + ) + except BlockingIOError: + return False + if not entry.exists(): + return False + entry.unlink() + return True diff --git a/python/sglang/srt/distributed/device_communicators/vmm_utils.py b/python/sglang/srt/distributed/device_communicators/vmm_utils.py index a286e9401a83..a5730d506699 100644 --- a/python/sglang/srt/distributed/device_communicators/vmm_utils.py +++ b/python/sglang/srt/distributed/device_communicators/vmm_utils.py @@ -175,53 +175,34 @@ def _recv_fd(sock): def export_shareable_handles(retained_handles, group: ProcessGroup, rank: int): - """Export retained VMM handles, preferring FABRIC and falling back to POSIX fds. - - FABRIC is used only if every rank can export it; otherwise all ranks use POSIX - fds. Returns ``(fabric_handles, posix_fds, use_fabric)`` (one list populated); - raises if both fail on any rank. Caller owns the returned ``posix_fds``. + """Export retained VMM handles through the strict same-host POSIX fd path. + + This topology has both ranks in containers on one Linux host. + CUDA FABRIC handles require an IMEX channel and only transport allocation + handles; they do not alter the peer mapping after import. Select the native + same-host SCM_RIGHTS path directly instead of probing another transport and + falling back. Returns ``([], posix_fds, False)`` and raises if export fails + on any rank. The caller owns the returned ``posix_fds``. """ drv = _get_cuda_driver() - FABRIC = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC POSIX_FD = drv.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR - fabric_handles: List[bytes] = [] - fabric_error: Optional[Exception] = None - try: - for alloc_h in retained_handles: - fabric_h = check_drv( - drv.cuMemExportToShareableHandle(alloc_h, FABRIC, 0), - "cuMemExportToShareableHandle(FABRIC)", - ) - fabric_handles.append(bytes(fabric_h.data)) - fabric_ok = True - except Exception as e: - fabric_error = e - fabric_ok = False - fabric_handles = [] - logger.info( - "FABRIC handle export failed on rank %s; falling back to " - "POSIX fd transport: %s", - rank, - e, - ) - - if all_ranks_ok(group, fabric_ok): - return fabric_handles, [], True - posix_fds: List[int] = [] posix_error: Optional[Exception] = None try: for alloc_h in retained_handles: fd = check_drv( drv.cuMemExportToShareableHandle(alloc_h, POSIX_FD, 0), - "cuMemExportToShareableHandle(POSIX_FD)", + "strict single-node cuMemExportToShareableHandle(POSIX_FD)", ) posix_fds.append(int(fd)) posix_ok = True except Exception as e: posix_error = e posix_ok = False + + all_posix_ok = all_ranks_ok(group, posix_ok) + if not all_posix_ok: for fd in posix_fds: try: os.close(fd) @@ -229,14 +210,11 @@ def export_shareable_handles(retained_handles, group: ProcessGroup, rank: int): pass posix_fds = [] - if not all_ranks_ok(group, posix_ok): - cause = posix_error or fabric_error - message = ( - "VMM handle export failed: FABRIC export failed on at least one " - "rank and POSIX fd export failed on at least one rank" - ) - if cause is not None: - message += f"; local rank {rank} error: {cause}" + message = "strict single-node VMM POSIX fd export failed on at least one rank" + if posix_error is not None: + message += f"; local rank {rank} error: {posix_error}" + else: + message += f"; local rank {rank} export succeeded but a peer failed" raise RuntimeError(message) from posix_error return [], posix_fds, False diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py index fd216d329ebb..a0bc09486540 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py @@ -188,6 +188,8 @@ # Global constants HEALTH_CHECK_TIMEOUT = int(os.getenv("SGLANG_HEALTH_CHECK_TIMEOUT", 20)) WAIT_WEIGHTS_READY_TIMEOUT = int(os.getenv("SGLANG_WAIT_WEIGHTS_READY_TIMEOUT", 120)) +SCHEDULER_HEALTH_MAX_STALENESS_SECONDS = 30.0 +SCHEDULER_HEALTH_MAX_FUTURE_SKEW_SECONDS = 5.0 # Store global states @@ -613,6 +615,63 @@ async def validate_json_request(raw_request: Request): ##### Native API endpoints ##### +@app.get("/health_scheduler") +async def health_scheduler() -> Response: + """Check every DP scheduler without submitting an inference request.""" + + global_state = _global_state + if global_state is None: + return Response(status_code=503) + + tokenizer_manager = global_state.tokenizer_manager + if tokenizer_manager.gracefully_exit: + logger.info("Scheduler health check received during shutdown. Returning 503.") + return Response(status_code=503) + + if tokenizer_manager.server_status == ServerStatus.Starting: + return Response(status_code=503) + + try: + snapshots = await tokenizer_manager.get_loads(include=["core"]) + except Exception: + logger.exception("Scheduler health check could not read load snapshots") + return Response(status_code=503) + + expected_ranks = set(range(tokenizer_manager.elastic_worker_count)) + actual_ranks = [snapshot.dp_rank for snapshot in snapshots] + if len(actual_ranks) != len(expected_ranks) or set(actual_ranks) != expected_ranks: + logger.warning( + "Scheduler health check found DP ranks %s; expected %s", + sorted(actual_ranks), + sorted(expected_ranks), + ) + return Response(status_code=503) + + now = time.time() + for snapshot in snapshots: + if snapshot.timestamp <= 0: + logger.warning( + "Scheduler health check found an invalid timestamp for DP rank %s", + snapshot.dp_rank, + ) + return Response(status_code=503) + + age = now - snapshot.timestamp + if ( + age > SCHEDULER_HEALTH_MAX_STALENESS_SECONDS + or age < -SCHEDULER_HEALTH_MAX_FUTURE_SKEW_SECONDS + ): + logger.warning( + "Scheduler health check found an out-of-range snapshot age " + "for DP rank %s: %.3fs", + snapshot.dp_rank, + age, + ) + return Response(status_code=503) + + return Response(status_code=200) + + @app.get("/health") @app.get("/health_generate") async def health_generate(request: Request) -> Response: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b4ff61057eb6..d277019ea152 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1050,6 +1050,7 @@ class Envs: SGLANG_FP8_PAGED_MQA_LOGITS_TORCH = EnvBool(False) SGLANG_TOPK_TRANSFORM_512_TORCH = EnvBool(False) SGLANG_OPT_FLASHMLA_SPARSE_PREFILL = EnvBool(True) + SGLANG_PREFILL_CUDA_GRAPH_EAGER_VALIDATION = EnvBool(False) # SWA radix cache # TODO(DSV4): @ispobock this has bug on main branch when retract diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index fadc8c1a46e7..1af1b3bf6ab6 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -55,6 +55,7 @@ ) from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( SparsePrefillChunkCache, + SparsePrefillOutputWorkspace, SparsePrefillWorkspace, ) from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool @@ -276,6 +277,47 @@ def refresh_for_breakable_cuda_graph_replay_(self, other: DSV4AttnMetadata) -> N for field_name in reference_assign_fields: setattr(self, field_name, getattr(other, field_name)) + def refresh_for_full_prefill_cuda_graph_replay_( + self, other: DSV4AttnMetadata + ) -> None: + assert self.c4_sparse_topk == other.c4_sparse_topk + assert self.page_size == other.page_size + assert self.cuda_int32_kwargs == other.cuda_int32_kwargs + + # A whole-forward graph captures every metadata tensor by address. + # Preserve those tensor owners and refresh only their contents. In + # particular, do not replace FlashMLASchedMeta: its lazily allocated + # scheduler tensors were created during capture and are graph inputs. + tensor_copy_fields = [ + "raw_out_loc", + "seq_lens_casual", + "positions_casual", + "c4_out_loc", + "c128_out_loc", + "page_table", + "swa_page_indices", + "swa_topk_lengths", + "c128_page_indices", + "c128_topk_lengths_clamp1", + "c4_topk_lengths_raw", + "c4_topk_lengths_clamp1", + "c4_sparse_topk_lengths", + "c4_sparse_page_indices", + "c4_sparse_raw_indices", + ] + for field_name in tensor_copy_fields: + src_val = getattr(other, field_name) + dst_val = getattr(self, field_name) + if src_val is None and dst_val is None: + continue + assert ( + src_val is not None and dst_val is not None + ), f"{field_name=} {src_val=} {dst_val=}" + assert ( + src_val.shape == dst_val.shape + ), f"{field_name=} {src_val.shape=} {dst_val.shape=}" + dst_val.copy_(src_val) + def init_compression_metadata(self): assert self.page_table.dim() == 2 assert ( @@ -424,6 +466,23 @@ def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata ) self.sparse_prefill_cache = None + def refresh_for_full_prefill_cuda_graph_replay_( + self, static_metadata: DSV4Metadata + ) -> None: + self.core_attn_metadata.refresh_for_full_prefill_cuda_graph_replay_( + static_metadata.core_attn_metadata + ) + maybe_copy_inplace(self.indexer_metadata, src=static_metadata.indexer_metadata) + maybe_copy_inplace( + self.c4_compress_metadata, src=static_metadata.c4_compress_metadata + ) + maybe_copy_inplace( + self.c128_compress_metadata, src=static_metadata.c128_compress_metadata + ) + # Sparse-prefill's lazy workspace is deliberately excluded from full + # graph capture; losing such an owner would be a use-after-free. + assert self.sparse_prefill_cache is None + @dataclass class DSV4RawVerifyMetadata: @@ -439,6 +498,15 @@ class DSV4RawVerifyMetadata: verify_lens: Optional[torch.Tensor] = None total_verify_tokens: int = 0 + # init_forward_metadata_in_graph materializes the full DSV4 metadata while + # CUDA capture is active. The graph retains only device addresses, not the + # Python tensor owners. Keep the captured object reachable from the raw + # per-shape metadata for the graph's entire lifetime; otherwise allocator + # reuse can turn c4_out_loc and the compressor plans into dangling pointers. + _captured_full_metadata: Optional[DSV4Metadata] = field( + default=None, repr=False, compare=False + ) + def copy_(self, other: DSV4RawVerifyMetadata): self.req_pool_indices.copy_(other.req_pool_indices) self.seq_lens.copy_(other.seq_lens) @@ -461,6 +529,13 @@ class DSV4RawDecodeMetadata: seq_lens: torch.Tensor out_cache_loc: torch.Tensor + # See DSV4RawVerifyMetadata._captured_full_metadata. Decode graphs create + # the same address-captured C4/C128 metadata and require identical lifetime + # ownership. + _captured_full_metadata: Optional[DSV4Metadata] = field( + default=None, repr=False, compare=False + ) + def copy_(self, other: DSV4RawDecodeMetadata): self.req_pool_indices.copy_(other.req_pool_indices) self.seq_lens.copy_(other.seq_lens) @@ -471,6 +546,7 @@ class _GraphBucket(enum.Enum): DECODE_OR_IDLE = "decode_or_idle" TARGET_VERIFY = "target_verify" DRAFT_EXTEND = "draft_extend" + PREFILL = "prefill" @classmethod def of(cls, forward_mode: ForwardMode) -> _GraphBucket: @@ -480,6 +556,8 @@ def of(cls, forward_mode: ForwardMode) -> _GraphBucket: return cls.TARGET_VERIFY if forward_mode.is_draft_extend_v2(): return cls.DRAFT_EXTEND + if forward_mode.is_extend_without_speculative(): + return cls.PREFILL raise NotImplementedError(f"unsupported {forward_mode=}") @@ -489,6 +567,17 @@ class DeepseekV4AttnBackend( use_captured_forward_metadata_for_breakable_cuda_graph: bool = True supports_ragged_verify_graph: bool = True needs_cpu_seq_lens: bool = False + # DSV4's measured useful TBO path is dense prefill. Decode and DSpark + # target verification retain their existing graph layouts (decode TBO + # regresses and compact-ragged verification has a different token axis). + tbo_supports_cuda_graph: bool = True + tbo_supports_decode_cuda_graph: bool = False + tbo_requires_decode_cuda_graph_state: bool = False + tbo_requires_global_cpu_seq_lens: bool = False + + @staticmethod + def tbo_supports_cuda_graph_for(forward_mode: ForwardMode) -> bool: + return forward_mode.is_extend_without_speculative() def __init__( self, @@ -563,6 +652,28 @@ def __init__( self.online_c128_mtp = OnlineC128MTPController(self) self.sparse_prefill_workspace = SparsePrefillWorkspace(self.device) spec_alg = model_runner.spec_algorithm + self.is_draft_runner = model_runner.is_draft_worker + self.sparse_prefill_output_workspace: Optional[SparsePrefillOutputWorkspace] = ( + None + ) + if not self.is_draft_runner: + local_num_heads = model_runner.model_config.get_num_attention_heads( + get_parallel().attn_tp_size + ) + self.sparse_prefill_output_workspace = SparsePrefillOutputWorkspace( + device=self.device, + capacity_tokens=model_runner.server_args.max_prefill_tokens, + num_heads=local_num_heads, + head_dim_v=self.head_dim_v, + ) + logger.info( + "Reserved fixed FlashMLA sparse-prefill outputs: " + "tokens=%d heads=%d d_v=%d bytes=%d", + self.sparse_prefill_output_workspace.capacity_tokens, + local_num_heads, + self.head_dim_v, + self.sparse_prefill_output_workspace.reserved_bytes, + ) self.needs_cpu_seq_lens = not spec_alg.is_dspark() and ( not _is_cuda or not envs.SGLANG_PREP_IN_CUDA_GRAPH.get() @@ -570,7 +681,6 @@ def __init__( ) self.is_dspark_draft = model_runner.is_draft_worker and spec_alg.is_dspark() - self.is_draft_runner = model_runner.is_draft_worker self.cuda_graph_custom_mask = None def _move_to_device(self, x: List[int]) -> torch.Tensor: @@ -1096,14 +1206,23 @@ def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None: # materialization is recorded inside the cuda graph; a no-op (Full # already) when PREP_IN_CUDA_GRAPH=0. if isinstance(self.forward_metadata, DSV4RawVerifyMetadata): - self.forward_metadata = self.make_forward_metadata_from_raw_verify( - raw_metadata=self.forward_metadata, + raw_metadata = self.forward_metadata + full_metadata = self.make_forward_metadata_from_raw_verify( + raw_metadata=raw_metadata, online_c128_state_slot_offset=self.online_c128_mtp.state_slot_offset(), ) + # A CUDA graph owns addresses, not Python tensor lifetimes. Pin all + # dynamically materialized output-location and compressor-plan + # tensors on the persistent raw metadata object for this graph key. + raw_metadata._captured_full_metadata = full_metadata + self.forward_metadata = full_metadata elif isinstance(self.forward_metadata, DSV4RawDecodeMetadata): - self.forward_metadata = self.make_forward_metadata_from_raw_decode( - raw_metadata=self.forward_metadata, + raw_metadata = self.forward_metadata + full_metadata = self.make_forward_metadata_from_raw_decode( + raw_metadata=raw_metadata, ) + raw_metadata._captured_full_metadata = full_metadata + self.forward_metadata = full_metadata # Compute the SWA KV-store write target once per forward and cache it on # the metadata for every layer's store. This is recorded inside the cuda @@ -1169,6 +1288,10 @@ def init_forward_metadata_out_graph( forward_batch: ForwardBatch, in_capture: bool = False, ) -> None: + # Prefill graphs are captured before the decode runner initializes its + # phase-specific CUDA-graph state. Create the shared metadata registry + # at its first use and preserve it when decode initialization follows. + self._ensure_cuda_graph_metadata_registry() bucket = _GraphBucket.of(forward_batch.forward_mode) bs = forward_batch.batch_size req_pool_indices = forward_batch.req_pool_indices @@ -1329,6 +1452,33 @@ def init_forward_metadata_out_graph( num_tokens_per_req=num_tokens_per_req, out_cache_loc=out_cache_loc, ) + elif bucket == _GraphBucket.PREFILL: + # Full prefill graphs are keyed by both their fixed request slots + # and padded token bucket. Build graph-compatible metadata against + # those stable buffers, then refresh it in place before replay. + # Large DSV4 prefill uses a different sparse-attention path whose + # lazily built workspace is not replay-stable under a whole-forward + # graph; keep that optimized path eager instead of capturing an + # invalid graph. + num_tokens = forward_batch.positions.numel() + if ( + num_tokens > _LARGE_INDEXER_QUERY_THRESHOLD + or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get() + ): + raise RuntimeError( + "DSV4 full prefill CUDA graphs require the non-sparse " + f"attention path, got {num_tokens=} and " + "SGLANG_OPT_FLASHMLA_SPARSE_PREFILL=" + f"{envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get()}" + ) + assert forward_batch.out_cache_loc is not None + assert forward_batch.out_cache_loc.numel() == num_tokens + graph_key = (bs, num_tokens) + temp_metadata = self._build_forward_metadata( + forward_batch, + max_seq_len_override=chosen_max_seq_len, + use_prefill_cuda_graph=True, + ) else: self.online_c128_mtp.clear() raise NotImplementedError @@ -1494,11 +1644,13 @@ def prepare_forward_metadata_for_breakable_cuda_graph_replay( capture_metadata.refresh_for_breakable_cuda_graph_replay_(static_metadata) self.forward_metadata = capture_metadata - def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: + def _ensure_cuda_graph_metadata_registry(self) -> None: + if hasattr(self, "cuda_graph_metadata_of_bucket_and_bs"): + return self.cuda_graph_metadata_of_bucket_and_bs: Dict[ _GraphBucket, Dict[ - int, + Union[int, Tuple[int, int]], Union[ DSV4Metadata, DSV4RawDecodeMetadata, @@ -1506,6 +1658,9 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: ], ], ] = {bucket: {} for bucket in _GraphBucket} + + def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: + self._ensure_cuda_graph_metadata_registry() self.draft_extend_num_tokens_per_req = ( max_num_tokens // max_bs if max_bs > 0 else 1 ) @@ -1526,7 +1681,7 @@ def get_verify_buffers_to_fill_after_draft(self): def replay_cuda_graph_metadata_from( self, - bs: int, + bs: Union[int, Tuple[int, int]], temp_metadata: Union[ DSV4Metadata, DSV4RawVerifyMetadata, @@ -1540,7 +1695,12 @@ def replay_cuda_graph_metadata_from( bucket_metadata[bs] = temp_metadata self.forward_metadata = temp_metadata return - chosen_metadata.copy_(temp_metadata) + if bucket == _GraphBucket.PREFILL: + assert isinstance(chosen_metadata, DSV4Metadata) + assert isinstance(temp_metadata, DSV4Metadata) + chosen_metadata.refresh_for_full_prefill_cuda_graph_replay_(temp_metadata) + else: + chosen_metadata.copy_(temp_metadata) self.forward_metadata = chosen_metadata def get_cuda_graph_seq_len_fill_value(self): @@ -1859,6 +2019,16 @@ def _forward_prefill_sparse( ) kv = workspace + if self.sparse_prefill_output_workspace is None: + raise RuntimeError( + "DSV4 sparse prefill reached a draft worker; the fixed target " + "output workspace is mandatory" + ) + out, max_logits, lse = self.sparse_prefill_output_workspace.get( + num_tokens=q_flat.shape[0], + num_heads=q_flat.shape[1], + head_dim_v=self.head_dim_v, + ) o, _, _ = flash_mla_sparse_fwd( q=q_flat, kv=kv, @@ -1867,6 +2037,9 @@ def _forward_prefill_sparse( d_v=self.head_dim_v, attn_sink=attn_sink, topk_length=combined_lens, + out=out, + max_logits=max_logits, + lse=lse, ) return o diff --git a/python/sglang/srt/layers/attention/dsv4/metadata.py b/python/sglang/srt/layers/attention/dsv4/metadata.py index d245ddce3c01..53e4125a35c8 100644 --- a/python/sglang/srt/layers/attention/dsv4/metadata.py +++ b/python/sglang/srt/layers/attention/dsv4/metadata.py @@ -48,7 +48,13 @@ c4_sparse: means "compressed by 4" but only attend to top-512 tokens. all related length will be clipped to 512. """ -_LARGE_INDEXER_QUERY_THRESHOLD = 11673 +# FlashMLA's dense scheduling-metadata kernel reserves +# sizeof(int) * (5 * num_queries + 1) bytes of dynamic shared memory. B300 +# exposes 232448 bytes per block, so its hard ceiling is 11622 queries. Keep +# the dense/graph range on a 64-token-aligned 11616 boundary; larger prefills +# use DSV4's dedicated SM100 sparse-attention kernel instead of FlashMLA's +# serial low-shared-memory metadata path. +_LARGE_INDEXER_QUERY_THRESHOLD = 11616 def copy_metadata( diff --git a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py index 931b791eeed4..79dc55d77fd1 100644 --- a/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py +++ b/python/sglang/srt/layers/attention/dsv4/sparse_prefill_utils.py @@ -80,6 +80,78 @@ def get(self, num_tokens: int) -> torch.Tensor: return self._buffer[:num_tokens] +class SparsePrefillOutputWorkspace: + """Fixed, long-lived FlashMLA sparse-prefill outputs. + + A full DSV4 TP2 prefill chunk produces a roughly 1 GiB BF16 attention + output. Allocating it inside every FlashMLA call makes runtime success + depend on a fresh contiguous block being available after variable decode + batches have fragmented the allocator. Reserve the exact maximum before + CUDA-graph capture and reuse it for every eager sparse-prefill layer. + """ + + def __init__( + self, + *, + device: torch.device, + capacity_tokens: int, + num_heads: int, + head_dim_v: int, + ) -> None: + assert capacity_tokens > 0 + assert num_heads > 0 + assert head_dim_v > 0 + self.capacity_tokens = capacity_tokens + self.num_heads = num_heads + self.head_dim_v = head_dim_v + self.out = torch.empty( + (capacity_tokens, num_heads, head_dim_v), + dtype=torch.bfloat16, + device=device, + ) + self.max_logits = torch.empty( + (capacity_tokens, num_heads), + dtype=torch.float32, + device=device, + ) + self.lse = torch.empty( + (capacity_tokens, num_heads), + dtype=torch.float32, + device=device, + ) + + @property + def reserved_bytes(self) -> int: + return sum( + tensor.numel() * tensor.element_size() + for tensor in (self.out, self.max_logits, self.lse) + ) + + def get( + self, + *, + num_tokens: int, + num_heads: int, + head_dim_v: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if num_tokens <= 0 or num_tokens > self.capacity_tokens: + raise RuntimeError( + "sparse-prefill output exceeds its fixed workspace capacity: " + f"tokens={num_tokens}, capacity={self.capacity_tokens}" + ) + if num_heads != self.num_heads or head_dim_v != self.head_dim_v: + raise RuntimeError( + "sparse-prefill output geometry changed after reservation: " + f"got heads={num_heads}, d_v={head_dim_v}; expected " + f"heads={self.num_heads}, d_v={self.head_dim_v}" + ) + return ( + self.out[:num_tokens], + self.max_logits[:num_tokens], + self.lse[:num_tokens], + ) + + def combined_topk_width(topk: int, window_size: int) -> int: """Width of the padded combined_indices last dim that ``combine_topk_swa_indices`` would produce for these args.""" diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index d595f352031e..9c3a501c7dd1 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -1,6 +1,8 @@ from types import SimpleNamespace from typing import TYPE_CHECKING, Callable, List +import torch + from sglang.srt.batch_overlap import two_batch_overlap from sglang.srt.layers.attention.base_attn_backend import AttentionBackend @@ -17,6 +19,11 @@ def __init__(self, primary: AttentionBackend, children: List[AttentionBackend]): # reads through TboAttnBackend resolve to the underlying pool. self.token_to_kv_pool = primary.token_to_kv_pool self.req_to_token_pool = primary.req_to_token_pool + # These are ABC class attributes, so __getattr__ cannot delegate them. + # Mirror the primary explicitly: otherwise the wrapper silently rejects + # captured ragged verify and forces an unnecessary decode seq-lens D2H. + self.supports_ragged_verify_graph = primary.supports_ragged_verify_graph + self.needs_cpu_seq_lens = primary.needs_cpu_seq_lens @classmethod def init_new(cls, creator: Callable[[], AttentionBackend]): @@ -25,7 +32,7 @@ def init_new(cls, creator: Callable[[], AttentionBackend]): children=[creator() for _ in range(2)], ) - def _children_use_cuda_graph(self) -> bool: + def _children_use_cuda_graph(self, forward_batch=None) -> bool: """Whether the TBO child backends participate in CUDA-graph capture/replay. Some models only run TBO in eager prefill and keep their graph-captured @@ -37,6 +44,10 @@ def _children_use_cuda_graph(self) -> bool: HSA_STATUS_ERROR_OUT_OF_RESOURCES. Eager prefill TBO (init_forward_metadata) is unaffected; only the *_graph paths are gated. """ + if forward_batch is not None: + supports_mode = getattr(self.primary, "tbo_supports_cuda_graph_for", None) + if supports_mode is not None: + return supports_mode(forward_batch.forward_mode) return getattr(self.primary, "tbo_supports_cuda_graph", True) def init_forward_metadata_out_graph( @@ -44,12 +55,20 @@ def init_forward_metadata_out_graph( forward_batch: "ForwardBatch", in_capture: bool = False, ): - self.primary.init_forward_metadata_out_graph( - forward_batch=forward_batch, in_capture=in_capture + children_supported = self._children_use_cuda_graph(forward_batch) + tbo_children = getattr(forward_batch, "tbo_children", None) + use_children = children_supported and ( + tbo_children is not None + or ( + not in_capture + and forward_batch.forward_mode.is_extend_without_speculative() + ) ) - if not self._children_use_cuda_graph(): + if not use_children: + self.primary.init_forward_metadata_out_graph( + forward_batch=forward_batch, in_capture=in_capture + ) return - tbo_children = getattr(forward_batch, "tbo_children", None) if tbo_children is not None: for child, forward_batch_child in zip( self.children, tbo_children, strict=True @@ -67,6 +86,10 @@ def init_forward_metadata_out_graph( self._dispatch_children_from_replay_view(forward_batch) def _dispatch_children_from_replay_view(self, fb_view) -> None: + if fb_view.forward_mode.is_extend_without_speculative(): + self._dispatch_prefill_children_from_replay_view(fb_view) + return + bs = fb_view.batch_size forward_mode = fb_view.forward_mode spec_info = fb_view.spec_info @@ -111,32 +134,60 @@ def _dispatch_children_from_replay_view(self, fb_view) -> None: forward_batch=child_fb_view, in_capture=False ) + def _dispatch_prefill_children_from_replay_view(self, fb_view) -> None: + """Refresh fixed-shape child metadata for a full prefill graph replay.""" + num_tokens = fb_view.positions.numel() + split_token_index = num_tokens // 2 + child_layouts = _split_prefill_replay_layout( + extend_seq_lens=fb_view.extend_seq_lens_cpu, + extend_prefix_lens=fb_view.extend_prefix_lens_cpu, + split_token_index=split_token_index, + ) + for child_backend, layout, tok_slice in zip( + self.children, + child_layouts, + (slice(None, split_token_index), slice(split_token_index, None)), + strict=True, + ): + child_fb_view = _build_tbo_prefill_child_replay_fb_view( + fb_view, + layout=layout, + tok_slice=tok_slice, + ) + child_backend.init_forward_metadata_out_graph( + forward_batch=child_fb_view, + in_capture=False, + ) + def init_forward_metadata_in_graph(self, forward_batch: "ForwardBatch"): - self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch) - if not self._children_use_cuda_graph(): - return tbo_children = getattr(forward_batch, "tbo_children", None) - if tbo_children is not None: - for child, forward_batch_child in zip( - self.children, tbo_children, strict=True - ): - if forward_batch_child.batch_size > 0: - child.init_forward_metadata_in_graph( - forward_batch=forward_batch_child - ) + use_children = ( + self._children_use_cuda_graph(forward_batch) and tbo_children is not None + ) + if not use_children: + self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch) + return + for child, forward_batch_child in zip(self.children, tbo_children, strict=True): + if forward_batch_child.batch_size > 0: + child.init_forward_metadata_in_graph(forward_batch=forward_batch_child) def init_forward_metadata(self, forward_batch: "ForwardBatch"): - self.primary.init_forward_metadata(forward_batch=forward_batch) - if forward_batch.tbo_children is not None: - for child, forward_batch_child in zip( - self.children, forward_batch.tbo_children, strict=True - ): - if forward_batch_child.batch_size > 0: - child.init_forward_metadata(forward_batch=forward_batch_child) + if forward_batch.tbo_children is None: + self.primary.init_forward_metadata(forward_batch=forward_batch) + return + for child, forward_batch_child in zip( + self.children, forward_batch.tbo_children, strict=True + ): + if forward_batch_child.batch_size > 0: + child.init_forward_metadata(forward_batch=forward_batch_child) def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): self.primary.init_cuda_graph_state(max_bs=max_bs, max_num_tokens=max_num_tokens) - if not self._children_use_cuda_graph(): + if not getattr( + self.primary, + "tbo_requires_decode_cuda_graph_state", + self._children_use_cuda_graph(), + ): return for item in self.children: # TODO for children, maybe can provide *smaller* max_bs to optimize @@ -239,3 +290,141 @@ def _build_tbo_child_replay_fb_view( ), spec_info=child_spec_info, ) + + +def _split_prefill_replay_layout( + *, + extend_seq_lens, + extend_prefix_lens, + split_token_index: int, +): + """Split live request metadata at a fixed flat-token graph boundary. + + Both returned layouts keep the parent's fixed request-slot count. A + request that has no query tokens in a child becomes a true zero sentinel; + a request straddling the boundary keeps the original prefix in child A and + advances child B's prefix by A's query-token contribution. + """ + extend_seq_lens = [int(value) for value in extend_seq_lens] + extend_prefix_lens = [int(value) for value in extend_prefix_lens] + if len(extend_seq_lens) != len(extend_prefix_lens): + raise RuntimeError( + "TBO prefill replay metadata has mismatched request axes: " + f"extend={len(extend_seq_lens)}, prefix={len(extend_prefix_lens)}" + ) + if split_token_index < 0: + raise RuntimeError( + f"TBO prefill replay split must be non-negative: {split_token_index}" + ) + + child_extend = [[], []] + child_prefix = [[], []] + child_seq = [[], []] + remaining_left = split_token_index + for extend_len, prefix_len in zip(extend_seq_lens, extend_prefix_lens, strict=True): + if extend_len < 0 or prefix_len < 0: + raise RuntimeError( + "TBO prefill replay lengths must be non-negative: " + f"extend={extend_len}, prefix={prefix_len}" + ) + left_len = min(extend_len, max(remaining_left, 0)) + right_len = extend_len - left_len + remaining_left -= left_len + + child_extend[0].append(left_len) + child_prefix[0].append(prefix_len if left_len else 0) + child_seq[0].append(prefix_len + left_len if left_len else 0) + + child_extend[1].append(right_len) + child_prefix[1].append(prefix_len + left_len if right_len else 0) + child_seq[1].append(prefix_len + extend_len if right_len else 0) + + if remaining_left < 0: + raise RuntimeError( + "TBO prefill replay split consumed more tokens than its boundary" + ) + + layouts = [] + for extend_lens, prefix_lens, seq_lens in zip( + child_extend, child_prefix, child_seq, strict=True + ): + start_locs = [] + cursor = 0 + for extend_len in extend_lens: + start_locs.append(cursor) + cursor += extend_len + layouts.append( + SimpleNamespace( + extend_seq_lens=extend_lens, + extend_prefix_lens=prefix_lens, + seq_lens=seq_lens, + extend_start_loc=start_locs, + extend_num_tokens=cursor, + ) + ) + return layouts + + +def _build_tbo_prefill_child_replay_fb_view( + fb_view, + *, + layout, + tok_slice: slice, +) -> SimpleNamespace: + device = fb_view.seq_lens.device + seq_dtype = fb_view.seq_lens.dtype + req_pool_indices = fb_view.req_pool_indices.clone() + active = torch.tensor( + [value > 0 for value in layout.extend_seq_lens], + dtype=torch.bool, + device=device, + ) + req_pool_indices.masked_fill_(~active, 0) + seq_lens = torch.tensor(layout.seq_lens, dtype=seq_dtype, device=device) + extend_seq_lens = torch.tensor( + layout.extend_seq_lens, dtype=fb_view.extend_seq_lens.dtype, device=device + ) + extend_prefix_lens = torch.tensor( + layout.extend_prefix_lens, + dtype=fb_view.extend_prefix_lens.dtype, + device=device, + ) + extend_start_loc = torch.tensor( + layout.extend_start_loc, + dtype=fb_view.extend_start_loc.dtype, + device=device, + ) + seq_lens_cpu = torch.tensor(layout.seq_lens, dtype=torch.int64, device="cpu") + parent_input_ids = getattr(fb_view, "input_ids", None) + parent_out_cache_loc = getattr(fb_view, "out_cache_loc", None) + parent_positions = getattr(fb_view, "positions", None) + return SimpleNamespace( + batch_size=fb_view.batch_size, + forward_mode=fb_view.forward_mode, + actual_forward_mode=getattr( + fb_view, "actual_forward_mode", fb_view.forward_mode + ), + input_ids=( + parent_input_ids[tok_slice] if parent_input_ids is not None else None + ), + positions=( + parent_positions[tok_slice] if parent_positions is not None else None + ), + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_sum=sum(layout.seq_lens), + seq_lens_cpu=seq_lens_cpu, + encoder_lens=None, + out_cache_loc=( + parent_out_cache_loc[tok_slice] + if parent_out_cache_loc is not None + else None + ), + extend_num_tokens=layout.extend_num_tokens, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=layout.extend_seq_lens, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=layout.extend_prefix_lens, + extend_start_loc=extend_start_loc, + spec_info=None, + ) diff --git a/python/sglang/srt/managers/data_parallel_controller.py b/python/sglang/srt/managers/data_parallel_controller.py index c43a4fc1bf06..ad362b038740 100644 --- a/python/sglang/srt/managers/data_parallel_controller.py +++ b/python/sglang/srt/managers/data_parallel_controller.py @@ -35,6 +35,7 @@ BlockReqInput, ElasticScaleUpdateReq, ProfileReq, + ShutdownReq, TokenizedEmbeddingReqInput, TokenizedGenerateReqInput, sock_recv, @@ -208,6 +209,7 @@ def __init__( self.control_message_step = 1 self.init_dispatcher() + self.gracefully_exit = False self.soft_watchdog = Watchdog.create( debug_name="DataParallelController", @@ -327,15 +329,57 @@ def dispatching_with_trace(self, req: Req, refresh_load_budget: bool = True): req.time_stats = time_stats req.time_stats.set_dp_dispatch_finish_time() + def _dispatch_atomic_routed_batch(self, batch_req): + """Keep an explicitly routed HTTP batch atomic through DP dispatch.""" + if len(batch_req) == 0: + return False + + routed_dp_rank = batch_req[0].routed_dp_rank + if routed_dp_rank is None or any( + req.routed_dp_rank != routed_dp_rank for req in batch_req + ): + return False + if ( + routed_dp_rank < 0 + or routed_dp_rank >= len(self.workers) + or routed_dp_rank not in self._active_workers + or self.workers[routed_dp_rank] is None + ): + raise ValueError(f"DP rank {routed_dp_rank} is not active.") + + time_stats = [] + for req in batch_req: + req_time_stats = DPControllerReqTimeStats.new_from_obj( + unwrap_from_pickle(req.time_stats) + ) + req_time_stats.set_dp_dispatch_time() + req.time_stats = wrap_as_pickle(req_time_stats) + time_stats.append(req_time_stats) + + # One ZMQ message means the scheduler receives and queues the entire + # batch before it can form the next prefill. Sending each item as its + # own message lets the scheduler wake after an arbitrary prefix, + # fragmenting a 64-request API batch into smaller GPU launches. + sock_send(self.workers[routed_dp_rank], batch_req) + + for req, req_time_stats in zip(batch_req, time_stats): + req.time_stats = req_time_stats + req.time_stats.set_dp_dispatch_finish_time() + return True + def dispatch_batch_generate(self, batch_req: BatchTokenizedGenerateReqInput): if self.refresh_load_budget_on_dispatch: self.refresh_load_budget() + if self._dispatch_atomic_routed_batch(batch_req): + return for req in batch_req: self.dispatching_with_trace(req, refresh_load_budget=False) def dispatch_batch_embedding(self, batch_req: BatchTokenizedEmbeddingReqInput): if self.refresh_load_budget_on_dispatch: self.refresh_load_budget() + if self._dispatch_atomic_routed_batch(batch_req): + return for req in batch_req: self.dispatching_with_trace(req, refresh_load_budget=False) @@ -348,6 +392,7 @@ def init_dispatcher(self): (BatchTokenizedEmbeddingReqInput, self.dispatch_batch_embedding), (BlockReqInput, self.send_to_all_workers), (ProfileReq, self.send_to_all_workers), + (ShutdownReq, self.handle_shutdown), (ActiveRanksOutput, self.update_active_ranks), ( ElasticScaleUpdateReq, @@ -359,6 +404,11 @@ def init_dispatcher(self): ) self._request_dispatcher.add_fallback_fn(self.send_control_message) + def handle_shutdown(self, recv_req: ShutdownReq) -> None: + logger.info("DataParallelController received graceful shutdown request.") + self.send_control_message(recv_req) + self.gracefully_exit = True + def launch_dp_schedulers(self, server_args, port_args): base_gpu_id = 0 @@ -794,7 +844,7 @@ def total_tokens_scheduler(self, req: Req): sock_send(self.workers[target_worker], req) def event_loop(self): - while True: + while not self.gracefully_exit: while True: self.soft_watchdog.feed() try: @@ -849,9 +899,12 @@ def run_data_parallel_controller_process( controller.event_loop() for proc in controller.scheduler_procs: proc.join() - logger.error( - f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}" - ) + if proc.exitcode != 0: + raise RuntimeError( + f"Scheduler {proc.pid} exited with status {proc.exitcode}" + ) + logger.info(f"Scheduler {proc.pid} exited cleanly.") + logger.info("DataParallelController graceful shutdown complete.") except Exception: traceback = get_exception_traceback() logger.error(f"DataParallelController hit an exception: {traceback}") diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py index 7b3ca52e5640..c5300f83c175 100644 --- a/python/sglang/srt/managers/detokenizer_manager.py +++ b/python/sglang/srt/managers/detokenizer_manager.py @@ -34,6 +34,7 @@ BatchTokenIDOutput, ConfigureLoggingReq, FreezeGCReq, + ShutdownReq, sock_recv, sock_send, ) @@ -142,6 +143,7 @@ def init_running_status(self, server_args: ServerArgs): self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES) self.disable_tokenizer_batch_decode = server_args.disable_tokenizer_batch_decode self.is_tool_call_parser_gpt_oss = server_args.tool_call_parser == "gpt-oss" + self.gracefully_exit = False self.soft_watchdog = Watchdog.create( debug_name="DetokenizerManager", @@ -159,13 +161,14 @@ def init_request_dispatcher(self): (BatchEmbeddingOutput, self.handle_batch_embedding_out), (BatchTokenIDOutput, self.handle_batch_token_id_out), (FreezeGCReq, self.handle_freeze_gc_req), + (ShutdownReq, self.handle_shutdown_req), (ConfigureLoggingReq, self.handle_configure_logging_req), ] ) def event_loop(self): """The event loop that handles requests""" - while True: + while not self.gracefully_exit: with self.soft_watchdog.disable(): recv_obj = sock_recv(self.recv_from_scheduler) output = self._request_dispatcher(recv_obj) @@ -173,6 +176,10 @@ def event_loop(self): sock_send(self.send_to_tokenizer, output) self.soft_watchdog.feed() + def handle_shutdown_req(self, recv_req: ShutdownReq): + logger.info("Detokenizer received graceful shutdown request.") + self.gracefully_exit = True + def trim_matched_stop( self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool ): diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 9e2f191aa2ff..7a6a180c0d43 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -33,8 +33,15 @@ def decide_needs_cpu_seq_lens( # importable everywhere; spec_info pulls in the spec/schedule_batch graph. from sglang.srt.speculative.spec_info import SpeculativeAlgorithm - if server_args.enable_two_batch_overlap: - # FIXME: support TBO without seq lens cpu value + if server_args.enable_two_batch_overlap and any( + getattr(backend, "tbo_requires_global_cpu_seq_lens", True) + for backend in attn_backends + if backend is not None + ): + # Legacy TBO models split decode-family batches from a host mirror. + # Phase-aware implementations such as DSV4 prefill-only TBO can opt + # out so their decode/speculative loop keeps the GPU-only seq-lens + # relay. return True algo = SpeculativeAlgorithm.from_string(server_args.speculative_algorithm) if algo.is_ngram(): diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f3e240c1a3ac..f6d705dfbdf3 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -995,6 +995,7 @@ def __init__( None ) self.grammar_wait_ct = 0 + self.grammar_stats_logged = False # The number of cached tokens that were already cached in the KV cache self.cached_tokens = 0 @@ -1149,6 +1150,18 @@ def finished(self) -> bool: # Whether request reached finished condition return self.finished_reason is not None + def log_grammar_stats_once(self, metrics_collector) -> None: + if self.grammar_stats_logged or metrics_collector is None: + return + grammar = self.grammar + if grammar is None or isinstance(grammar, Future): + return + grammar_stats = grammar.grammar_stats + if grammar_stats is None: + return + metrics_collector.log_grammar_stats(grammar_stats) + self.grammar_stats_logged = True + def set_extend_range(self, start: int, end: int) -> None: self.extend_range = Range(start, end) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index c5ac9a6e7a20..34e46f451b21 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -67,7 +67,10 @@ prepare_abort, ) from sglang.srt.distributed import get_pp_group, get_world_group -from sglang.srt.distributed.parallel_state import get_tp_group +from sglang.srt.distributed.parallel_state import ( + cleanup_dist_env_and_memory, + get_tp_group, +) from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin from sglang.srt.environ import envs @@ -228,6 +231,7 @@ validate_input_length, ) from sglang.srt.mem_cache import kv_cache_builder +from sglang.srt.mem_cache.base_prefix_cache import EvictParams from sglang.srt.mem_cache.common import maybe_cache_unfinished_req, release_kv_cache from sglang.srt.model_executor.forward_batch_info import PPProxyTensors from sglang.srt.model_loader.utils import get_resolved_model_impl @@ -378,6 +382,16 @@ def __init__( ) self.page_size = server_args.page_size self.enable_hierarchical_cache = server_args.enable_hierarchical_cache + hicache_gate_control = os.environ.get("SGLANG_HICACHE_GATE_CONTROL", "disabled") + if hicache_gate_control not in {"disabled", "armed"}: + raise ValueError( + "SGLANG_HICACHE_GATE_CONTROL must be 'disabled' or 'armed', " + f"got {hicache_gate_control!r}" + ) + self.hicache_gate_control_armed = hicache_gate_control == "armed" + self.hicache_gate_last_requested_tokens = 0 + self.hicache_gate_last_evictable_tokens = 0 + self.hicache_gate_last_evicted_tokens = 0 self.enable_hicache_storage = server_args.hicache_storage_backend is not None self.enable_decode_hicache = ( server_args.disaggregation_decode_enable_radix_cache @@ -425,6 +439,8 @@ def __init__( # Init metrics stats self.init_metrics_collector(tp_rank, pp_rank, dp_rank) + self._scheduler_phase_accumulator: dict[str, tuple[float, int, float]] = {} + self._scheduler_phase_flush_deadline = time.perf_counter() + 1.0 # Init inter-process communication self.init_ipc_channels(port_args) @@ -807,6 +823,8 @@ def maybe_init_draft_worker(self): nccl_port=self.nccl_port, target_worker=self.tp_worker, ) + if self.spec_algorithm.is_dspark(): + draft_worker_kwargs["metrics_collector"] = self.metrics_collector if self.server_args.speculative_draft_load_format is not None: # Write the draft load_format onto server_args (not just the bag): @@ -1582,6 +1600,46 @@ def event_loop_normal(self): if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.invariant_checker.self_check_during_busy() + def _finish_scheduler_phase(self, phase: str, started: float) -> float: + """Record one latency-critical scheduler phase and return the next start.""" + now = time.perf_counter() + duration = now - started + if self.server_args.enable_metrics: + total, calls, maximum = self._scheduler_phase_accumulator.get( + phase, (0.0, 0, 0.0) + ) + self._scheduler_phase_accumulator[phase] = ( + total + duration, + calls + 1, + max(maximum, duration), + ) + if phase == "loop_total" and now >= self._scheduler_phase_flush_deadline: + for ( + accumulated_phase, + ( + accumulated_seconds, + accumulated_calls, + accumulated_max, + ), + ) in self._scheduler_phase_accumulator.items(): + self.metrics_collector.add_scheduler_phase( + accumulated_phase, + accumulated_seconds, + accumulated_calls, + accumulated_max, + ) + self._scheduler_phase_accumulator.clear() + self._scheduler_phase_flush_deadline = now + 1.0 + if duration >= 1.0: + logger.warning( + "Slow scheduler phase: phase=%s duration=%.3fs running=%d waiting=%d", + phase, + duration, + len(self.running_batch.reqs), + len(self.waiting_queue), + ) + return now + @DynamicGradMode() def event_loop_overlap(self): """A scheduler loop that overlaps the CPU processing and GPU computation.""" @@ -1598,10 +1656,16 @@ def pop_and_process(): if self.gracefully_exit: break + loop_started = phase_started = time.perf_counter() + # Receive requests recv_reqs = self.request_receiver.recv_requests() self.process_input_requests(recv_reqs) + phase_started = self._finish_scheduler_phase( + "request_ingest", phase_started + ) if self._engine_paused: + self._finish_scheduler_phase("loop_total", loop_started) continue # Get the next batch to run @@ -1614,6 +1678,7 @@ def pop_and_process(): disable_overlap_for_batch = self.is_disable_overlap_for_batch( batch, last_batch=self.last_batch ) + phase_started = self._finish_scheduler_phase("schedule_plan", phase_started) # If we do not need to overlap the current batch with the last batch, # we can process the last batch immediately. @@ -1627,6 +1692,7 @@ def pop_and_process(): self.token_to_kv_pool_allocator.flush_opportunistic() except Exception: pass + phase_started = self._finish_scheduler_phase("result_sync", phase_started) # Launch the current batch if batch: @@ -1636,6 +1702,9 @@ def pop_and_process(): self.result_queue.append((batch.copy(), batch_result)) else: batch_result = None + phase_started = self._finish_scheduler_phase( + "forward_submit", phase_started + ) # Process the last batch if self.last_batch: @@ -1644,17 +1713,22 @@ def pop_and_process(): elif batch is None: # When the server is idle, do self-check and re-init some states self.on_idle() + phase_started = self._finish_scheduler_phase( + "result_process", phase_started + ) # Run sample of the current batch # It depends on the result of the last batch (e.g., grammar), so we run it after the last batch is processed. if self.is_generation: self.launch_batch_sample_if_needed(batch_result, batch) + phase_started = self._finish_scheduler_phase("sample", phase_started) # Update last_batch self.last_batch = batch if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): self.invariant_checker.self_check_during_busy() + self._finish_scheduler_phase("loop_total", loop_started) def is_disable_overlap_for_batch( self, batch: ScheduleBatch, last_batch: Optional[ScheduleBatch] @@ -1701,8 +1775,15 @@ def _advance_pending_grammar(self): (before generate_token_bitmask) so the CPU advance overlaps the target verify forward. Idempotent; no-op when the queue is empty or has no grammar. """ + phase_seconds = {} for prev_batch, prev_result in self.result_queue: - self.batch_result_processor.advance_grammar_fsm(prev_result, prev_batch) + result_phase_seconds = self.batch_result_processor.advance_grammar_fsm( + prev_result, prev_batch + ) + if result_phase_seconds is not None: + for phase, seconds in result_phase_seconds.items(): + phase_seconds[phase] = phase_seconds.get(phase, 0.0) + seconds + return phase_seconds @scheduler_nvtx_method("scheduler.process_input_requests") def process_input_requests(self, recv_reqs: List): @@ -2253,7 +2334,11 @@ def handle_generate_request( self._maybe_namespace_elastic_radix_cache(req) if self.spec_algorithm.is_dflash_family(): - error_msg = validate_dflash_request(req, self.enable_overlap) + error_msg = validate_dflash_request( + req, + self.enable_overlap, + self.spec_algorithm, + ) if error_msg is not None: req.set_finish_with_abort(error_msg) self.init_req_max_new_tokens(req) @@ -2798,6 +2883,23 @@ def get_next_batch_to_run( # We need to discard it. chunked_req_to_exclude.add(last_batch.chunked_req) + # A one-token DSpark logprob request is completed by this target + # prefill. Under overlap scheduling its result is still queued, so + # merging the live batch into running_batch would launch a stale + # speculative decode before the prefill result retires the request. + # The result queue owns a shallow batch snapshot and still streams + # the sampled token and exact target logprobs. Exclude only final + # prefill chunks; middle chunks must continue until prefill ends. + if ( + self.spec_algorithm.is_dspark() + and last_batch.contains_last_prefill_chunk + ): + chunked_req_to_exclude.update( + req + for req in last_batch.reqs + if req.return_logprob and req.sampling_params.max_new_tokens == 1 + ) + if self.dllm_config is not None and last_batch.reqs: chunked_req_to_exclude.update(last_batch.reqs) @@ -2825,12 +2927,14 @@ def get_next_batch_to_run( if running_batch.is_empty(): running_batch.batch_is_full = False + plan_phase_started = time.perf_counter() if self.dllm_config is not None: new_batch = self.get_new_batch_dllm(running_batch) else: prefill_plan = self.get_new_batch_prefill(running_batch) new_batch = prefill_plan.batch_to_run running_batch = prefill_plan.running_batch + self._finish_scheduler_phase("plan_prefill", plan_phase_started) need_mlp_sync = self.require_mlp_sync if ( @@ -2842,7 +2946,9 @@ def get_next_batch_to_run( # Before merging the new batch into running batch: # 1. All new batches are none -> need_mlp_sync remains true (sync is needed for decode batch). # 2. All new batches are some (prefill / idle) -> we do not need prepare mlp sync one more time. + plan_phase_started = time.perf_counter() new_batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(new_batch) + self._finish_scheduler_phase("plan_mlp_sync_prefill", plan_phase_started) need_mlp_sync = new_batch is None if new_batch is not None: @@ -2851,20 +2957,26 @@ def get_next_batch_to_run( else: # Run decode (skip for prefill-only batches) if not running_batch.is_empty() and not running_batch.is_prefill_only: + plan_phase_started = time.perf_counter() running_batch = self.update_running_batch(running_batch) + self._finish_scheduler_phase("plan_decode_update", plan_phase_started) ret = running_batch if not running_batch.is_empty() else None else: ret = None # Handle DP attention and log stats + plan_phase_started = time.perf_counter() ret = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch( ret, need_sync=need_mlp_sync ) + self._finish_scheduler_phase("plan_mlp_sync_final", plan_phase_started) # Handle ngram embedding + plan_phase_started = time.perf_counter() ret = self.ngram_embedding_manager.prepare_for_forward( ret, chunked_req=self.chunked_req ) + self._finish_scheduler_phase("plan_ngram", plan_phase_started) if ret: set_schedule_time_batch(ret) @@ -3998,6 +4110,12 @@ def get_internal_state(self, recv_req: GetInternalStateReq): / self.metrics_reporter.spec_total_num_forward_ct ) + ret["hicache_gate_control"] = { + "armed": self.hicache_gate_control_armed, + "last_requested_tokens": self.hicache_gate_last_requested_tokens, + "last_evictable_tokens": self.hicache_gate_last_evictable_tokens, + "last_evicted_tokens": self.hicache_gate_last_evicted_tokens, + } if RECORD_STEP_TIME: ret["step_time_dict"] = self.metrics_reporter.step_time_dict @@ -4022,6 +4140,7 @@ def set_internal_state(self, recv_req: SetInternalStateReq): "speculative_accept_threshold_acc", "dspark_force_budget_frac", "dspark_clear_info_records", + "hicache_gate_evict_device_tokens", ] ) @@ -4063,6 +4182,37 @@ def set_internal_state(self, recv_req: SetInternalStateReq): ) if_success = False break + elif k == "hicache_gate_evict_device_tokens": + if not self.hicache_gate_control_armed: + logging.warning( + "hicache_gate_evict_device_tokens rejected: the one-shot " + "startup control is not armed." + ) + if_success = False + break + if isinstance(v, bool) or not isinstance(v, int) or v <= 0: + logging.warning( + "hicache_gate_evict_device_tokens must be a positive integer, " + f"got {v!r}." + ) + if_success = False + break + if ( + not self.enable_hierarchical_cache + or getattr(self.tree_cache, "cache_controller", None) is None + ): + logging.warning( + "hicache_gate_evict_device_tokens requires an active " + "hierarchical cache." + ) + if_success = False + break + if not self.is_fully_idle(): + logging.warning( + "hicache_gate_evict_device_tokens requires an idle scheduler." + ) + if_success = False + break if if_success: if ( @@ -4079,7 +4229,54 @@ def set_internal_state(self, recv_req: SetInternalStateReq): ) = 0 # DSpark control keys are worker commands, not server args; route # them to the draft worker and keep them out of the override. + # The HiCache startup-gate command is likewise not a server arg. + # Disarm it before touching cache state: a failed operation is a + # fail-closed startup error, never permission to retry through a + # different or degraded path. remaining = dict(server_args_dict) + hicache_gate_evict_tokens = remaining.pop( + "hicache_gate_evict_device_tokens", None + ) + if hicache_gate_evict_tokens is not None: + self.hicache_gate_control_armed = False + self.hicache_gate_last_requested_tokens = int(hicache_gate_evict_tokens) + try: + self.tree_cache.writing_check(write_back=True) + self.hicache_gate_last_evictable_tokens = int( + self.tree_cache.evictable_size() + ) + # Selective write-through may intentionally leave cold + # device leaves unbacked. The isolation proof must never + # delete and silently recompute those leaves: temporarily + # force the existing write-back eviction path so every + # evicted leaf is preserved on host first. + cache_controller = self.tree_cache.cache_controller + original_write_policy = cache_controller.write_policy + try: + cache_controller.write_policy = "write_back" + evict_result = self.tree_cache.evict( + EvictParams(num_tokens=int(hicache_gate_evict_tokens)) + ) + finally: + cache_controller.write_policy = original_write_policy + self.hicache_gate_last_evicted_tokens = int( + evict_result.num_tokens_evicted + ) + if self.hicache_gate_last_evicted_tokens <= 0: + raise RuntimeError( + "one-shot HiCache gate control did not evict any " + "device tokens" + ) + logger.info( + "HiCache startup gate device eviction completed: " + "requested=%d evictable=%d evicted=%d", + self.hicache_gate_last_requested_tokens, + self.hicache_gate_last_evictable_tokens, + self.hicache_gate_last_evicted_tokens, + ) + except Exception: + logger.exception("HiCache startup gate device eviction failed.") + if_success = False frac = remaining.pop("dspark_force_budget_frac", None) if "dspark_force_budget_frac" in server_args_dict: self.draft_worker.set_dspark_forced_budget_frac( @@ -4527,10 +4724,15 @@ def maybe_sleep_on_idle(self): def handle_freeze_gc(self, recv_req: FreezeGCReq): """Handle freeze_gc request: freeze scheduler's GC and forward to detokenizer.""" freeze_gc("Scheduler") + if self.server_args.enable_metrics: + self.metrics_collector.set_runtime_gc_frozen() self.ipc_channels.send_to_detokenizer.send_output(recv_req, recv_req) return None def handle_shutdown(self, recv_req: ShutdownReq): + # The tokenizer waits for every tracked child, so let the detokenizer + # leave its blocking receive loop before this scheduler tears down. + self.ipc_channels.send_to_detokenizer.send_output(recv_req, recv_req) # Break the event loop; the finally in run_scheduler_process releases resources. self.gracefully_exit = True return None @@ -4751,4 +4953,17 @@ def run_scheduler_process( # Graceful path only: on the exception path the GPU may be wedged # and the synchronize() in destroy() could itself hang. if scheduler.gracefully_exit: + logger.info("Graceful scheduler teardown: releasing host resources.") scheduler.release_host_resources() + # Host-buffer release times differ slightly between ranks. + # Rendezvous on the CPU group before destroying every tracked + # subgroup and the default process group, otherwise the first + # rank to exit tears down TCPStore while its peer's NCCL + # heartbeat is still active. + logger.info("Graceful scheduler teardown: waiting for peer rank.") + get_world_group().barrier() + logger.info( + "Graceful scheduler teardown: destroying distributed groups." + ) + cleanup_dist_env_and_memory() + logger.info("Graceful scheduler teardown complete.") diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index 248a92939e28..56be3370b84e 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -1,10 +1,12 @@ from __future__ import annotations import logging +import time from dataclasses import dataclass from typing import ( TYPE_CHECKING, Callable, + Dict, List, Optional, Tuple, @@ -95,6 +97,7 @@ def process_batch_result_prebuilt(self, batch: ScheduleBatch): if self.server_args.enable_hisparse: self.hisparse_coordinator.request_finished(req) release_kv_cache(req, self.tree_cache) + req.log_grammar_stats_once(self.metrics_collector) # Note: Logprobs should be handled on the prefill engine. self.output_streamer.stream_output(batch.reqs, batch.return_logprob) @@ -278,6 +281,8 @@ def process_batch_result_prefill( next_token_id=next_token_id, already_advanced=result.grammar_advanced, ) + if req.finished(): + req.log_grammar_stats_once(self.metrics_collector) else: # being chunked reqs' prefill is not finished @@ -653,7 +658,7 @@ def _accept_grammar_tokens( def advance_grammar_fsm( self, result: GenerationBatchResult, batch: ScheduleBatch - ) -> None: + ) -> Optional[Dict[str, float]]: """Advance each req's grammar FSM over the tokens THIS batch committed, and (for decode) memoize the grammar-truncated run on ``result``. @@ -666,15 +671,39 @@ def advance_grammar_fsm( batch was — e.g. the extend->decode boundary. """ if result.grammar_advanced or not batch.has_grammar: - return + return None is_decode = batch.forward_mode.is_decode() if not (is_decode or batch.forward_mode.is_extend()): - return - if result.copy_done is not None: - result.copy_done.synchronize() - next_token_ids = result.next_token_ids.tolist() + return None + total_started = time.perf_counter() + copy_wait_started = time.perf_counter() + prefetch_phase_seconds = {} + if batch.spec_algorithm.is_dspark(): + if result.grammar_result_future is None: + raise RuntimeError( + "DSpark grammar overlap requires the asynchronous result path" + ) + prepared = result.grammar_result_future.result() + copy_wait_seconds = time.perf_counter() - copy_wait_started + next_token_ids = prepared.next_token_ids + tensor_to_list_seconds = 0.0 + prefetch_phase_seconds = { + "barrier_prefetch_queue_delay": prepared.queue_delay_seconds, + "barrier_prefetch_copy_wait": prepared.copy_wait_seconds, + "barrier_prefetch_tensor_to_list": (prepared.tensor_to_list_seconds), + "barrier_prefetch_submit_to_ready": (prepared.submit_to_ready_seconds), + } + else: + if result.copy_done is not None: + result.copy_done.synchronize() + next_token_ids_tensor = result.next_token_ids + copy_wait_seconds = time.perf_counter() - copy_wait_started + tensor_to_list_started = time.perf_counter() + next_token_ids = next_token_ids_tensor.tolist() + tensor_to_list_seconds = time.perf_counter() - tensor_to_list_started if not is_decode: + fsm_accept_started = time.perf_counter() # Extend: advance over the single token each completed-prefill req emitted # (mirrors process_batch_result_prefill's per-req token indexing). for i, req in enumerate(batch.reqs): @@ -687,18 +716,33 @@ def advance_grammar_fsm( continue self._accept_grammar_tokens(req, next_token_ids[i]) result.grammar_advanced = True - return + return { + "barrier_result_copy_wait": copy_wait_seconds, + "barrier_tensor_to_list": tensor_to_list_seconds, + **prefetch_phase_seconds, + "barrier_fsm_accept": time.perf_counter() - fsm_accept_started, + "barrier_advance_total": time.perf_counter() - total_started, + } # Decode: only the spec-v2 path reaches here (the grammar barrier for # spec-overlap workers and _resolve_spec_v2_tokens). Non-spec grammar decode # advances its FSM in process_batch_result_decode and has no accept_lens, so # bail out defensively. if result.accept_lens is None: - return - accept_lens = result.accept_lens.tolist() + return None + if batch.spec_algorithm.is_dspark(): + if prepared.accept_lens is None: + raise RuntimeError( + "DSpark grammar decode requires asynchronous accept lengths" + ) + accept_lens = prepared.accept_lens + else: + accept_lens = result.accept_lens.tolist() + tensor_to_list_seconds = time.perf_counter() - tensor_to_list_started stride = result.speculative_num_draft_tokens assert stride is not None, "spec-v2 result missing speculative_num_draft_tokens" retained = [None] * len(batch.reqs) + fsm_accept_started = time.perf_counter() for i, req in enumerate(batch.reqs): if req.grammar is None or req.is_retracted or req.finished(): continue @@ -708,6 +752,13 @@ def advance_grammar_fsm( retained[i] = self._accept_grammar_tokens(req, accept_tokens) result.grammar_retained_tokens = retained result.grammar_advanced = True + return { + "barrier_result_copy_wait": copy_wait_seconds, + "barrier_tensor_to_list": tensor_to_list_seconds, + **prefetch_phase_seconds, + "barrier_fsm_accept": time.perf_counter() - fsm_accept_started, + "barrier_advance_total": time.perf_counter() - total_started, + } def process_batch_result_idle( self, @@ -820,6 +871,8 @@ def process_batch_result_decode( # here; spec already advanced it in _resolve_spec_v2_tokens. self._accept_grammar_tokens(req, next_token_id) req.grammar.finished = req.finished() + if req.finished(): + req.log_grammar_stats_once(self.metrics_collector) self.output_streamer.stream_output(batch.reqs, batch.return_logprob) self.token_to_kv_pool_allocator.free_group_end() diff --git a/python/sglang/srt/managers/scheduler_components/dp_attn.py b/python/sglang/srt/managers/scheduler_components/dp_attn.py index 01a1d2adb504..b92b999be87d 100644 --- a/python/sglang/srt/managers/scheduler_components/dp_attn.py +++ b/python/sglang/srt/managers/scheduler_components/dp_attn.py @@ -1,5 +1,6 @@ from __future__ import annotations +from bisect import bisect_left from dataclasses import dataclass from typing import TYPE_CHECKING, Callable, Optional @@ -15,13 +16,15 @@ from sglang.srt.managers.scheduler_components.recv_skipper import ( SchedulerRecvSkipper, ) +from sglang.srt.managers.scheduler_components.single_node_dp2_sync import ( + exchange_single_node_dp2_mlp_info, + single_node_dp2_sync_enabled, +) 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 ( Backend, - Phase, - check_cuda_graph_backend, cuda_graph_fully_disabled, ) from sglang.srt.model_executor.forward_batch_info import ForwardMode @@ -35,6 +38,59 @@ _ENABLE_METRICS_DP_ATTENTION = envs.SGLANG_ENABLE_METRICS_DP_ATTENTION.get() +_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR = 4 + + +def _can_run_dp_prefill_cuda_graph_locally( + local_batch: Optional[ScheduleBatch], + server_args: ServerArgs, +) -> bool: + """Return this rank's shape-independent prefill graph eligibility. + + The result is min-reduced by the DP MLP sync. Idle ranks are permissive: + once any peer contributes an admitted prefill, MAX_LEN padding converts + them to the same captured EXTEND shape. + """ + prefill_config = server_args.cuda_graph_config.prefill + if prefill_config.backend not in (Backend.BREAKABLE, Backend.FULL): + return False + if local_batch is None or local_batch.forward_mode.is_idle(): + return True + if local_batch.forward_mode not in (ForwardMode.EXTEND, ForwardMode.MIXED): + return False + if local_batch.input_embeds is not None or local_batch.replace_embeds is not None: + return False + if prefill_config.backend == Backend.FULL: + max_requests = prefill_config.full_prefill_max_req + if max_requests is None: + max_requests = max(server_args.chunked_prefill_size // 512, 1) + if local_batch.batch_size() > max_requests: + return False + return True + + +def _finalize_dp_prefill_cuda_graph_admission( + local_admission: bool, + *, + global_num_tokens: list[int], + is_extend_in_batch: bool, + server_args: ServerArgs, +) -> bool: + """Apply the shared token-bucket crossover to a DP-reduced admission.""" + if not local_admission or not is_extend_in_batch: + return False + + capture_tokens = server_args.cuda_graph_config.prefill.bs + if not capture_tokens: + return False + num_tokens = max(global_num_tokens) + if num_tokens <= 0: + return False + bucket_index = bisect_left(capture_tokens, num_tokens) + if bucket_index == len(capture_tokens): + return False + padded_num_tokens = capture_tokens[bucket_index] + return padded_num_tokens <= num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR def _resolve_elastic_world_dp_size( @@ -90,8 +146,11 @@ class MLPSyncBatchInfo: # some gathered elements tp0_info: torch.Tensor = None + global_forward_modes: list[int] = None + used_native_sync: bool = False global_num_tokens: list[int] = None global_num_tokens_for_logprob: list[int] = None + global_local_can_run_tbo: list[int] = None tbo_split_seq_index: torch.Tensor = None global_forward_mode: int = None dp_cooperation_info: Optional[DPCooperationInfo] = None @@ -132,6 +191,42 @@ def all_gather( group: torch.distributed.ProcessGroup, use_all_reduce: bool = False, ): + if single_node_dp2_sync_enabled(): + if use_all_reduce or device != "cpu": + raise RuntimeError( + "strict DP2 shared-memory MLP sync requires the fixed CPU " + "all-gather path" + ) + global_info = exchange_single_node_dp2_mlp_info( + self.num_tokens, + self.num_tokens_for_logprob, + self.can_cuda_graph, + self.is_extend_in_batch, + self.local_can_run_tbo, + self.local_forward_mode, + self.can_run_breakable_cuda_graph, + group=group, + dp_size=self.dp_size, + tp_size=self.tp_size, + cp_size=self.cp_size, + ) + self.used_native_sync = True + self.global_num_tokens = [global_info[0], global_info[7]] + self.global_num_tokens_for_logprob = [ + global_info[1], + global_info[8], + ] + self.can_cuda_graph = bool(global_info[2] and global_info[9]) + self.is_extend_in_batch = bool(global_info[3] or global_info[10]) + self.global_forward_modes = [global_info[5], global_info[12]] + self.global_local_can_run_tbo = [global_info[4], global_info[11]] + self.can_run_breakable_cuda_graph = bool(global_info[6] and global_info[13]) + if _ENABLE_METRICS_DP_ATTENTION: + self.dp_cooperation_info = DPCooperationInfo.create( + self.global_forward_modes + ) + return + local_info_tensor = self._get_local_tensor(device=device) fallback_tensor = self._get_fallback_tensor(device=device) info_width = local_info_tensor.numel() @@ -226,6 +321,7 @@ def prepare_mlp_sync_batch_raw( require_mlp_tp_gather: bool, disable_overlap_schedule: bool, offload_tags: set[str], + server_args: ServerArgs, dwdp: bool = False, ): # Check if other DP workers have running batches @@ -260,14 +356,12 @@ 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 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) + # Historical field name: this is now the rank-local admission for either + # captured prefill backend. The seven-field DP exchange min-reduces it + # before the shared token-bucket decision below. + can_run_breakable_cuda_graph = _can_run_dp_prefill_cuda_graph_locally( + local_batch, server_args + ) is_extend_in_batch = local_batch.forward_mode.is_extend() if local_batch else False if local_batch is not None: @@ -320,11 +414,33 @@ def prepare_mlp_sync_batch_raw( use_all_reduce=use_world_group, ) - mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = ( - tbo_preparer.compute_output( - mlp_sync_info.tp0_info[:, 4:6], + if mlp_sync_info.used_native_sync: + ( + mlp_sync_info.tbo_split_seq_index, + mlp_sync_info.global_forward_mode, + ) = tbo_preparer.compute_output_from_values( + mlp_sync_info.global_local_can_run_tbo, + mlp_sync_info.global_forward_modes, + ) + else: + mlp_sync_info.tbo_split_seq_index, mlp_sync_info.global_forward_mode = ( + tbo_preparer.compute_output( + mlp_sync_info.tp0_info[:, 4:6], + ) ) + + mlp_sync_info.can_run_breakable_cuda_graph = ( + _finalize_dp_prefill_cuda_graph_admission( + mlp_sync_info.can_run_breakable_cuda_graph, + global_num_tokens=( + mlp_sync_info.global_num_tokens + if not skip_all_gather + else [mlp_sync_info.num_tokens] + ), + is_extend_in_batch=mlp_sync_info.is_extend_in_batch, + server_args=server_args, ) + ) # Decide whether to emit idle batch if skip_all_gather: @@ -349,10 +465,11 @@ def prepare_mlp_sync_batch_raw( # Set on `local_batch`, not `batch_to_gather`: for PREBUILT batches the # scheduler's `last_batch` is the prebuilt batch, not its inner idle batch. if local_batch is not None and not skip_all_gather: + forward_modes = mlp_sync_info.global_forward_modes + if forward_modes is None: + forward_modes = mlp_sync_info.tp0_info[:, 5].tolist() local_batch.recv_skipper_forward_mode = ( - SchedulerRecvSkipper.derive_forward_mode( - mlp_sync_info.tp0_info[:, 5].tolist() - ) + SchedulerRecvSkipper.derive_forward_mode(forward_modes) ) if _ENABLE_METRICS_DP_ATTENTION and local_batch is not None: @@ -387,6 +504,7 @@ def prepare_mlp_sync_batch(self, local_batch: ScheduleBatch): require_mlp_tp_gather=require_mlp_tp_gather(self.server_args), disable_overlap_schedule=self.server_args.disable_overlap_schedule, offload_tags=self.offload_tags, + server_args=self.server_args, dwdp=self.server_args.dwdp_size > 1, ) diff --git a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py index 15e1b814d738..04370ebf8bb7 100644 --- a/python/sglang/srt/managers/scheduler_components/metrics_reporter.py +++ b/python/sglang/srt/managers/scheduler_components/metrics_reporter.py @@ -2,8 +2,10 @@ import dataclasses import logging +import os import tempfile import time +from bisect import bisect_left from collections import defaultdict from dataclasses import dataclass from typing import ( @@ -18,6 +20,7 @@ from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.observability.metrics_collector import ( DPCooperationInfo, QueueCount, @@ -26,6 +29,7 @@ SchedulerStats, compute_routing_key_stats, ) +from sglang.srt.server_args import ServerArgs from sglang.srt.utils.device_timer import DeviceTimer from sglang.srt.utils.scheduler_status_logger import SchedulerStatusLogger @@ -42,6 +46,7 @@ RECORD_STEP_TIME = envs.SGLANG_RECORD_STEP_TIME.get() LOG_FORWARD_ITERS = envs.SGLANG_LOG_FORWARD_ITERS.get() ENABLE_METRICS_DEVICE_TIMER = envs.SGLANG_ENABLE_METRICS_DEVICE_TIMER.get() +_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR = 4 def _decode_total_seq_lens(batch: ScheduleBatch) -> int: @@ -86,6 +91,109 @@ def from_adder( ) +@dataclass(frozen=True, slots=True) +class PrefillExecutionObservation: + outcome: str + scheduled_tokens: int + executed_tokens: int + requests: int + bucket_tokens: int + + +def classify_prefill_execution( + *, + batch: ScheduleBatch, + can_run_cuda_graph: bool, + server_args: ServerArgs, + dp_rank: int, +) -> PrefillExecutionObservation: + """Classify one local prefill without synchronizing the hot path.""" + prefill_config = server_args.cuda_graph_config.prefill + capture_tokens = prefill_config.bs + global_num_tokens = batch.global_num_tokens + if global_num_tokens is not None and dp_rank < len(global_num_tokens): + scheduled_tokens = max(int(global_num_tokens[dp_rank]), 0) + else: + # Non-DP callers do not carry the shared geometry. Keep their metrics + # useful without pretending a page-rounded log counter is exact. + scheduled_tokens = max(int(getattr(batch, "extend_num_tokens", 0) or 0), 0) + requests = batch.batch_size() + global_tokens = ( + max(int(tokens) for tokens in global_num_tokens) + if global_num_tokens + else scheduled_tokens + ) + bucket_index = bisect_left(capture_tokens, global_tokens) if capture_tokens else 0 + bucket_tokens = ( + capture_tokens[bucket_index] + if capture_tokens and bucket_index < len(capture_tokens) + else 0 + ) + + if can_run_cuda_graph: + if bucket_tokens <= 0: + raise RuntimeError( + "prefill CUDA graph executed without a captured token bucket" + ) + return PrefillExecutionObservation( + outcome="cuda_graph", + scheduled_tokens=scheduled_tokens, + executed_tokens=bucket_tokens, + requests=requests, + bucket_tokens=bucket_tokens, + ) + + if prefill_config.backend not in (Backend.BREAKABLE, Backend.FULL): + outcome = "backend_disabled" + elif not capture_tokens: + outcome = "no_capture_buckets" + elif ( + prefill_config.backend == Backend.FULL + and bucket_tokens > 0 + and requests + > min( + ( + prefill_config.full_prefill_max_req + or max(server_args.chunked_prefill_size // 512, 1) + ), + bucket_tokens, + ) + ): + outcome = "request_slots" + elif global_tokens <= 0: + outcome = "no_tokens" + elif bucket_tokens <= 0: + sparse_threshold_text = os.environ.get("SGLANG_DSV4_SPARSE_PREFILL_THRESHOLD") + if sparse_threshold_text is not None and global_tokens > int( + sparse_threshold_text + ): + # The DSpark launcher validates this threshold against the pinned + # DSV4 backend and requires dense graphs through the exact + # transition before ingress. A larger batch therefore takes the + # model's dedicated FlashMLA sparse-attention implementation, not + # an accidental generic eager fallback. + outcome = "dsv4_sparse_attention" + else: + outcome = "token_oversize" + elif bucket_tokens > global_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR: + outcome = "padding_crossover" + elif not batch.can_run_dp_breakable_cuda_graph: + outcome = "synchronized_constraint" + else: + # The fail-closed runner rejects a local graph/eager divergence before + # reaching metrics. Retain a distinct label in case a future backend + # deliberately applies an additional shared constraint. + outcome = "runner_constraint" + + return PrefillExecutionObservation( + outcome=outcome, + scheduled_tokens=scheduled_tokens, + executed_tokens=scheduled_tokens, + requests=requests, + bucket_tokens=0, + ) + + @dataclass(kw_only=True) class SchedulerMetricsReporter: scheduler: Scheduler @@ -157,7 +265,7 @@ def _init_metrics( self._mfu_log_read_bytes = 0.0 self._mfu_log_write_bytes = 0.0 - self.fwd_occupancy = float("nan") + self.fwd_occupancy = 0.0 self.forward_pass_device_timer: Optional[DeviceTimer] = None @@ -610,6 +718,16 @@ def report_prefill_stats( self.metrics_collector.increment_prefill_cuda_graph_pass( value=can_run_cuda_graph ) + if batch is not None: + observation = classify_prefill_execution( + batch=batch, + can_run_cuda_graph=can_run_cuda_graph, + server_args=self.scheduler.server_args, + dp_rank=self.scheduler.ps.dp_rank, + ) + self.metrics_collector.observe_prefill_execution( + **dataclasses.asdict(observation) + ) self.metrics_collector.increment_realtime_tokens( prefill_compute_tokens=prefill_stats.log_input_tokens, prefill_cache_tokens=prefill_stats.log_hit_tokens, @@ -1078,11 +1196,10 @@ def update_device_timer(self): self.forward_pass_device_timer._report() now = time.perf_counter() if self._device_timer_window_batch_count == 0: - # Window start: keep the last published value instead of NaN-ing - # the gauge. Readers sample it asynchronously, and the window - # boundary can phase-lock with the decode-log cadence, turning a - # one-tick NaN into NaN on every log line. NaN is published only - # when truly stale (reset_device_timer_window after idle). + # Keep the last finite value while the new window gets its second + # sample. Readers scrape asynchronously, so publishing a sentinel + # here can phase-lock with the log cadence and poison every scrape. + # True scheduler idle is represented explicitly as zero below. self._device_timer_window_start = now self._device_timer_window_gpu_time = 0.0 else: @@ -1098,7 +1215,10 @@ def update_device_timer(self): def reset_device_timer_window(self): if ENABLE_METRICS_DEVICE_TIMER: self._device_timer_window_batch_count = 0 - self.fwd_occupancy = float("nan") + self._device_timer_window_gpu_time = 0.0 + self._device_timer_window_start = None + # Scheduler idle is a measured zero, not an undefined sample. + self.fwd_occupancy = 0.0 def _maybe_log_idle_metrics(self): """Collect and log metrics every 30 seconds during idle.""" @@ -1150,4 +1270,7 @@ def _maybe_log_idle_metrics(self): self.stats.num_decode_transfer_queue_reqs = QueueCount.from_reqs( self.scheduler.disagg_decode_transfer_queue.queue, priority_enabled ) + # Do not retain an active-window sample (or publish NaN) while idle. + self.stats.fwd_occupancy = 0.0 + self.fwd_occupancy = 0.0 self.metrics_collector.log_stats(self.stats) diff --git a/python/sglang/srt/managers/scheduler_components/single_node_dp2_sync.py b/python/sglang/srt/managers/scheduler_components/single_node_dp2_sync.py new file mode 100644 index 000000000000..1334859a34ed --- /dev/null +++ b/python/sglang/srt/managers/scheduler_components/single_node_dp2_sync.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +import atexit +import ctypes +import importlib.util +import os +import time +from collections.abc import Sequence +from pathlib import Path + +import torch + +_ENABLE_ENV = "SGLANG_DSPARK_DP2_SHM_MLP_SYNC" +_SESSION_ENV = "SGLANG_DSPARK_DP2_SHM_SESSION_ID" +_TIMEOUT_ENV = "SGLANG_DSPARK_DP2_SHM_TIMEOUT_MS" +_METRICS_ENV = "SGLANG_DSPARK_DP2_SHM_METRICS" +_LIBRARY_ENV = "SGLANG_DSPARK_DP2_SHM_LIBRARY" +_SKIP_GATHER_ENV = "SGLANG_SCHEDULER_SKIP_ALL_GATHER" +_NCCL_GATHER_ENV = "SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH" +_LIBRARY_BASENAME = "sglang_dp2_sync.so" +_EXPECTED_ABI_VERSION = 2 +_WORLD_SIZE = 2 +_PAYLOAD_WIDTH = 7 +_MLP_CHANNEL = "mlp" +_VERIFY_TIER_CHANNEL = "verify_tier" +_VERIFY_TIER_PAYLOAD_MAGIC = 0x4453504B54494552 +_ERROR_BUFFER_SIZE = 512 +_METRIC_PHASES = ( + "exchange_total", + "peer_wait", + "arrival_skew", + "post_latest_arrival", +) +_enabled: bool | None = None + + +def _strict_bool_env(name: str, default: str) -> bool: + value = os.environ.get(name, default) + if value not in {"0", "1"}: + raise RuntimeError(f"{name} must be exactly 0 or 1, got {value!r}") + return value == "1" + + +def single_node_dp2_sync_enabled() -> bool: + global _enabled + if _enabled is None: + _enabled = _strict_bool_env(_ENABLE_ENV, "0") + return _enabled + + +def _timeout_ns() -> int: + value = os.environ.get(_TIMEOUT_ENV, "30000") + try: + timeout_ms = int(value) + except ValueError as error: + raise RuntimeError( + f"{_TIMEOUT_ENV} must be a positive integer, got {value!r}" + ) from error + if timeout_ms <= 0: + raise RuntimeError(f"{_TIMEOUT_ENV} must be a positive integer, got {value!r}") + return timeout_ms * 1_000_000 + + +class _NativeStats(ctypes.Structure): + _fields_ = [ + ("sequence", ctypes.c_uint64), + ("total_ns", ctypes.c_uint64), + ("peer_wait_ns", ctypes.c_uint64), + ("arrival_skew_ns", ctypes.c_uint64), + ("post_latest_arrival_ns", ctypes.c_uint64), + ] + + +def _resolve_library_path() -> Path: + override = os.environ.get(_LIBRARY_ENV) + if override: + return Path(override) + + package = importlib.util.find_spec("sgl_kernel") + if package is None or package.submodule_search_locations is None: + raise RuntimeError( + "strict DP2 shared-memory sync requires the sgl-kernel package" + ) + locations = tuple(package.submodule_search_locations) + if len(locations) != 1: + raise RuntimeError( + "strict DP2 shared-memory sync expected one sgl-kernel package " + f"location, got {locations}" + ) + return Path(locations[0]) / _LIBRARY_BASENAME + + +def _load_library() -> ctypes.CDLL: + path = _resolve_library_path() + if not path.is_file(): + raise RuntimeError( + f"strict DP2 shared-memory MLP sync library is missing: {path}" + ) + library = ctypes.CDLL(str(path)) + library.sglang_dp2_sync_abi_version.argtypes = [] + library.sglang_dp2_sync_abi_version.restype = ctypes.c_uint32 + library.sglang_dp2_sync_open.argtypes = [ + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_uint64, + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.sglang_dp2_sync_open.restype = ctypes.c_int + library.sglang_dp2_sync_exchange.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(_NativeStats), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.sglang_dp2_sync_exchange.restype = ctypes.c_int + library.sglang_dp2_sync_exchange_values.argtypes = [ + ctypes.c_void_p, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_int64, + ctypes.c_void_p, + ctypes.POINTER(_NativeStats), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.sglang_dp2_sync_exchange_values.restype = ctypes.c_int + library.sglang_dp2_sync_close.argtypes = [ctypes.c_void_p] + library.sglang_dp2_sync_close.restype = None + actual_abi = library.sglang_dp2_sync_abi_version() + if actual_abi != _EXPECTED_ABI_VERSION: + raise RuntimeError( + "strict DP2 shared-memory MLP sync ABI mismatch: " + f"expected {_EXPECTED_ABI_VERSION}, got {actual_abi}" + ) + return library + + +def validate_single_node_dp2_sync_runtime() -> None: + if not single_node_dp2_sync_enabled(): + raise RuntimeError(f"{_ENABLE_ENV}=1 is required; Gloo fallback is forbidden") + session_id = os.environ.get(_SESSION_ENV, "") + if not session_id: + raise RuntimeError(f"{_SESSION_ENV} is required") + _timeout_ns() + if not _strict_bool_env(_METRICS_ENV, "1"): + raise RuntimeError(f"{_METRICS_ENV}=1 is required for runtime observability") + for incompatible_env in (_SKIP_GATHER_ENV, _NCCL_GATHER_ENV): + if os.environ.get(incompatible_env) != "0": + raise RuntimeError( + f"{incompatible_env}=0 must be explicit; alternate MLP sync " + "paths are forbidden" + ) + _load_library() + + +class _SyncMetrics: + def __init__(self, rank: int, channel: str) -> None: + from prometheus_client import Counter, Gauge + + if channel == _MLP_CHANNEL: + metric_prefix = "sglang:dp2_mlp_sync" + subject = "MLP geometry" + elif channel == _VERIFY_TIER_CHANNEL: + metric_prefix = "sglang:dp2_verify_tier_sync" + subject = "DSpark verify-tier" + else: + raise RuntimeError(f"unknown strict DP2 sync channel {channel!r}") + self._rank = str(rank) + seconds_total = Counter( + name=f"{metric_prefix}_seconds_total", + documentation=( + f"Cumulative native DP2 shared-memory {subject} exchange " + "time. peer_wait includes rank arrival skew; " + "post_latest_arrival isolates transport and wakeup cost." + ), + labelnames=["dp_rank", "phase"], + ) + calls_total = Counter( + name=f"{metric_prefix}_calls_total", + documentation=(f"Number of native DP2 shared-memory {subject} exchanges."), + labelnames=["dp_rank"], + ) + max_seconds = Gauge( + name=f"{metric_prefix}_max_seconds", + documentation=( + f"Maximum native DP2 {subject} exchange phase time in the latest " + "one-second reporting window." + ), + labelnames=["dp_rank", "phase"], + multiprocess_mode="mostrecent", + ) + enabled = Gauge( + name=f"{metric_prefix}_enabled", + documentation=( + f"One when the strict native DP2 shared-memory {subject} sync path " + "is initialized." + ), + labelnames=["dp_rank"], + multiprocess_mode="mostrecent", + ) + sequence = Gauge( + name=f"{metric_prefix}_sequence", + documentation=( + f"Latest completed strict native DP2 {subject} sync sequence." + ), + labelnames=["dp_rank"], + multiprocess_mode="mostrecent", + ) + self._seconds = [ + seconds_total.labels(dp_rank=self._rank, phase=phase) + for phase in _METRIC_PHASES + ] + self._max = [ + max_seconds.labels(dp_rank=self._rank, phase=phase) + for phase in _METRIC_PHASES + ] + self._calls = calls_total.labels(dp_rank=self._rank) + self._sequence = sequence.labels(dp_rank=self._rank) + enabled.labels(dp_rank=self._rank).set(1) + self._totals_ns = [0, 0, 0, 0] + self._max_ns = [0, 0, 0, 0] + self._calls_pending = 0 + self._flush_deadline_ns = time.monotonic_ns() + 1_000_000_000 + + def observe(self, stats: _NativeStats) -> None: + values = ( + stats.total_ns, + stats.peer_wait_ns, + stats.arrival_skew_ns, + stats.post_latest_arrival_ns, + ) + for index, value in enumerate(values): + self._totals_ns[index] += value + self._max_ns[index] = max(self._max_ns[index], value) + self._calls_pending += 1 + now_ns = time.monotonic_ns() + if now_ns < self._flush_deadline_ns: + return + for index in range(len(_METRIC_PHASES)): + self._seconds[index].inc(self._totals_ns[index] / 1_000_000_000) + self._max[index].set(self._max_ns[index] / 1_000_000_000) + self._totals_ns[index] = 0 + self._max_ns[index] = 0 + self._calls.inc(self._calls_pending) + self._sequence.set(stats.sequence) + self._calls_pending = 0 + self._flush_deadline_ns = now_ns + 1_000_000_000 + + +class _SingleNodeDP2Sync: + def __init__(self, rank: int, channel: str) -> None: + validate_single_node_dp2_sync_runtime() + self._pid = os.getpid() + self._rank = rank + self._channel = channel + self._library = _load_library() + self._handle = ctypes.c_void_p() + self._stats = _NativeStats() + self._global_payload = (ctypes.c_int64 * (_WORLD_SIZE * _PAYLOAD_WIDTH))() + self._error = ctypes.create_string_buffer(_ERROR_BUFFER_SIZE) + session_id_raw = os.environ[_SESSION_ENV] + if channel == _MLP_CHANNEL: + session_id = session_id_raw.encode() + elif channel == _VERIFY_TIER_CHANNEL: + session_id = f"{session_id_raw}:{channel}".encode() + else: + raise RuntimeError(f"unknown strict DP2 sync channel {channel!r}") + result = self._library.sglang_dp2_sync_open( + session_id, + rank, + _timeout_ns(), + ctypes.byref(self._handle), + self._error, + len(self._error), + ) + if result != 0: + raise RuntimeError(self._error.value.decode()) + if self._handle.value is None: + raise RuntimeError("strict DP2 shared-memory sync returned a null handle") + self._metrics = _SyncMetrics(rank, channel) + + def exchange( + self, + num_tokens: int, + num_tokens_for_logprob: int, + can_cuda_graph: bool, + is_extend_in_batch: bool, + local_can_run_tbo: bool, + local_forward_mode: int, + can_run_breakable_cuda_graph: bool, + ) -> Sequence[int]: + if os.getpid() != self._pid: + raise RuntimeError( + "strict DP2 shared-memory sync handle crossed a process fork" + ) + self._error[0] = 0 + result = self._library.sglang_dp2_sync_exchange_values( + self._handle, + num_tokens, + num_tokens_for_logprob, + int(can_cuda_graph), + int(is_extend_in_batch), + int(local_can_run_tbo), + local_forward_mode, + int(can_run_breakable_cuda_graph), + self._global_payload, + ctypes.byref(self._stats), + self._error, + len(self._error), + ) + if result != 0: + message = self._error.value.decode() + raise RuntimeError( + f"strict DP2 shared-memory {self._channel} sync failed: " + message + ) + self._metrics.observe(self._stats) + return self._global_payload + + def close(self) -> None: + if self._handle.value is not None and os.getpid() == self._pid: + self._library.sglang_dp2_sync_close(self._handle) + self._handle = ctypes.c_void_p() + + +_managers: dict[str, _SingleNodeDP2Sync] = {} + + +def _get_manager( + *, + channel: str, + group: torch.distributed.ProcessGroup, + dp_size: int, + tp_size: int, + cp_size: int, +) -> _SingleNodeDP2Sync: + manager = _managers.get(channel) + if manager is not None: + return manager + if not single_node_dp2_sync_enabled(): + raise RuntimeError( + f"strict DP2 shared-memory {channel} sync was called while disabled" + ) + if (dp_size, tp_size, cp_size) != (2, 1, 1): + raise RuntimeError( + f"strict shared-memory {channel} sync only supports exact " + f"DP2/TP1/CP1 geometry, got DP{dp_size}/TP{tp_size}/CP{cp_size}" + ) + if torch.distributed.get_world_size(group) != _WORLD_SIZE: + raise RuntimeError( + f"strict DP2 shared-memory {channel} sync requires group world size 2" + ) + rank = torch.distributed.get_rank(group) + if rank not in (0, 1): + raise RuntimeError( + f"strict DP2 shared-memory {channel} sync got invalid group rank {rank}" + ) + manager = _SingleNodeDP2Sync(rank, channel) + _managers[channel] = manager + atexit.register(manager.close) + return manager + + +def exchange_single_node_dp2_mlp_info( + num_tokens: int, + num_tokens_for_logprob: int, + can_cuda_graph: bool, + is_extend_in_batch: bool, + local_can_run_tbo: bool, + local_forward_mode: int, + can_run_breakable_cuda_graph: bool, + *, + group: torch.distributed.ProcessGroup, + dp_size: int, + tp_size: int, + cp_size: int, +) -> Sequence[int]: + manager = _get_manager( + channel=_MLP_CHANNEL, + group=group, + dp_size=dp_size, + tp_size=tp_size, + cp_size=cp_size, + ) + return manager.exchange( + num_tokens, + num_tokens_for_logprob, + can_cuda_graph, + is_extend_in_batch, + local_can_run_tbo, + local_forward_mode, + can_run_breakable_cuda_graph, + ) + + +def exchange_single_node_dp2_verify_tier( + local_tier_num_tokens: int, + *, + group: torch.distributed.ProcessGroup, + dp_size: int, + tp_size: int, + cp_size: int, +) -> list[int]: + if local_tier_num_tokens < -1: + raise RuntimeError( + "strict DP2 verify-tier sync requires a tier >= -1, got " + f"{local_tier_num_tokens}" + ) + manager = _get_manager( + channel=_VERIFY_TIER_CHANNEL, + group=group, + dp_size=dp_size, + tp_size=tp_size, + cp_size=cp_size, + ) + payload = manager.exchange( + local_tier_num_tokens, + _VERIFY_TIER_PAYLOAD_MAGIC, + False, + False, + False, + 0, + False, + ) + for rank in range(_WORLD_SIZE): + magic = int(payload[rank * _PAYLOAD_WIDTH + 1]) + if magic != _VERIFY_TIER_PAYLOAD_MAGIC: + raise RuntimeError( + "strict DP2 verify-tier sync channel mismatch: " + f"rank {rank} published 0x{magic:016x}" + ) + return [int(payload[rank * _PAYLOAD_WIDTH]) for rank in range(_WORLD_SIZE)] diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 7fe6e9491797..c6d94481bcd5 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py @@ -1359,13 +1359,23 @@ def _should_use_batch_tokenization(self, batch_size, requests) -> bool: Current policy: - Respect explicit server flag `enable_tokenizer_batch_encode`. - - Or, if no request has text or multimodal input (all use pre-tokenized input_ids or input_embeds), batch the requests without tokenization. - - Batch tokenization does not support DP attention yet, and it will make everything goes to the first rank currently + - If no request has text or multimodal input (all use pre-tokenized + input_ids or input_embeds), batch the requests without tokenization. + - Under DP attention, the fast batch transport is safe when every item + has the same explicit routed DP rank. """ + has_atomic_dp_route = ( + batch_size > 0 + and requests[0].routed_dp_rank is not None + and all( + requests[i].routed_dp_rank == requests[0].routed_dp_rank + for i in range(1, batch_size) + ) + ) return batch_size > 0 and ( self.server_args.enable_tokenizer_batch_encode or ( - (not self.server_args.enable_dp_attention) + (not self.server_args.enable_dp_attention or has_atomic_dp_route) and (not self._batch_has_text(batch_size, requests)) ) ) @@ -2765,16 +2775,22 @@ async def sigterm_watchdog(self): break # Stop the watchdog: child exits are expected during shutdown, not crashes. - if self._subprocess_watchdog is not None: - self._subprocess_watchdog.stop() - # Ask schedulers to release resources in userspace and exit (see - # ShutdownReq), then wait for them before hard-killing the rest. + subprocess_watchdog = self._subprocess_watchdog + if subprocess_watchdog is None: + raise RuntimeError( + "graceful shutdown requires the SGLang subprocess watchdog" + ) + subprocess_watchdog.stop() + # Each scheduler forwards ShutdownReq to the detokenizer before it + # releases host memory and destroys its distributed process groups. self._dispatch_to_scheduler(ShutdownReq()) - deadline = time.monotonic() + 15 - while time.monotonic() < deadline and collect_scheduler_processes(): - time.sleep(0.1) - kill_process_tree(os.getpid(), include_parent=True) - sys.exit(0) + subprocess_watchdog.wait_for_clean_exit(timeout=180) + logger.info("All SGLang child processes exited cleanly.") + # The signal handler replaced uvicorn's SIGTERM callback, so there is + # no server.should_exit path left to drive. At this point every child + # and GPU/distributed resource has already been joined; terminate the + # CPU-only parent directly with a successful container status. + os._exit(0) def force_exit_handler(self): """Put some custom force exit logic here.""" diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 94ae2a307eb4..015c65aef175 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -2,6 +2,7 @@ import dataclasses import logging +from concurrent.futures import Future from dataclasses import dataclass from typing import TYPE_CHECKING, Any, List, Optional, Union @@ -22,6 +23,18 @@ logger = logging.getLogger(__name__) +@dataclass(frozen=True, slots=True) +class PreparedGrammarResult: + """Minimal DSpark result materialized off the scheduler critical path.""" + + next_token_ids: list[int] + accept_lens: Optional[list[int]] + queue_delay_seconds: float + copy_wait_seconds: float + tensor_to_list_seconds: float + submit_to_ready_seconds: float + + def _async_d2h(t: torch.Tensor) -> torch.Tensor: """Async D2H copy for overlap scheduling. On CUDA the dest is pinned (a D2H to pageable host memory blocks the caller until done) and record_stream keeps @@ -64,6 +77,14 @@ class GenerationBatchResult: copy_done: Optional[torch.cuda.Event] = None delay_sample_func: Optional[callable] = None future_indices: Optional[torch.Tensor] = None + # DSpark grammar overlap needs the previous accepted tokens before the next + # target-verify mask can be built. The worker publishes this minimal D2H + # result immediately after accept, before the hidden-state commit and the + # rest of the forward tail. The normal result copy remains independent. + grammar_next_token_ids: Optional[torch.Tensor] = None + grammar_accept_lens: Optional[torch.Tensor] = None + grammar_copy_done: Optional[torch.cuda.Event] = None + grammar_result_future: Optional[Future[PreparedGrammarResult]] = None speculative_num_draft_tokens: Optional[int] = None # Grammar FSM advance memoization (spec-v2 overlap). advance_grammar_fsm sets diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 244eef333b70..f9bd3401f361 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -237,6 +237,8 @@ def init_metrics_collector(self): server_args = get_server_args() labels = {"cache_type": self.__class__.__name__} + if torch.distributed.is_available() and torch.distributed.is_initialized(): + labels["rank"] = str(torch.distributed.get_rank()) if server_args.extra_metric_labels: labels.update(server_args.extra_metric_labels) radix_cache_cls = resolve_collector_class( diff --git a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py index 14c9e52cb52f..b773e8587a64 100644 --- a/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py +++ b/python/sglang/srt/mem_cache/deepseek_v4_memory_pool.py @@ -49,6 +49,8 @@ def get_compress_state_ring_size( class DeepSeekV4SingleKVPool(KVCache): + dynamic_kv_cache_scale_block_size = 64 + def __init__( self, size: int, @@ -76,7 +78,7 @@ def __init__( self.qk_rope_head_dim = qk_rope_head_dim self.scale_pad = 1 - self.quantize_block_size = 64 + self.quantize_block_size = self.dynamic_kv_cache_scale_block_size self.rope_storage_dtype = torch.bfloat16 self.k_with_scale_buffer_dtype = torch.int8 self._create_buffers() diff --git a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py index af35e964af45..84c4b106fb4e 100644 --- a/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py +++ b/python/sglang/srt/mem_cache/hybrid_cache/hybrid_cache_controller.py @@ -411,8 +411,7 @@ def start_writing(self) -> None: self.move_hybrid_indices(op) ) self.write_queue.clear() - start_event = device_module.Event() - finish_event = device_module.Event() + start_event, finish_event, timing_enabled = make_timing_event_pair() start_event.record() with device_module.stream(self.write_stream): start_event.wait(self.write_stream) @@ -437,7 +436,15 @@ def start_writing(self) -> None: device_indices, resolved_pool_transfers, ) - self.ack_write_queue.append(HiCacheAck(start_event, finish_event, op.node_ids)) + self.ack_write_queue.append( + HiCacheAck( + start_event, + finish_event, + op.node_ids, + num_tokens=len(op.device_indices), + timing_enabled=timing_enabled, + ) + ) def load( self, diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 268976191925..c4154f92b961 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -668,6 +668,12 @@ def __init__(self, size: int, page_size: int, layout: str = "layer_first"): def clear(self): self.free_slots = torch.arange(self.size, dtype=torch.int64) + def destroy(self) -> None: + """Release host resources; this logical index pool owns no buffers. + + HostPoolGroup.destroy() requires every member to provide this method. + """ + def available_size(self): return len(self.free_slots) diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 69477a27a42f..7dadc2495703 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -129,6 +129,9 @@ def __init__( if params.enable_metrics: self.init_metrics_collector() self._enable_metrics_flag = params.enable_metrics + self._hicache_phase_accumulator: dict[str, tuple[float, int, float]] = {} + self._hicache_pending_snapshot: dict[str, int] = {} + self._hicache_metrics_flush_deadline = time.perf_counter() + 1.0 self.enable_storage_metrics = False self.storage_metrics_collector: Optional[StorageMetricsCollector] = None self.extra_metric_labels = None @@ -858,9 +861,12 @@ def _execute_kv_backup(self, node_id, device_value, comp_xfers, sidecar_xfers): return None aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] aux_xfers.extend(sidecar_xfers) - return self.cache_controller.write( + submit_started = time.perf_counter() + host_indices = self.cache_controller.write( device_value, node_id=node_id, extra_pools=aux_xfers or None ) + self._finish_hicache_phase("write_submit", submit_started) + return host_indices def _track_write_through_node( self, @@ -990,11 +996,13 @@ def _load_back_transfers( # Load H→D aux_xfers = [x for xfers in comp_xfers.values() for x in xfers] aux_xfers.extend(sidecar_xfers) + submit_started = time.perf_counter() device_indices = self.cache_controller.load( host_indices=kv_xfer.host_indices, node_id=node_id, extra_pools=aux_xfers or None, ) + self._finish_hicache_phase("load_prepare", submit_started) self.dec_lock_ref(node_id, ancestor_lock_params) if device_indices is None: @@ -1700,6 +1708,72 @@ def clear_storage_backend(self) -> bool: # ---- HiCache: Async Event Management ---- + def _finish_hicache_phase(self, phase: str, started: float) -> float: + now = time.perf_counter() + duration = now - started + if self.metrics_collector is not None: + total, calls, maximum = self._hicache_phase_accumulator.get( + phase, (0.0, 0, 0.0) + ) + self._hicache_phase_accumulator[phase] = ( + total + duration, + calls + 1, + max(maximum, duration), + ) + if now >= self._hicache_metrics_flush_deadline: + for ( + accumulated_phase, + ( + accumulated_seconds, + accumulated_calls, + accumulated_max, + ), + ) in self._hicache_phase_accumulator.items(): + self.metrics_collector.observe_hicache_scheduler_phase( + accumulated_phase, + accumulated_seconds, + accumulated_calls, + accumulated_max, + ) + for kind, count in self._hicache_pending_snapshot.items(): + self.metrics_collector.set_hicache_pending_operations(kind, count) + self._hicache_phase_accumulator.clear() + self._hicache_metrics_flush_deadline = now + 1.0 + if duration >= 0.5: + logger.warning( + "Slow HiCache scheduler phase: phase=%s duration=%.3fs " + "pending_write=%d pending_load=%d", + phase, + duration, + len(self.ongoing_write_through), + len(self.ongoing_load_back), + ) + return now + + def _record_hicache_backup_ack(self, ack) -> None: + if self.metrics_collector is None: + return + duration_seconds = None + if ack.timing_enabled: + duration_seconds = ack.start_event.elapsed_time(ack.finish_event) / 1000.0 + self.metrics_collector.observe_hicache_backup( + ack.num_tokens, + duration_seconds, + ) + + def _update_hicache_pending_metrics(self) -> None: + if self.metrics_collector is None or self.cache_controller is None: + return + cc = self.cache_controller + for kind, count in ( + ("write_acks", len(cc.ack_write_queue)), + ("write_nodes", len(self.ongoing_write_through)), + ("load_acks", len(cc.ack_load_queue)), + ("load_nodes", len(self.ongoing_load_back)), + ("pp_sync", len(self.work_list)), + ): + self._hicache_pending_snapshot[kind] = count + def writing_check(self, write_back: bool = False) -> None: """Poll write-through completions.""" cc = self.cache_controller @@ -1708,36 +1782,48 @@ def writing_check(self, write_back: bool = False) -> None: if write_back: # Blocking: wait for all pending write-backs + flush_started = time.perf_counter() while self.ongoing_write_through: for ack in cc.ack_write_queue: ack.finish_event.synchronize() + self._record_hicache_backup_ack(ack) for ack_id in ack.node_ids: if ack_id in self.ongoing_write_through: self._finish_write_through_ack(ack_id) cc.ack_write_queue.clear() assert len(self.ongoing_write_through) == 0 + self._finish_hicache_phase("write_back_flush", flush_started) + self._update_hicache_pending_metrics() return # Every rank must enter the all_reduce below; ongoing_write_through can # diverge across ranks (e.g. a backup returning 0 on a subset). + phase_started = time.perf_counter() finish_count = 0 if self.pp_rank == 0: for ack in cc.ack_write_queue: if not ack.finish_event.query(): break finish_count += 1 + phase_started = self._finish_hicache_phase("write_query", phase_started) finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu") self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) finish_count = finish_count_tensor.item() + phase_started = self._finish_hicache_phase( + "write_completion_sync", phase_started + ) # Process completed acks while finish_count > 0: ack = cc.ack_write_queue.pop(0) ack.finish_event.synchronize() + self._record_hicache_backup_ack(ack) for ack_id in ack.node_ids: self._finish_write_through_ack(ack_id) finish_count -= 1 + self._finish_hicache_phase("write_finalize", phase_started) + self._update_hicache_pending_metrics() def loading_check(self) -> None: """Poll load-back completions.""" @@ -1746,15 +1832,20 @@ def loading_check(self) -> None: return # Every rank must enter the all_reduce below; ongoing_load_back can # diverge across ranks. + phase_started = time.perf_counter() finish_count = 0 if self.pp_rank == 0: for ack in cc.ack_load_queue: if not ack.finish_event.query(): break finish_count += 1 + phase_started = self._finish_hicache_phase("load_query", phase_started) finish_count_tensor = torch.tensor(finish_count, dtype=torch.int, device="cpu") self._all_reduce(finish_count_tensor, torch.distributed.ReduceOp.MIN) finish_count = finish_count_tensor.item() + phase_started = self._finish_hicache_phase( + "load_completion_sync", phase_started + ) while finish_count > 0: ack = cc.ack_load_queue.pop(0) @@ -1772,6 +1863,8 @@ def loading_check(self) -> None: duration_ms / 1000.0 ) finish_count -= 1 + self._finish_hicache_phase("load_finalize", phase_started) + self._update_hicache_pending_metrics() # ---- HiCache: Scheduler Entry Points ---- @@ -1819,16 +1912,21 @@ def init_load_back( def check_hicache_events(self) -> None: """Called per scheduler step to poll async HiCache events.""" + round_started = phase_started = time.perf_counter() # Reap the previous round's PP-sync sends before issuing new ones. self._drain_async_work() + phase_started = self._finish_hicache_phase("pp_sync_drain", phase_started) self.writing_check() self.loading_check() + phase_started = time.perf_counter() if self.enable_storage: self.drain_storage_control_queues() if self.enable_storage_metrics and self.storage_metrics_collector is not None: self.storage_metrics_collector.log_storage_metrics( self.cache_controller.storage_backend.get_stats() ) + self._finish_hicache_phase("storage_control", phase_started) + self._finish_hicache_phase("event_round_total", round_started) def flush_write_through_acks(self) -> None: """Flush pending write-through acknowledgements.""" @@ -1837,7 +1935,10 @@ def flush_write_through_acks(self) -> None: def ready_to_load_host_cache(self) -> int: """Notify the cache controller to start the KV cache loading.""" if self.cache_controller is not None: - return self.cache_controller.start_loading() + submit_started = time.perf_counter() + producer_id = self.cache_controller.start_loading() + self._finish_hicache_phase("load_submit", submit_started) + return producer_id return 0 # ---- Query / Inspection APIs ---- diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 277d950385f0..17c55421d14e 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -1238,25 +1238,14 @@ def prepare_mlp_sync_batch(self, model_runner: ModelRunner): ): # 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 - # 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) - ): + # A captured prefill CUDA graph requires every DP rank to run the SAME + # token shape. Under SUM_LEN each rank can select a different capture + # bucket, so in-graph DP collectives mismatch and corrupt or deadlock + # the forward. Scheduler-side admission has already min-reduced all + # local constraints and applied the shared bucket crossover. Honor + # that single decision unconditionally by forcing MAX_LEN. + if self.can_run_dp_breakable_cuda_graph: + assert self.is_extend_in_batch dp_padding_mode = DpPaddingMode.MAX_LEN self.dp_padding_mode = dp_padding_mode diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index c919cd39429e..73e39de01eb4 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1542,13 +1542,14 @@ def _forward_raw( forward_batch.forward_mode.is_extend(include_draft_extend_v2=True) and not isinstance(self.prefill_cuda_graph_runner, EagerRunner) and self.prefill_cuda_graph_runner is not None + and not envs.SGLANG_PREFILL_CUDA_GRAPH_EAGER_VALIDATION.get() and self.prefill_cuda_graph_runner.can_run_graph(forward_batch) and get_cp_strategy() is None ): category = ( "target_verify" if forward_batch.forward_mode.is_target_verify() - else "extend" + else "extend_cuda_graph" ) # Prefill cuda graph (piecewise). kwargs = self._extend_forward_kwargs(forward_batch, pp_proxy_tensors) diff --git a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py index f91339d5e4e5..b4b4dca800f3 100644 --- a/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py +++ b/python/sglang/srt/model_executor/model_runner_components/load_model_utils.py @@ -87,6 +87,34 @@ def maybe_trigger_remote_instance_nccl_send_group( def load_kv_cache_scales(*, model, server_args: ServerArgs) -> None: if server_args.kv_cache_dtype == "fp8_e4m3": + dynamic_scale_block_size = getattr( + model, + "dynamic_kv_cache_scale_block_size", + None, + ) + if dynamic_scale_block_size is not None: + if ( + isinstance(dynamic_scale_block_size, bool) + or not isinstance(dynamic_scale_block_size, int) + or dynamic_scale_block_size <= 0 + ): + raise RuntimeError( + f"Model {model.__class__} advertises an invalid dynamic " + "FP8 KV cache scale block size: " + f"{dynamic_scale_block_size!r}" + ) + if server_args.quantization_param_path is not None: + raise RuntimeError( + f"Model {model.__class__} uses native dynamic per-" + f"{dynamic_scale_block_size}-element FP8 KV cache scales; " + "an external static scaling-factor file is invalid." + ) + logger.info( + "Using model-native dynamic per-%d-element FP8 KV cache " + "scales; no external scaling-factor file is required.", + dynamic_scale_block_size, + ) + return if server_args.quantization_param_path is not None: if callable(getattr(model, "load_kv_cache_scales", None)): model.load_kv_cache_scales(server_args.quantization_param_path) diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py index 55f6fac786bd..0d84b36781e4 100644 --- a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py @@ -56,7 +56,10 @@ def freeze_gc(enable_cudagraph_gc: bool): def get_batch_sizes_to_capture( - model_runner: ModelRunner, captured_req_width: int = 1 + model_runner: ModelRunner, + captured_req_width: int = 1, + *, + enable_two_batch_overlap: bool | None = None, ) -> Tuple[List[int], List[int]]: """Build the (capture_bs, compile_bs) lists for the decode runner. @@ -72,7 +75,9 @@ def get_batch_sizes_to_capture( # TBO splits each request's rows across two micro-batches, so the # alignment constraint applies per request rather than per token row. alignment_width = captured_req_width - if server_args.enable_two_batch_overlap: + if enable_two_batch_overlap is None: + enable_two_batch_overlap = server_args.enable_two_batch_overlap + if enable_two_batch_overlap: mul_base *= 2 alignment_width = 1 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 60f4cab53634..71689d46f8a6 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 @@ -214,8 +214,13 @@ def __init__( self.require_mlp_sync = ( model_runner.server_args.enable_dp_attention or self.require_gathered_buffer ) - self.enable_two_batch_overlap = ( - model_runner.server_args.enable_two_batch_overlap + configured_tbo = model_runner.server_args.enable_two_batch_overlap + decode_attn_backend = attn_backend or model_runner.attn_backend + decode_attn_backend = getattr( + decode_attn_backend, "primary", decode_attn_backend + ) + self.enable_two_batch_overlap = configured_tbo and getattr( + decode_attn_backend, "tbo_supports_decode_cuda_graph", True ) self.use_ngram_embedding = model_runner.ngram_embedding_manager.enabled if self.use_ngram_embedding: @@ -275,7 +280,9 @@ def __init__( # --- bucket sizes --------------------------------------------- self.capture_bs, self.compile_bs = get_batch_sizes_to_capture( - model_runner, self.captured_req_width + model_runner, + self.captured_req_width, + enable_two_batch_overlap=self.enable_two_batch_overlap, ) if KTRANSFORMERS_AVAILABLE: KTMoEWrapper.set_capture_batch_sizes(self.capture_bs) @@ -831,7 +838,9 @@ def capture(self) -> None: if self.enable_torch_compile and not (get_flags().capture.enable_torch_compile): self.enable_torch_compile = False _, self.compile_bs = get_batch_sizes_to_capture( - self.model_runner, self.captured_req_width + self.model_runner, + self.captured_req_width, + enable_two_batch_overlap=self.enable_two_batch_overlap, ) profile_context = empty_context() if self.enable_profile_cuda_graph: diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 21173ba0092c..8dfe3d60f62a 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -291,7 +291,7 @@ def _execute_extend( category = ( "target_verify" if forward_batch.forward_mode.is_target_verify() - else "extend" + else "extend_eager" ) ctx = ( model_runner.device_timer.wrap(metadata={"category": category}) 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 7f8f108cfa55..2e0b6600e6a0 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 @@ -25,7 +25,8 @@ cuda_graph_config.prefill.full_prefill_max_req request slots (auto-derived when unset); replay pads num_tokens to the nearest bucket and pads unused request slots with - zero-length sentinels. bs > slots falls back to eager. + zero-length sentinels. bs > slots is routed to eager by + the scheduler's shared admission policy. Attention metadata is refreshed out-of-graph against the slot-padded batch before capture/replay. - "tc_piecewise" — TcPiecewiseCudaGraphBackend: torch.compile @@ -117,7 +118,7 @@ # A replay executes every padded token in its capture bucket. Sparse bucket # lists can otherwise turn the lower launch overhead into substantially more # model work than an exact-shape eager forward. -_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR = 2 +_MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR = 4 def prefill_failure_msg(backend_name: str) -> str: @@ -280,6 +281,30 @@ def __init__(self, model_runner: ModelRunner): ) from e self._is_full_backend = isinstance(self.backend, FullCudaGraphBackend) + self._share_dspark_full_prefill_outputs = ( + self._is_full_backend + and model_runner.spec_algorithm.is_dspark() + and not model_runner.is_draft_worker + ) + self._dspark_full_prefill_output_buffers: Optional[ + tuple[tuple[torch.Tensor, torch.Tensor], list[torch.Tensor]] + ] = None + self._dspark_full_prefill_expected_aux_outputs = 0 + primary_attn_backend = getattr( + model_runner.attn_backend, "primary", model_runner.attn_backend + ) + supports_prefill_tbo_graph = getattr( + primary_attn_backend, "tbo_supports_cuda_graph_for", None + ) + self.enable_tbo_prefill_graph = ( + model_runner.server_args.enable_two_batch_overlap + and self._is_full_backend + and ( + supports_prefill_tbo_graph(ForwardMode.EXTEND) + if supports_prefill_tbo_graph is not None + else getattr(primary_attn_backend, "tbo_supports_cuda_graph", True) + ) + ) if self._is_full_backend: max_req = prefill_config.full_prefill_max_req if max_req is None: @@ -359,6 +384,19 @@ def __init__(self, model_runner: ModelRunner): self._input_embeds_arg_idx = ( params.index("input_embeds") if "input_embeds" in params else None ) + if self._share_dspark_full_prefill_outputs: + target_layer_ids = getattr( + self.layer_model, "dspark_layers_to_capture", None + ) + if ( + not isinstance(target_layer_ids, (list, tuple)) + or not target_layer_ids + ): + raise RuntimeError( + "DSpark FullCG shared outputs require a non-empty " + "dspark_layers_to_capture configuration" + ) + self._dspark_full_prefill_expected_aux_outputs = len(target_layer_ids) # --- aiter chip info pre-warming (AMD) ------------------------- maybe_pre_warm_aiter_chip_info() @@ -395,6 +433,16 @@ def _next_token_logits_buffer(self, rows: int) -> Optional[torch.Tensor]: def _uses_eager_prefill_tail(self) -> bool: return self.prefill_backend_name in (Backend.BREAKABLE, Backend.FULL) + def _capture_req_slots_for_tokens(self, num_tokens: int) -> int: + if not self._is_full_backend: + return self._capture_req_slots + # A prefill batch cannot contain more requests than token-axis + # elements unless it contains zero-length requests, which the graph + # path rejects. Shrinking the request axis with each token bucket lets + # Full capture useful sub-64-token shapes without reducing the global + # request-slot ceiling for larger batches. + return min(self._capture_req_slots, num_tokens) + def _prefill_logits_buffer_rows(self, forward_batch: ForwardBatch) -> int: if not forward_batch.return_logprob: return forward_batch.batch_size @@ -471,6 +519,120 @@ def _prefill_forward_context( ): yield + def _stage_dspark_full_prefill_body_output( + self, + output: tuple[ + tuple[torch.Tensor, torch.Tensor], + list[torch.Tensor], + ], + *, + num_tokens: int, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + list[torch.Tensor], + ]: + """Copy DSpark body outputs into one address-stable maximum-size tree. + + Full CUDA graphs retain every tensor returned from capture. DSpark's + target body returns the final hidden state plus one full token-axis + tensor for each configured target tap. Retaining that tree separately + for every token bucket consumes the sum of all bucket sizes even + though a scheduler rank replays only one body graph at a time. + + The first (largest) shape allocates a single shared tree during graph + warmup. Every graph records a copy into its matching prefix and returns + views of those same buffers. The eager tail consumes this tree + synchronously, then the DSpark prefill epilogue enqueues target-hidden + injection on the same stream before another target forward can replay, + so reuse cannot overlap a live consumer. + """ + if ( + not isinstance(output, tuple) + or len(output) != 2 + or not isinstance(output[0], tuple) + or len(output[0]) != 2 + ): + raise RuntimeError( + "DSpark FullCG body output must be " + "((hidden_states, pre_hc_head), auxiliary_hidden_states)" + ) + main_outputs, auxiliary_hidden_states = output + hidden_states, pre_hc_head = main_outputs + if ( + not isinstance(hidden_states, torch.Tensor) + or not isinstance(pre_hc_head, torch.Tensor) + or not isinstance(auxiliary_hidden_states, list) + ): + raise RuntimeError( + "DSpark FullCG body output has an incompatible tensor tree" + ) + if len( + auxiliary_hidden_states + ) != self._dspark_full_prefill_expected_aux_outputs or not all( + isinstance(tensor, torch.Tensor) for tensor in auxiliary_hidden_states + ): + raise RuntimeError( + "DSpark FullCG body output does not match the configured " + "target hidden-state taps" + ) + + outputs = [hidden_states, pre_hc_head, *auxiliary_hidden_states] + for index, tensor in enumerate(outputs): + if tensor.ndim == 0 or tensor.shape[0] != num_tokens: + raise RuntimeError( + "DSpark FullCG body output must use the captured token " + f"axis: tensor={index} shape={tuple(tensor.shape)} " + f"tokens={num_tokens}" + ) + + buffers = self._dspark_full_prefill_output_buffers + if buffers is None: + if num_tokens != self.max_num_tokens: + raise RuntimeError( + "DSpark FullCG shared outputs must initialize from the " + f"largest capture shape: got={num_tokens} " + f"expected={self.max_num_tokens}" + ) + allocated = [ + torch.empty( + (self.max_num_tokens, *tensor.shape[1:]), + dtype=tensor.dtype, + device=tensor.device, + ) + for tensor in outputs + ] + buffers = ((allocated[0], allocated[1]), allocated[2:]) + self._dspark_full_prefill_output_buffers = buffers + total_bytes = sum( + tensor.numel() * tensor.element_size() for tensor in allocated + ) + logger.info( + "Reserved shared DSpark FullCG body outputs: " + "tokens=%d tensors=%d bytes=%d", + self.max_num_tokens, + len(allocated), + total_bytes, + ) + + flat_buffers = [*buffers[0], *buffers[1]] + if len(flat_buffers) != len(outputs): + raise RuntimeError("DSpark FullCG shared output tree changed after warmup") + staged = [] + for index, (buffer, tensor) in enumerate(zip(flat_buffers, outputs)): + if ( + buffer.shape[1:] != tensor.shape[1:] + or buffer.dtype != tensor.dtype + or buffer.device != tensor.device + ): + raise RuntimeError( + "DSpark FullCG shared output tensor changed after warmup: " + f"tensor={index}" + ) + view = buffer[:num_tokens] + view.copy_(tensor) + staged.append(view) + return (staged[0], staged[1]), staged[2:] + @torch.no_grad() def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int): """Run forward inside the prefill set_tc_piecewise_forward_context. @@ -499,16 +661,27 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int): ) set_is_extend_in_batch(False) + if self._is_full_backend: + # Record address-stable metadata work exactly once per replay. + # DSV4 uses this to translate the SWA write locations once rather + # than allocating/recomputing them in every transformer layer. + self.model_runner.attn_backend.init_forward_metadata_in_graph(forward_batch) + with self._prefill_forward_context(forward_batch): if self._uses_eager_prefill_tail(): # BCG / Full: capture the transformer body only. positions = self._get_layer_model_positions(forward_batch) - return self.layer_model.forward( + body_output = self.layer_model.forward( forward_batch.input_ids, positions, forward_batch, forward_batch.input_embeds, ) + if self._share_dspark_full_prefill_outputs: + return self._stage_dspark_full_prefill_body_output( + body_output, num_tokens=num_tokens + ) + return body_output # tc_piecewise: compile/capture the outer model.forward path. return self.model_runner.model.forward( forward_batch.input_ids, @@ -594,8 +767,9 @@ def run_dummy_multimodal_deepstack_forward( def _has_inactive_dp_rank(self, forward_batch: ForwardBatch) -> bool: # DSV4 DP attention / DeepEP collectives need every DP rank to enter - # the same replay path. Sparse-DP batches (one or more ranks with - # zero local tokens) fall back to eager to avoid hanging ranks. + # the same replay path. Admitted sparse-DP prefills are converted to + # MAX_LEN before this check, so an inactive rank here is a contract + # violation rather than a reason to silently choose eager. global_num_tokens = forward_batch.global_num_tokens_cpu if global_num_tokens is None: return False @@ -603,6 +777,18 @@ def _has_inactive_dp_rank(self, forward_batch: ForwardBatch) -> bool: int(num_tokens) == 0 for num_tokens in global_num_tokens ) + @staticmethod + def _reject_graph(forward_batch: ForwardBatch, reason: str) -> bool: + # The scheduler synchronizes this admission across DP ranks before + # padding. Once true, a local rejection would make one rank enter the + # graph while another enters eager collectives. Fail closed instead. + if forward_batch.can_run_dp_breakable_cuda_graph: + raise RuntimeError( + "DP prefill CUDA graph admission diverged after synchronization: " + f"{reason}" + ) + return False + def _init_forward_metadata_for_capture( self, forward_batch: ForwardBatch, num_tokens: int ) -> None: @@ -636,20 +822,40 @@ def _prepare_forward_metadata_for_replay( attn_backend = self.model_runner.attn_backend if self._is_full_backend: # Slot-padded shallow view: plan() must see exactly req_slots - # entries (real values in [:bs], sentinels in [bs:req_slots] - # already populated by replay_prepare). - r = self._capture_req_slots + # entries and the captured token bucket. Build it from the static + # batch so out_cache_loc/positions retain the padded token-axis + # storage captured by the graph. The live batch is intentionally + # raw-sized and cannot be used to refresh fixed-shape metadata. + r = self._capture_req_slots_for_tokens(num_tokens) bs = forward_batch.batch_size s = self._prefill_static_buffers - self._full_cg_seq_lens_cpu.zero_() + assert s is not None + assert self._full_cg_seq_lens_cpu is not None + if bs > r: + raise RuntimeError( + "full prefill graph request count exceeds the selected " + f"shape slots: requests={bs}, slots={r}, tokens={num_tokens}" + ) + assert forward_batch.seq_lens_cpu is not None + assert forward_batch.extend_seq_lens_cpu is not None + assert forward_batch.extend_prefix_lens_cpu is not None + self._full_cg_seq_lens_cpu[:r].zero_() self._full_cg_seq_lens_cpu[:bs].copy_(forward_batch.seq_lens_cpu) - padded_view = copy.copy(forward_batch) + padded_view = copy.copy(static_forward_batch) padded_view.batch_size = r padded_view.seq_lens = s["seq_lens"][:r] - padded_view.seq_lens_cpu = self._full_cg_seq_lens_cpu + padded_view.seq_lens_cpu = self._full_cg_seq_lens_cpu[:r] padded_view.req_pool_indices = s["req_pool_indices"][:r] padded_view.extend_seq_lens = s["extend_seq_lens"][:r] padded_view.extend_prefix_lens = s["extend_prefix_lens"][:r] + padded_view.extend_start_loc = s["extend_start_loc"][:r] + padded_view.orig_seq_lens = s["orig_seq_lens"][:r] + padded_view.extend_seq_lens_cpu = list( + forward_batch.extend_seq_lens_cpu + ) + [0] * (r - bs) + padded_view.extend_prefix_lens_cpu = list( + forward_batch.extend_prefix_lens_cpu + ) + [0] * (r - bs) attn_backend.init_forward_metadata_out_graph(padded_view) return if not self.use_captured_attn_metadata: @@ -679,38 +885,65 @@ def _restore_mha_capture_state(forward_batch: ForwardBatch) -> None: forward_batch.set_attn_attend_prefix_cache(False) def can_run_graph(self, forward_batch: ForwardBatch) -> bool: - if self._is_full_backend and forward_batch.batch_size > self._capture_req_slots: - return False if forward_batch.input_embeds is not None: - return False + return self._reject_graph(forward_batch, "input embeddings are unsupported") if forward_batch.replace_embeds is not None: - return False + return self._reject_graph( + forward_batch, "replacement embeddings are unsupported" + ) if self._has_unsupported_mha_prefix(forward_batch): - return False + return self._reject_graph( + forward_batch, "MHA companion prefix is unsupported" + ) # tc_piecewise captures with ForwardMode.EXTEND and spec_info=None. if forward_batch.forward_mode.is_target_verify(): - return False + return self._reject_graph(forward_batch, "target verify is unsupported") + if self.enable_tbo_prefill_graph and not forward_batch.can_run_tbo: + return self._reject_graph( + forward_batch, + "captured dense-prefill graph requires synchronized TBO", + ) if forward_batch.capture_hidden_mode != self.capture_hidden_mode: - return False - # BCG-with-captured-metadata under DP attention: every rank must - # have local tokens, and the batch must declare itself replayable. - # These gates are no-ops for non-DP / non-opt-in paths because - # global_num_tokens_cpu stays None. + return self._reject_graph( + forward_batch, "capture-hidden mode differs from capture" + ) + # Captured prefill under DP attention: every rank must have local + # tokens after MAX_LEN conversion, and the batch must carry the shared + # admission decision. if self._has_inactive_dp_rank(forward_batch): - return False + return self._reject_graph( + forward_batch, "a DP rank remained inactive after padding" + ) if ( forward_batch.global_num_tokens_cpu is not None and not forward_batch.can_run_dp_breakable_cuda_graph ): - return False + return self._reject_graph( + forward_batch, "scheduler routed this shape to eager" + ) num_tokens = len(forward_batch.input_ids) if forward_batch.return_logprob and not self._uses_eager_prefill_tail(): - return False + return self._reject_graph( + forward_batch, "return_logprob requires an eager logits tail" + ) if num_tokens > self.max_num_tokens: - return False + return self._reject_graph( + forward_batch, "padded token count exceeds captured buckets" + ) padded_num_tokens = self._pad_to_bucket(num_tokens, self.capture_num_tokens) if padded_num_tokens > num_tokens * _MAX_PREFILL_CUDA_GRAPH_PADDING_FACTOR: - return False + return self._reject_graph( + forward_batch, "capture-bucket padding exceeds the admitted factor" + ) + if ( + self._is_full_backend + and forward_batch.batch_size + > self._capture_req_slots_for_tokens(padded_num_tokens) + ): + return self._reject_graph( + forward_batch, + "request count exceeds the selected full-graph shape slots", + ) # No exact-shape check here: load_batch bucket-pads to the nearest # captured shape. The factor above only rejects replays whose padded # model work is disproportionate to the useful token count. @@ -740,7 +973,7 @@ def capture_prepare(self, num_tokens: int) -> tuple[ForwardBatch, AttentionBacke Returns ``(forward_batch, attn_backend)`` to mirror decode's capture_prepare signature. """ - bs = self._capture_req_slots + bs = self._capture_req_slots_for_tokens(num_tokens) # Slot 0 carries num_tokens; slots 1..bs-1 are zero-length sentinels. lens_cpu = [num_tokens] + [0] * (bs - 1) start_loc_cpu = [0] + [num_tokens] * (bs - 1) @@ -853,7 +1086,12 @@ def _slot(name): lora_ids=None, return_pooled_hidden_states=self.capture_return_pooled_hidden_states, ) - self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens) + self.tbo_plugin.capture_one_batch_size( + forward_batch, + num_tokens=num_tokens, + enabled=self.enable_tbo_prefill_graph, + attn_backend=self.model_runner.attn_backend, + ) return forward_batch, self.model_runner.attn_backend def capture(self) -> None: @@ -944,6 +1182,11 @@ def load_batch(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch: raw_num_tokens=num_tokens, padded_num_tokens=static_num_tokens, ) + if self.enable_tbo_prefill_graph: + self.tbo_plugin.replay_prepare_prefill( + split_token_index=static_num_tokens // 2, + num_token_non_padded=num_tokens, + ) registry = self.buffer_registry @@ -1078,20 +1321,27 @@ def _slot(name): s["req_pool_indices"][:bs].copy_(forward_batch.req_pool_indices) if forward_batch.orig_seq_lens is not None: s["orig_seq_lens"][:bs].copy_(forward_batch.orig_seq_lens) - if self._is_full_backend and bs < self._capture_req_slots: + if self._is_full_backend: # Sentinel tail for slots [bs:req_slots]: the captured graph # reads all req_slots entries (e.g. the logits-processor # cumsum), so stale values from the previous replay must be # cleared. Zero lengths make the sentinels no-ops; # extend_start_loc sentinels sit at the flat end of the real # tokens. - r = self._capture_req_slots - s["seq_lens"][bs:r].zero_() - s["extend_seq_lens"][bs:r].zero_() - s["extend_prefix_lens"][bs:r].zero_() - s["extend_start_loc"][bs:r].fill_(self.raw_num_tokens) - s["req_pool_indices"][bs:r].zero_() - s["orig_seq_lens"][bs:r].zero_() + r = self._capture_req_slots_for_tokens(static_num_tokens) + if bs > r: + raise RuntimeError( + "full prefill graph request count exceeds the selected " + f"shape slots: requests={bs}, slots={r}, " + f"tokens={static_num_tokens}" + ) + if bs < r: + s["seq_lens"][bs:r].zero_() + s["extend_seq_lens"][bs:r].zero_() + s["extend_prefix_lens"][bs:r].zero_() + s["extend_start_loc"][bs:r].fill_(self.raw_num_tokens) + s["req_pool_indices"][bs:r].zero_() + s["orig_seq_lens"][bs:r].zero_() # Refresh the static buffer the captured graph reads from. if ( @@ -1109,6 +1359,46 @@ def _slot(name): self._static_num_tokens = static_num_tokens return static_forward_batch + @classmethod + def _trim_replayed_body_output( + cls, + output, + *, + raw_num_tokens: int, + static_num_tokens: int, + ): + """Trim token-axis tensors without changing the model output tree. + + Transformer bodies may return a tensor, nested tuples, or lists of + auxiliary hidden states. Slicing the top-level object corrupts tuple + structure (and turns a two-item tuple into a one-item tuple when the + raw batch has one token). Only tensors whose leading dimension is the + captured token bucket are token-axis outputs. + """ + if isinstance(output, torch.Tensor): + if output.ndim > 0 and output.shape[0] == static_num_tokens: + return output[:raw_num_tokens] + return output + if isinstance(output, tuple): + return tuple( + cls._trim_replayed_body_output( + item, + raw_num_tokens=raw_num_tokens, + static_num_tokens=static_num_tokens, + ) + for item in output + ) + if isinstance(output, list): + return [ + cls._trim_replayed_body_output( + item, + raw_num_tokens=raw_num_tokens, + static_num_tokens=static_num_tokens, + ) + for item in output + ] + return output + def _execute_body_capture( self, forward_batch: ForwardBatch, @@ -1142,7 +1432,15 @@ def replay_layer_forward(*args, **layer_kwargs): 1, static_n )[: ie.shape[0]].copy_(ie) hs = self.backend.replay(shape_key, static_forward_batch, **kwargs) - return hs[:raw_num_tokens] if full_path else hs + return ( + self._trim_replayed_body_output( + hs, + raw_num_tokens=raw_num_tokens, + static_num_tokens=static_n, + ) + if full_path + else hs + ) original_layer_forward = self.layer_model.forward self.layer_model.forward = replay_layer_forward diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py index e4f13381e238..7fc0ff61adc0 100644 --- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py @@ -57,6 +57,7 @@ def __init__( ) -> None: self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {} self._outputs: Dict[Any, Any] = {} + self._capture_inputs: Dict[Any, Any] = {} self._pool = None self._device_module = cuda_graph_runner.device_module self._tp_group = cuda_graph_runner.model_runner.tp_group @@ -90,6 +91,12 @@ def capture_one( self._device_module.synchronize() self._tp_group.barrier() forward_fn() + # Some attention backends materialize per-shape CUDA tensors during + # the warmup and let post_warmup_hook replace their Python owners. + # Their custom kernels are asynchronous and are not necessarily + # visible to PyTorch's allocator stream tracking. Finish all work + # before the hook can release/reuse those captured addresses. + self._device_module.synchronize() if post_warmup_hook is not None: post_warmup_hook() @@ -110,6 +117,12 @@ def capture_one( with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream): out = forward_fn() + # A CUDA graph retains device addresses, not Python tensor owners. + # Keep every per-shape ForwardBatch reachable for the graph's lifetime; + # callers deliberately pass capture_inputs for tensors that are not in + # a runner-owned static buffer (for example DP global token counts). + if capture_inputs is not None: + self._capture_inputs[shape_key] = capture_inputs self._graphs[shape_key] = graph self._outputs[shape_key] = out @@ -132,4 +145,5 @@ def replay( def cleanup(self) -> None: self._graphs.clear() self._outputs.clear() + self._capture_inputs.clear() self._pool = None diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 4014e4e6c89b..b1caa7966650 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -705,8 +705,12 @@ def _compute_q_b( ) -> torch.Tensor: q, _ = self.wq_b(q) q = q.view(-1, self.n_local_heads, self.head_dim) + # Eager prefill has no padded attention-head output buffer to target. + # The fused kernel loads a complete head before storing and explicitly + # supports aliasing, so normalize/RoPE in place instead of allocating + # another full Q tensor (1 GiB at the 16K B300 prefill chunk). if q_out is None: - q_out = torch.empty_like(q) + q_out = q # Fused warp-per-(token, head) rmsnorm-self + RoPE + write to q_out. fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions) return q_out @@ -2211,7 +2215,7 @@ def _forward_layers_tbo( positions: torch.Tensor, hidden_states: torch.Tensor, forward_batch: ForwardBatch, - ) -> torch.Tensor: + ) -> Tuple[torch.Tensor, List[torch.Tensor]]: from sglang.srt.batch_overlap.operations import execute_overlapped_operations from sglang.srt.batch_overlap.operations_strategy import OperationsStrategy from sglang.srt.batch_overlap.two_batch_overlap import ( @@ -2219,10 +2223,17 @@ def _forward_layers_tbo( _model_forward_tbo_merge_outputs, ) - layers = [self.layers[i] for i in range(self.start_layer, self.end_layer)] - operations_strategy = OperationsStrategy.init_new_tbo( - layers, forward_batch.global_forward_mode - ) + capture_layer_ids = set(self.dspark_layers_to_capture or []) + layer_groups = [] + current_group = [] + for layer_id in range(self.start_layer, self.end_layer): + current_group.append(self.layers[layer_id]) + if layer_id in capture_layer_ids: + layer_groups.append((current_group, layer_id)) + current_group = [] + if current_group: + layer_groups.append((current_group, None)) + assert layer_groups, "DSV4 TBO requires at least one local decoder layer" # Split the per-rank batch into the 2 ubatches (token-range slice + pad # to tbo_padded_len). residual is unused by the DSV4 non-fused layer ops. @@ -2246,24 +2257,29 @@ def _forward_layers_tbo( tp_group = get_tp_group() world = tp_group.world_size children = forward_batch.tbo_children - local_lens = torch.tensor( - [int(c.tbo_padded_len) for c in children], - dtype=torch.int64, - device=hidden_states.device, - ) - gathered = torch.empty( - (world, local_lens.shape[0]), - dtype=torch.int64, - device=hidden_states.device, - ) - tp_group.all_gather_into_tensor(gathered, local_lens) - gathered_cpu = gathered.tolist() + if any(child.global_num_tokens_cpu is None for child in children): + local_lens = torch.tensor( + [int(c.tbo_padded_len) for c in children], + dtype=torch.int64, + device=hidden_states.device, + ) + gathered = torch.empty( + (world, local_lens.shape[0]), + dtype=torch.int64, + device=hidden_states.device, + ) + tp_group.all_gather_into_tensor(gathered, local_lens) + gathered_cpu = gathered.tolist() + for idx, child in enumerate(children): + sizes = [gathered_cpu[r][idx] for r in range(world)] + child.global_num_tokens_cpu = sizes + child.global_num_tokens_gpu = gathered[:, idx].contiguous() + child.global_dp_buffer_len = sum(sizes) + rank = tp_group.rank_in_group - for idx, child in enumerate(children): - sizes = [gathered_cpu[r][idx] for r in range(world)] - child.global_num_tokens_cpu = sizes - child.global_num_tokens_gpu = gathered[:, idx].contiguous() - child.global_dp_buffer_len = sum(sizes) + for child in children: + sizes = child.global_num_tokens_cpu + assert sizes is not None and len(sizes) == world # Gather the ubatch's input_ids -> global ONCE here (cached on the # child) instead of per-layer in op_gather_a. The hash MoE reads # the SAME global ids every layer, so 61x2 per-layer all_gatherv of @@ -2284,16 +2300,40 @@ def _forward_layers_tbo( tp_group.all_gatherv(padded_ids, sizes=sizes, output=gids) child._tbo_global_input_ids = gids - outputs_arr = execute_overlapped_operations( - inputs_arr=inputs_arr, - operations_arr=[operations_strategy.operations] * 2, - delta_stages=[0, operations_strategy.tbo_delta_stages], - ) + dspark_aux_hidden_states = [] + outputs_arr = inputs_arr + + def merge_dspark_aux() -> torch.Tensor: + merged = hidden_states.new_zeros( + (hidden_states.shape[0], hidden_states.shape[-1]) + ) + for output in outputs_arr: + start, end = output["forward_batch"].tbo_parent_token_range + merged[start:end] = output["hidden_states"][: end - start].mean(dim=1) + return merged + + for group_layers, captured_layer_id in layer_groups: + operations_strategy = OperationsStrategy.init_new_tbo( + group_layers, forward_batch.global_forward_mode + ) + outputs_arr = execute_overlapped_operations( + inputs_arr=outputs_arr, + operations_arr=[operations_strategy.operations] * 2, + delta_stages=[0, operations_strategy.tbo_delta_stages], + ) + if captured_layer_id is not None: + dspark_aux_hidden_states.append(merge_dspark_aux()) hidden_states, _ = _model_forward_tbo_merge_outputs( outputs_arr[0], outputs_arr[1], hidden_states.shape[0] ) - return hidden_states + if capture_layer_ids: + assert len(dspark_aux_hidden_states) == len(capture_layer_ids), ( + "DSpark TBO did not capture every configured target layer: " + f"configured={sorted(capture_layer_ids)}, " + f"captured={len(dspark_aux_hidden_states)}" + ) + return hidden_states, dspark_aux_hidden_states def forward( self, @@ -2348,13 +2388,10 @@ def forward( "of them: DSpark static-verify is CP-off for v1." ) dspark_aux_hidden_states: List[torch.Tensor] = [] - # DSpark aux capture needs the per-layer eager loop (TBO's overlapped - # execution cannot expose per-layer completed hidden states), so skip - # TBO when capturing -- a perf-only downgrade, not a correctness one. - if self._can_run_tbo(forward_batch) and not capture_dspark: + if self._can_run_tbo(forward_batch): # Two-batch-overlap prefill (EP / mori). Cross-layer mHC fusion is # disabled here (each layer self-contained), so no trailing hc_post. - hidden_states = self._forward_layers_tbo( + hidden_states, dspark_aux_hidden_states = self._forward_layers_tbo( positions=positions, hidden_states=hidden_states, forward_batch=forward_batch, @@ -2422,6 +2459,8 @@ def forward( class DeepseekV4ForCausalLM(nn.Module): + dynamic_kv_cache_scale_block_size = 64 + def __init__( self, config: DeepSeekV4Config, diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index dc2bdbf19dcf..8fa7aea42e5a 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -156,8 +156,10 @@ def _compute_q( q, _ = self.wq_b(q) q = q.view(-1, self.n_local_heads, self.head_dim) if self._use_fast_kernel: + # Match the target model's allocation-free eager path. Capture and + # TP-padded calls still supply their fixed q_out buffer. if q_out is None: - q_out = torch.empty_like(q) + q_out = q fused_q_norm_rope(q, q_out, self.eps, self.freqs_cis, positions) return q_out else: @@ -553,6 +555,7 @@ def _run_ffn(self, x: torch.Tensor, forward_batch: ForwardBatch) -> torch.Tensor class DeepseekV4ForCausalLMDSpark(nn.Module): + dynamic_kv_cache_scale_block_size = 64 def __init__( self, diff --git a/python/sglang/srt/observability/metrics_collector.py b/python/sglang/srt/observability/metrics_collector.py index 0552c46a54fa..a43c131611db 100644 --- a/python/sglang/srt/observability/metrics_collector.py +++ b/python/sglang/srt/observability/metrics_collector.py @@ -302,6 +302,39 @@ def __init__( labelnames=labels.keys(), multiprocess_mode="mostrecent", ) + self.scheduler_phase_seconds_total = Counter( + name="sglang:scheduler_phase_seconds_total", + documentation=( + "Cumulative scheduler-thread wall time by phase. The phases " + "partition the overlap event loop except that loop_total is " + "also emitted as an authoritative end-to-end total." + ), + labelnames=[*labels.keys(), "phase"], + ) + self.scheduler_phase_calls_total = Counter( + name="sglang:scheduler_phase_calls_total", + documentation="Number of scheduler-thread phase executions.", + labelnames=[*labels.keys(), "phase"], + ) + self.scheduler_phase_max_seconds = Gauge( + name="sglang:scheduler_phase_max_seconds", + documentation=( + "Maximum scheduler-thread phase wall time in the latest " + "one-second reporting window." + ), + labelnames=[*labels.keys(), "phase"], + multiprocess_mode="mostrecent", + ) + self.runtime_gc_frozen = Gauge( + name="sglang:runtime_gc_frozen", + documentation=( + "1 after the warmed scheduler object graph has been moved to " + "Python's permanent GC generation, else 0." + ), + labelnames=labels.keys(), + multiprocess_mode="mostrecent", + ) + self.runtime_gc_frozen.labels(**labels).set(0) # ================================================================= # Memory pool usage ratios @@ -786,6 +819,79 @@ def __init__( documentation="Number of the total grammar requests.", labelnames=labels.keys(), ) + self.grammar_cache_lookup_total = Counter( + name="sglang:grammar_cache_lookup_total", + documentation=( + "Grammar cache resolutions by source: compile, local_compile, " + "disk, inflight, or memory." + ), + labelnames=list(labels.keys()) + ["source"], + ) + grammar_cache_time_buckets = [ + 0.0, + 0.00001, + 0.00002, + 0.00005, + 0.0001, + 0.0002, + 0.0005, + 0.001, + 0.002, + 0.005, + 0.01, + 0.02, + 0.05, + 0.1, + 0.2, + 0.5, + 1, + 2, + 5, + ] + self.grammar_cache_resolution_time = Histogram( + name="sglang:grammar_cache_resolution_time_seconds", + documentation="End-to-end grammar cache resolution latency.", + labelnames=list(labels.keys()) + ["source"], + buckets=grammar_cache_time_buckets, + ) + self.grammar_cache_lock_wait_time = Histogram( + name="sglang:grammar_cache_lock_wait_time_seconds", + documentation="Cross-process persistent grammar cache lock wait time.", + labelnames=labels.keys(), + buckets=grammar_cache_time_buckets, + ) + self.grammar_cache_phase_time = Histogram( + name="sglang:grammar_cache_phase_time_seconds", + documentation="Internal persistent grammar cache phase latency.", + labelnames=list(labels.keys()) + ["phase"], + buckets=grammar_cache_time_buckets, + ) + self.dspark_grammar_pipeline_seconds = Histogram( + name="sglang:dspark_grammar_pipeline_seconds", + documentation=( + "Post-launch wall latency of each DSpark grammar " + "pipeline phase; phases overlap target verification." + ), + labelnames=list(labels.keys()) + ["phase"], + buckets=grammar_cache_time_buckets, + ) + self.dspark_grammar_active_matchers = Histogram( + name="sglang:dspark_grammar_active_matchers", + documentation="Native XGrammar matchers traversed per DSpark decode step.", + labelnames=labels.keys(), + buckets=[0, 1, 2, 4, 8, 16, 32, 64, 96, 128, 192, 256, 384, 512], + ) + self.dspark_grammar_batch_size = Histogram( + name="sglang:dspark_grammar_batch_size", + documentation="Requests in each structured DSpark decode step.", + labelnames=labels.keys(), + buckets=[1, 2, 4, 8, 16, 32, 64, 96, 128, 192, 256, 384, 512], + ) + self.dspark_grammar_steps_total = Counter( + name="sglang:dspark_grammar_steps_total", + documentation="Optimized DSpark grammar steps by mask outcome.", + labelnames=list(labels.keys()) + ["outcome"], + ) self.grammar_schema_count = Histogram( name="sglang:grammar_schema_count", documentation="Histogram of grammar schema count.", @@ -897,6 +1003,93 @@ def __init__( ), labelnames=list(labels.keys()) + ["category"], ) + self.prefill_graph_admissions_total = Counter( + name="sglang:prefill_graph_admissions_total", + documentation=( + "Prefill batches by CUDA graph admission outcome and captured " + "token bucket. bucket='none' denotes eager execution." + ), + labelnames=list(labels.keys()) + ["outcome", "bucket"], + ) + self.prefill_graph_shapes_total = Counter( + name="sglang:prefill_graph_shapes_total", + documentation=( + "Full-prefill CUDA graph admissions jointly classified by " + "captured token bucket and bounded request-slot shape." + ), + labelnames=list(labels.keys()) + ["bucket", "request_slots"], + ) + self.prefill_execution_tokens_total = Counter( + name="sglang:prefill_execution_tokens_total", + documentation=( + "Scheduled and actually executed prefill token-axis elements " + "by graph/eager path and capture bucket. Their ratio exposes " + "graph padding cost at the exact shape that caused it." + ), + labelnames=list(labels.keys()) + ["path", "kind", "bucket"], + ) + prefill_token_buckets = ( + 1, + 2, + 4, + 8, + 16, + 32, + 64, + 96, + 128, + 192, + 256, + 384, + 512, + 768, + 1024, + 1536, + 2048, + 3072, + 4096, + 6144, + 8192, + 12288, + 16384, + ) + self.prefill_batch_scheduled_tokens = Histogram( + name="sglang:prefill_batch_scheduled_tokens", + documentation="Local scheduled tokens in each prefill batch.", + labelnames=list(labels.keys()) + ["path"], + buckets=prefill_token_buckets, + ) + self.prefill_batch_requests = Histogram( + name="sglang:prefill_batch_requests", + documentation="Local request count in each prefill batch.", + labelnames=list(labels.keys()) + ["path"], + # Include each captured request-slot boundary and its immediate + # predecessor. Their cumulative-bucket difference proves that a + # live batch exercised the exact shape instead of merely landing + # somewhere below it. + buckets=( + 1, + 2, + 3, + 4, + 7, + 8, + 15, + 16, + 31, + 32, + 63, + 64, + 95, + 96, + 127, + 128, + 191, + 192, + 255, + 256, + ), + ) self.estimated_flops_per_gpu_total = Counter( name="sglang:estimated_flops_per_gpu_total", documentation=( @@ -1154,6 +1347,21 @@ def observe_per_stage_req_latency(self, stage: str, latency: float) -> None: def observe_queue_time(self, latency: float) -> None: self._log_histogram(self.queue_time, latency) + def add_scheduler_phase( + self, + phase: str, + duration_seconds: float, + calls: int, + max_seconds: float, + ) -> None: + labels = {**self.labels, "phase": phase} + self.scheduler_phase_seconds_total.labels(**labels).inc(duration_seconds) + self.scheduler_phase_calls_total.labels(**labels).inc(calls) + self.scheduler_phase_max_seconds.labels(**labels).set(max_seconds) + + def set_runtime_gc_frozen(self) -> None: + self.runtime_gc_frozen.labels(**self.labels).set(1) + def observe_weight_load(self, duration_seconds: float, source: str) -> None: # Edge-triggered: engine is paused during the update, so log_stats # won't fire — write the gauge inline at end of update_weights_from_*. @@ -1207,6 +1415,42 @@ def increment_prefill_cuda_graph_pass(self, value: bool) -> None: mode = "prefill_cuda_graph" if value else "prefill_none" self.cuda_graph_passes_total.labels(**self.labels, mode=mode).inc(1) + def observe_prefill_execution( + self, + *, + outcome: str, + scheduled_tokens: int, + executed_tokens: int, + requests: int, + bucket_tokens: int, + ) -> None: + path = "cuda_graph" if outcome == "cuda_graph" else "eager" + bucket = str(bucket_tokens) if bucket_tokens > 0 else "none" + self.prefill_graph_admissions_total.labels( + **self.labels, outcome=outcome, bucket=bucket + ).inc(1) + if outcome == "cuda_graph": + request_slots = ( + str(requests) + if requests in (0, 1, 2, 4, 8, 16, 32, 64, 96, 128, 192, 256) + else "other" + ) + self.prefill_graph_shapes_total.labels( + **self.labels, + bucket=bucket, + request_slots=request_slots, + ).inc(1) + self.prefill_execution_tokens_total.labels( + **self.labels, path=path, kind="scheduled", bucket=bucket + ).inc(scheduled_tokens) + self.prefill_execution_tokens_total.labels( + **self.labels, path=path, kind="executed", bucket=bucket + ).inc(executed_tokens) + self.prefill_batch_scheduled_tokens.labels(**self.labels, path=path).observe( + scheduled_tokens + ) + self.prefill_batch_requests.labels(**self.labels, path=path).observe(requests) + def increment_eplb_balancedness( self, forward_mode: str, balancedness: float ) -> None: @@ -1396,8 +1640,59 @@ def log_grammar_stats(self, grammar_stats) -> None: self.num_grammar_timeout.labels(**self.labels).inc( grammar_stats.num_timeout ) + source = grammar_stats.cache_source + self.grammar_cache_lookup_total.labels( + **self.labels, + source=source, + ).inc(1) + if grammar_stats.cache_resolution_time is not None: + self.grammar_cache_resolution_time.labels( + **self.labels, + source=source, + ).observe(grammar_stats.cache_resolution_time) + if grammar_stats.cache_lock_wait_time is not None: + self._log_histogram( + self.grammar_cache_lock_wait_time, + grammar_stats.cache_lock_wait_time, + ) + for phase, phase_seconds in grammar_stats.cache_phase_seconds.items(): + self.grammar_cache_phase_time.labels( + **self.labels, + phase=phase, + ).observe(phase_seconds) self.num_grammar_total.labels(**self.labels).inc(1) + def observe_dspark_grammar_step( + self, + *, + batch_size: int, + active_matchers: int, + outcome_counts: Dict[str, int], + phase_seconds: Dict[str, float], + ) -> None: + self._log_histogram(self.dspark_grammar_batch_size, batch_size) + self._log_histogram( + self.dspark_grammar_active_matchers, + active_matchers, + ) + for outcome, count in outcome_counts.items(): + if count > 0: + self.dspark_grammar_steps_total.labels( + **self.labels, + outcome=outcome, + ).inc(count) + for phase, seconds in phase_seconds.items(): + self.dspark_grammar_pipeline_seconds.labels( + **self.labels, + phase=phase, + ).observe(seconds) + + def observe_dspark_grammar_gpu_gap(self, *, seconds: float) -> None: + self.dspark_grammar_pipeline_seconds.labels( + **self.labels, + phase="gpu_target_to_mask_ready", + ).observe(seconds) + def emit_constants( self, max_total_num_tokens: int, @@ -1869,9 +2164,11 @@ def __init__( ) -> None: # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` from prometheus_client import Counter as _PromCounter + from prometheus_client import Gauge as _PromGauge from prometheus_client import Histogram as _PromHistogram Counter = self._counter_cls or _PromCounter + Gauge = self._gauge_cls or _PromGauge Histogram = self._histogram_cls or _PromHistogram self.labels = labels @@ -1949,6 +2246,44 @@ def __init__( documentation="The number of tokens loaded from CPU to GPU.", labelnames=labels.keys(), ) + self.hicache_backup_duration_seconds = Histogram( + name="sglang:hicache_backup_duration_seconds", + documentation="GPU duration of asynchronous HiCache D-to-H backups.", + labelnames=labels.keys(), + buckets=bucket_load_back_duration, + ) + self.hicache_backup_num_tokens = Counter( + name="sglang:hicache_backup_tokens_total", + documentation="The number of tokens backed up from GPU to host HiCache.", + labelnames=labels.keys(), + ) + self.hicache_scheduler_phase_seconds_total = Counter( + name="sglang:hicache_scheduler_phase_seconds_total", + documentation=( + "Scheduler-thread wall time spent in HiCache control phases." + ), + labelnames=[*labels.keys(), "phase"], + ) + self.hicache_scheduler_phase_calls_total = Counter( + name="sglang:hicache_scheduler_phase_calls_total", + documentation="Number of HiCache scheduler control-phase executions.", + labelnames=[*labels.keys(), "phase"], + ) + self.hicache_scheduler_phase_max_seconds = Gauge( + name="sglang:hicache_scheduler_phase_max_seconds", + documentation=( + "Maximum HiCache scheduler control-phase wall time in the " + "latest one-second reporting window." + ), + labelnames=[*labels.keys(), "phase"], + multiprocess_mode="mostrecent", + ) + self.hicache_pending_operations = Gauge( + name="sglang:hicache_pending_operations", + documentation="Current number of pending HiCache operations by kind.", + labelnames=[*labels.keys(), "kind"], + multiprocess_mode="mostrecent", + ) def increment_eviction_num_tokens(self, num_tokens: int) -> None: self.eviction_num_tokens.labels(**self.labels).inc(num_tokens) @@ -1962,6 +2297,35 @@ def observe_eviction_duration(self, duration_seconds: float) -> None: def observe_load_back_duration(self, duration_seconds: float) -> None: self.load_back_duration_seconds.labels(**self.labels).observe(duration_seconds) + def observe_hicache_backup( + self, num_tokens: int, duration_seconds: Optional[float] + ) -> None: + self.hicache_backup_num_tokens.labels(**self.labels).inc(num_tokens) + if duration_seconds is not None: + self.hicache_backup_duration_seconds.labels(**self.labels).observe( + duration_seconds + ) + + def observe_hicache_scheduler_phase( + self, + phase: str, + duration_seconds: float, + calls: int, + max_seconds: float, + ) -> None: + labels = {**self.labels, "phase": phase} + self.hicache_scheduler_phase_seconds_total.labels(**labels).inc( + duration_seconds + ) + self.hicache_scheduler_phase_calls_total.labels(**labels).inc(calls) + self.hicache_scheduler_phase_max_seconds.labels(**labels).set(max_seconds) + + def set_hicache_pending_operations(self, kind: str, count: int) -> None: + self.hicache_pending_operations.labels( + **self.labels, + kind=kind, + ).set(count) + class EncoderMetricsCollector(_StatLoggerDIMixin): """Metrics collector for the EPD encoder server (--encoder-only).""" diff --git a/python/sglang/srt/observability/ray_wrappers.py b/python/sglang/srt/observability/ray_wrappers.py index deb43ae74ed5..2ec484449d12 100644 --- a/python/sglang/srt/observability/ray_wrappers.py +++ b/python/sglang/srt/observability/ray_wrappers.py @@ -298,6 +298,7 @@ class RayRadixCacheMetricsCollector(RadixCacheMetricsCollector): """``RadixCacheMetricsCollector`` that emits via Ray's metric system.""" _counter_cls = RayCounterWrapper + _gauge_cls = RayGaugeWrapper _histogram_cls = RayHistogramWrapper diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 72162ba51556..16cd0b51a89d 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -2,7 +2,7 @@ import contextlib from dataclasses import dataclass -from typing import List, Optional +from typing import TYPE_CHECKING, List, Optional import torch @@ -13,6 +13,9 @@ from sglang.srt.speculative.spec_info import SpecInput, SpecInputType from sglang.srt.utils.common import is_pin_memory_available +if TYPE_CHECKING: + from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject + _OVERLAP_PLAN_STREAMS: dict[str, torch.cuda.Stream] = {} @@ -56,6 +59,8 @@ class DFlashDraftInputV2(SpecInput): verify_token_budget: Optional[int] = None + grammar: Optional["BaseGrammarObject"] = None + def __post_init__(self): super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT) # Spec v2 draft state itself does not change token accounting. diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 7df8edbd0f65..890ca4043be9 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -12,6 +12,7 @@ from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod from sglang.srt.layers.sampler import apply_custom_logit_processor from sglang.srt.managers.schedule_batch import Req +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.utils import is_cuda, is_musa DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" @@ -793,9 +794,19 @@ def build_dflash_verify_target_probs( return target_probs.view(bs, draft_token_num, -1).contiguous() -def validate_dflash_request(req: Req, enable_overlap: bool) -> Optional[str]: +def validate_dflash_request( + req: Req, + enable_overlap: bool, + spec_algorithm: SpeculativeAlgorithm, +) -> Optional[str]: if req.return_logprob: - return "DFLASH speculative decoding does not support return_logprob yet." + if not spec_algorithm.is_dspark(): + return "DFLASH speculative decoding does not support return_logprob yet." + # max_new_tokens=1 completes from the pure target-model prefill result, + # before speculative decode. Anything longer could reach DSpark + # verification, whose accepted-token logprobs are not reconstructed. + if req.sampling_params.max_new_tokens != 1: + return "DSpark return_logprob requires a one-token, prefill-only request." if enable_overlap and req.return_hidden_states: return "DFLASH speculative decoding does not support return_hidden_states yet." diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index acdbd7451340..4580daf79cd5 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -338,25 +338,51 @@ def propose( def run_idle_participation(self, batch: ScheduleBatch) -> None: if not self._dp_moe_sync or batch.global_num_tokens is None: return + global_bs = max(int(num_reqs) for num_reqs in batch.global_num_tokens) + if global_bs == 0: + return + + # An IDLE ForwardBatch selects the ordinary one-token decode graph, + # while the peer runs the gamma-wide DSpark draft TARGET_VERIFY graph. + # Those graphs contain different DP all-gather geometries; once a DP + # rank drains before its peer, launching them at the same collective + # sequence deadlocks both GPUs. Materialize the idle rank as a full + # dummy draft tier so graph selection, token width, and every in-graph + # collective exactly match the active rank. Slot/cache index zero is + # the allocator's reserved dummy target. device = self.draft_model_runner.device - empty_long = torch.empty((0,), dtype=torch.int64, device=device) + num_dummy_tokens = global_bs * self.gamma + dummy_tokens = torch.full( + (num_dummy_tokens,), + int(self._mask_token_id), + dtype=torch.int64, + device=device, + ) + dummy_slots = torch.zeros((global_bs,), dtype=torch.int64, device=device) + dummy_seq_lens = torch.ones((global_bs,), dtype=torch.int64, device=device) idle_batch = ForwardBatch( - forward_mode=ForwardMode.IDLE, - batch_size=0, - input_ids=empty_long, - req_pool_indices=empty_long, - seq_lens=empty_long, - out_cache_loc=empty_long, - seq_lens_sum=0, - seq_lens_cpu=torch.empty((0,), dtype=torch.int64), - positions=empty_long, + forward_mode=ForwardMode.TARGET_VERIFY, + batch_size=global_bs, + input_ids=dummy_tokens, + req_pool_indices=dummy_slots, + seq_lens=dummy_seq_lens, + out_cache_loc=torch.zeros_like(dummy_tokens), + seq_lens_sum=global_bs * (self.gamma + 1), + seq_lens_cpu=torch.full((global_bs,), self.gamma + 1, dtype=torch.int64), + positions=torch.zeros_like(dummy_tokens), spec_algorithm=SpeculativeAlgorithm.DSPARK, spec_info=self._draft_block_spec_info, capture_hidden_mode=CaptureHiddenMode.NULL, ) self._fill_dp_moe_sync_metadata(idle_batch, batch) with torch.inference_mode(): - self.draft_model_runner.forward(idle_batch) + idle_out = self.draft_model_runner.forward(idle_batch) + if not idle_out.can_run_graph: + raise RuntimeError( + "DSpark DP idle draft participation missed the required CUDA " + f"graph for {global_bs=} and gamma={self.gamma}; an eager idle " + "forward cannot safely share collectives with an active graph." + ) def _run_forward( self, @@ -430,6 +456,7 @@ def _fill_dp_moe_sync_metadata( ) -> None: if not self._dp_moe_sync or batch.global_num_tokens is None: return + forward_batch.original_global_num_tokens_cpu = batch.global_num_tokens gnt, gnt_logprob = spec_scale_global_num_tokens( self._draft_block_spec_info, batch.global_num_tokens, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_grammar_pipeline.py b/python/sglang/srt/speculative/dspark_components/dspark_grammar_pipeline.py new file mode 100644 index 000000000000..38db49e857f6 --- /dev/null +++ b/python/sglang/srt/speculative/dspark_components/dspark_grammar_pipeline.py @@ -0,0 +1,544 @@ +"""Fail-closed, overlapped XGrammar mask construction for DSpark. + +The pipeline deliberately has no scalar or Python traversal fallback: +XGrammar must provide the batched draft-tree traversal API or startup fails. +""" + +from __future__ import annotations + +import logging +import math +import os +import time +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from typing import Protocol + +import torch +from xgrammar import BatchGrammarMatcher, GrammarMatcher + +from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject +from sglang.srt.constrained.reasoner_grammar_backend import ReasonerGrammarObject +from sglang.srt.constrained.xgrammar_backend import XGrammarGrammar + +logger = logging.getLogger(__name__) + + +class _RequestWithGrammar(Protocol): + grammar: BaseGrammarObject | None + + +class _DraftInputWithGrammar(Protocol): + grammar: BaseGrammarObject | None + + +class _GrammarMetricsCollector(Protocol): + def observe_dspark_grammar_step( + self, + *, + batch_size: int, + active_matchers: int, + outcome_counts: dict[str, int], + phase_seconds: dict[str, float], + ) -> None: ... + + def observe_dspark_grammar_gpu_gap(self, *, seconds: float) -> None: ... + + +@dataclass(slots=True) +class _GrammarBufferSlot: + draft_tokens_cpu: torch.Tensor + vocab_mask_cpu: torch.Tensor + vocab_mask_device: torch.Tensor + d2h_done: torch.cuda.Event + mask_ready: torch.cuda.Event + mask_consumed: torch.cuda.Event + has_mask_consumer: bool = False + + +@dataclass(slots=True) +class _GrammarGpuTimingSlot: + target_done: torch.cuda.Event + mask_dependency_done: torch.cuda.Event + reserved: bool = False + target_recorded: bool = False + dependency_recorded: bool = False + emit: bool = False + + +@dataclass(slots=True) +class GrammarMaskStep: + slot: _GrammarBufferSlot + batch_size: int + row_count: int + requests: Sequence[_RequestWithGrammar] + draft_input: _DraftInputWithGrammar + buffer_wait_seconds: float + gpu_timing: _GrammarGpuTimingSlot | None + + +@dataclass(slots=True) +class GrammarMaskResult: + step: GrammarMaskStep + grammar: XGrammarGrammar | ReasonerGrammarObject + vocab_mask: torch.Tensor + active_matchers: int + traversal_seconds: float + d2h_wait_seconds: float + + +class DSparkGrammarPipeline: + """Double-buffered CPU traversal and async mask transfer for one DP rank.""" + + def __init__( + self, + *, + device: str, + max_batch_size: int, + chain_length: int, + vocab_size: int, + traversal_threads: int, + metrics_collector: _GrammarMetricsCollector | None = None, + ) -> None: + if not hasattr(BatchGrammarMatcher, "batch_traverse_draft_tree"): + raise RuntimeError( + "DSpark grammar batching requires XGrammar BatchTraverseDraftTree" + ) + if not str(device).startswith("cuda"): + raise RuntimeError("DSpark grammar pipeline requires CUDA") + if max_batch_size <= 0 or chain_length <= 0 or vocab_size <= 0: + raise ValueError("invalid DSpark grammar buffer geometry") + if traversal_threads <= 0: + raise ValueError("grammar traversal thread count must be positive") + + self.device = device + self.max_batch_size = max_batch_size + self.chain_length = chain_length + self.vocab_size = vocab_size + self.mask_words = (vocab_size + 31) // 32 + self._device_module = torch.get_device_module(device) + self._copy_stream = self._device_module.Stream() + self._batch_matcher = BatchGrammarMatcher(max_threads=traversal_threads) + self._metrics_collector = metrics_collector + self._metrics_interval = int( + os.environ.get("SGLANG_DSPARK_GRAMMAR_METRICS_INTERVAL", "16") + ) + if self._metrics_interval <= 0: + raise ValueError("grammar metrics interval must be positive") + self._pending_outcomes = { + "masked": 0, + "thinking_only": 0, + } + self._next_token = torch.arange(1, chain_length + 1, dtype=torch.int64) + self._next_token[-1] = -1 + self._next_sibling = torch.full( + (chain_length,), + -1, + dtype=torch.int64, + ) + self._slot_index = 0 + self._steps = 0 + self._traversal_seconds = 0.0 + self._traversal_max_seconds = 0.0 + self._d2h_wait_seconds = 0.0 + self._buffer_wait_seconds = 0.0 + self._active_matchers = 0 + self._gpu_timing_step = 0 + self._gpu_timing_cursor = 0 + + mask_shape = (max_batch_size * chain_length, self.mask_words) + self._slots = [ + _GrammarBufferSlot( + draft_tokens_cpu=torch.empty( + (max_batch_size, chain_length), + dtype=torch.int64, + pin_memory=True, + ), + vocab_mask_cpu=torch.empty( + mask_shape, + dtype=torch.int32, + pin_memory=True, + ), + vocab_mask_device=torch.empty( + mask_shape, + dtype=torch.int32, + device=device, + ), + d2h_done=self._device_module.Event(), + mask_ready=self._device_module.Event(), + mask_consumed=self._device_module.Event(), + ) + for _ in range(2) + ] + self._gpu_timing_slots = [ + _GrammarGpuTimingSlot( + target_done=self._device_module.Event(enable_timing=True), + mask_dependency_done=self._device_module.Event(enable_timing=True), + ) + for _ in range(8) + ] + logger.info( + "DSpark grammar pipeline enabled: max_bs=%d chain=%d " + "vocab=%d threads=%d metrics_interval=%d double_buffer_bytes=%d", + max_batch_size, + chain_length, + vocab_size, + traversal_threads, + self._metrics_interval, + 2 + * ( + max_batch_size * chain_length * torch.int64.itemsize + + max_batch_size + * chain_length + * self.mask_words + * (torch.int32.itemsize * 2) + ), + ) + + def poll_gpu_timing_metrics(self) -> None: + for slot in self._gpu_timing_slots: + if not slot.dependency_recorded or not slot.mask_dependency_done.query(): + continue + if slot.emit: + elapsed_seconds = ( + slot.target_done.elapsed_time(slot.mask_dependency_done) / 1000 + ) + if not math.isfinite(elapsed_seconds) or elapsed_seconds < 0: + raise RuntimeError( + "DSpark grammar CUDA timing returned an invalid " + f"target-to-mask gap: {elapsed_seconds!r}" + ) + if self._metrics_collector is None: + raise RuntimeError( + "DSpark grammar CUDA timing has no metrics collector" + ) + self._metrics_collector.observe_dspark_grammar_gpu_gap( + seconds=elapsed_seconds + ) + slot.reserved = False + slot.target_recorded = False + slot.dependency_recorded = False + slot.emit = False + + def _reserve_gpu_timing(self) -> _GrammarGpuTimingSlot | None: + self._gpu_timing_step += 1 + if ( + self._metrics_collector is None + or self._gpu_timing_step % self._metrics_interval != 0 + ): + return None + self.poll_gpu_timing_metrics() + for offset in range(len(self._gpu_timing_slots)): + index = (self._gpu_timing_cursor + offset) % len(self._gpu_timing_slots) + slot = self._gpu_timing_slots[index] + if slot.reserved: + continue + slot.reserved = True + self._gpu_timing_cursor = (index + 1) % len(self._gpu_timing_slots) + return slot + raise RuntimeError( + "DSpark grammar CUDA timing ring exhausted before the GPU " + "completed prior samples" + ) + + def begin( + self, + *, + verify_ids_2d: torch.Tensor, + requests: Sequence[_RequestWithGrammar], + draft_input: _DraftInputWithGrammar, + ) -> GrammarMaskStep: + if verify_ids_2d.device.type != "cuda": + raise RuntimeError("DSpark verify ids must be CUDA-resident") + if verify_ids_2d.dtype != torch.int64 or not verify_ids_2d.is_contiguous(): + raise RuntimeError( + "DSpark grammar batching requires contiguous int64 verify ids" + ) + batch_size, chain_length = verify_ids_2d.shape + if chain_length != self.chain_length: + raise RuntimeError( + f"DSpark grammar chain changed: expected {self.chain_length}, " + f"got {chain_length}" + ) + if batch_size > self.max_batch_size or len(requests) != batch_size: + raise RuntimeError( + f"DSpark grammar batch geometry is invalid: bs={batch_size}, " + f"requests={len(requests)}, max={self.max_batch_size}" + ) + + slot = self._slots[self._slot_index] + self._slot_index = (self._slot_index + 1) % len(self._slots) + wait_started = time.perf_counter() + + current_stream = self._device_module.current_stream() + with self._device_module.stream(self._copy_stream): + # The copy stream is persistent, so this D2H is naturally ordered + # after the slot's preceding mask H2D. Waiting on the host here + # would serialize before target verify; the d2h_done event below + # instead lets any residual slot reuse overlap target compute. + self._copy_stream.wait_stream(current_stream) + slot.draft_tokens_cpu[:batch_size].copy_( + verify_ids_2d, + non_blocking=True, + ) + slot.d2h_done.record(self._copy_stream) + buffer_wait_seconds = time.perf_counter() - wait_started + + return GrammarMaskStep( + slot=slot, + batch_size=batch_size, + row_count=batch_size * chain_length, + requests=requests, + draft_input=draft_input, + buffer_wait_seconds=buffer_wait_seconds, + gpu_timing=self._reserve_gpu_timing(), + ) + + def mark_target_verify_enqueued(self, step: GrammarMaskStep) -> None: + slot = step.gpu_timing + if slot is None: + return + if not slot.reserved or slot.target_recorded or slot.dependency_recorded: + raise RuntimeError("DSpark grammar CUDA target timing state is invalid") + slot.target_done.record(self._device_module.current_stream()) + slot.target_recorded = True + + def _mark_mask_dependency( + self, + step: GrammarMaskStep, + *, + emit: bool, + ) -> None: + slot = step.gpu_timing + if slot is None: + return + if not slot.reserved or not slot.target_recorded or slot.dependency_recorded: + raise RuntimeError("DSpark grammar CUDA mask timing state is invalid") + slot.mask_dependency_done.record(self._device_module.current_stream()) + slot.dependency_recorded = True + slot.emit = emit + + @staticmethod + def _resolve_native_matcher( + grammar: BaseGrammarObject, + draft_tokens: Sequence[int], + ) -> tuple[GrammarMatcher, int] | None: + if isinstance(grammar, XGrammarGrammar): + return grammar.matcher, 0 + if not isinstance(grammar, ReasonerGrammarObject): + raise RuntimeError("DSpark grammar batching requires the XGrammar backend") + if grammar.enable_token_filter: + raise RuntimeError( + "DSpark grammar batching does not admit strict-thinking token filters" + ) + inner = grammar.grammar + if not isinstance(inner, XGrammarGrammar): + raise RuntimeError("DSpark reasoning grammar has no XGrammar matcher") + if grammar._is_generation(): + return inner.matcher, 0 + if not grammar._is_thinking(): + raise RuntimeError("reasoning grammar is in an invalid state") + + think_end_id = grammar.think_end_id + for position in range(1, len(draft_tokens)): + if int(draft_tokens[position]) == think_end_id: + return inner.matcher, position + return None + + def finish( + self, + step: GrammarMaskStep, + *, + grammar_barrier: Callable[[], dict[str, float]] | None, + ) -> GrammarMaskResult | None: + if grammar_barrier is None: + raise RuntimeError("DSpark grammar overlap requires the scheduler barrier") + if step.gpu_timing is not None and not step.gpu_timing.target_recorded: + raise RuntimeError("DSpark grammar target completion was not recorded") + + finish_started = time.perf_counter() + barrier_started = time.perf_counter() + barrier_phase_seconds = grammar_barrier() + barrier_seconds = time.perf_counter() - barrier_started + if not isinstance(barrier_phase_seconds, dict): + raise RuntimeError("DSpark grammar barrier must return phase timings") + for phase, seconds in barrier_phase_seconds.items(): + if ( + not isinstance(phase, str) + or not phase.startswith("barrier_") + or isinstance(seconds, bool) + or not isinstance(seconds, (int, float)) + or not math.isfinite(seconds) + or seconds < 0 + ): + raise RuntimeError( + "DSpark grammar barrier returned invalid timing: " + f"{phase!r}={seconds!r}" + ) + d2h_wait_started = time.perf_counter() + step.slot.d2h_done.synchronize() + d2h_wait_seconds = time.perf_counter() - d2h_wait_started + + cpu_mask = step.slot.vocab_mask_cpu[: step.row_count] + matcher_resolution_started = time.perf_counter() + matchers: list[GrammarMatcher] = [] + indices: list[int] = [] + root_positions: list[int] = [] + apply_grammar: XGrammarGrammar | ReasonerGrammarObject | None = None + draft_tokens = step.slot.draft_tokens_cpu[: step.batch_size] + # NumPy exposes the synchronized pinned tensor as a zero-copy CPU view. + # Its scalar indexing avoids one PyTorch dispatcher trip per request + # and draft position during reasoner transition detection. + draft_token_rows = draft_tokens.numpy() + + for index, request in enumerate(step.requests): + grammar = request.grammar + if grammar is None: + continue + if apply_grammar is None: + apply_grammar = grammar + resolved = self._resolve_native_matcher( + grammar, + draft_token_rows[index], + ) + if resolved is None: + continue + matcher, root_position = resolved + matchers.append(matcher) + indices.append(index) + root_positions.append(root_position) + matcher_resolution_seconds = time.perf_counter() - matcher_resolution_started + + if apply_grammar is None: + raise RuntimeError( + "batch.has_grammar was true but no request carried a grammar" + ) + if not matchers: + self._mark_mask_dependency(step, emit=False) + self._record( + batch_size=step.batch_size, + active_matchers=0, + outcome="thinking_only", + traversal_seconds=0.0, + d2h_wait_seconds=d2h_wait_seconds, + buffer_wait_seconds=step.buffer_wait_seconds, + phase_seconds={ + "buffer_wait": step.buffer_wait_seconds, + "grammar_barrier": barrier_seconds, + **barrier_phase_seconds, + "d2h_wait": d2h_wait_seconds, + "matcher_resolution": matcher_resolution_seconds, + "native_traversal": 0.0, + "h2d_enqueue": 0.0, + "finish_total": time.perf_counter() - finish_started, + }, + ) + return None + + traversal_started = time.perf_counter() + completed = self._batch_matcher.batch_traverse_draft_tree( + matchers, + self._next_token, + self._next_sibling, + draft_tokens, + cpu_mask, + indices=indices, + root_positions=root_positions, + ) + traversal_seconds = time.perf_counter() - traversal_started + if not all(completed): + raise RuntimeError("native DSpark grammar traversal timed out") + + h2d_enqueue_started = time.perf_counter() + current_stream = self._device_module.current_stream() + with self._device_module.stream(self._copy_stream): + if step.slot.has_mask_consumer: + self._copy_stream.wait_event(step.slot.mask_consumed) + step.slot.vocab_mask_device[: step.row_count].copy_( + cpu_mask, + non_blocking=True, + ) + step.slot.mask_ready.record(self._copy_stream) + current_stream.wait_event(step.slot.mask_ready) + self._mark_mask_dependency(step, emit=True) + h2d_enqueue_seconds = time.perf_counter() - h2d_enqueue_started + + step.draft_input.grammar = apply_grammar + self._record( + batch_size=step.batch_size, + active_matchers=len(matchers), + outcome="masked", + traversal_seconds=traversal_seconds, + d2h_wait_seconds=d2h_wait_seconds, + buffer_wait_seconds=step.buffer_wait_seconds, + phase_seconds={ + "buffer_wait": step.buffer_wait_seconds, + "grammar_barrier": barrier_seconds, + **barrier_phase_seconds, + "d2h_wait": d2h_wait_seconds, + "matcher_resolution": matcher_resolution_seconds, + "native_traversal": traversal_seconds, + "h2d_enqueue": h2d_enqueue_seconds, + "finish_total": time.perf_counter() - finish_started, + }, + ) + return GrammarMaskResult( + step=step, + grammar=apply_grammar, + vocab_mask=step.slot.vocab_mask_device[: step.row_count], + active_matchers=len(matchers), + traversal_seconds=traversal_seconds, + d2h_wait_seconds=d2h_wait_seconds, + ) + + def mark_consumed(self, result: GrammarMaskResult) -> None: + result.step.slot.mask_consumed.record(self._device_module.current_stream()) + result.step.slot.has_mask_consumer = True + + def _record( + self, + *, + batch_size: int, + active_matchers: int, + outcome: str, + traversal_seconds: float, + d2h_wait_seconds: float, + buffer_wait_seconds: float, + phase_seconds: dict[str, float], + ) -> None: + self._steps += 1 + self._active_matchers += active_matchers + self._traversal_seconds += traversal_seconds + self._traversal_max_seconds = max( + self._traversal_max_seconds, + traversal_seconds, + ) + self._d2h_wait_seconds += d2h_wait_seconds + self._buffer_wait_seconds += buffer_wait_seconds + self._pending_outcomes[outcome] += 1 + if ( + self._metrics_collector is not None + and self._steps % self._metrics_interval == 0 + ): + self._metrics_collector.observe_dspark_grammar_step( + batch_size=batch_size, + active_matchers=active_matchers, + outcome_counts=self._pending_outcomes.copy(), + phase_seconds=phase_seconds, + ) + self._pending_outcomes = { + "masked": 0, + "thinking_only": 0, + } + if self._steps % 1000 == 0: + logger.info( + "DSpark grammar stats: steps=%d active_per_step=%.2f " + "traversal_avg_ms=%.3f traversal_max_ms=%.3f " + "d2h_wait_avg_ms=%.3f buffer_wait_avg_ms=%.3f", + self._steps, + self._active_matchers / self._steps, + self._traversal_seconds / self._steps * 1000, + self._traversal_max_seconds * 1000, + self._d2h_wait_seconds / self._steps * 1000, + self._buffer_wait_seconds / self._steps * 1000, + ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py index 3fbd56fe2b10..07ff9bc08c62 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_planner.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_planner.py @@ -20,6 +20,9 @@ ResolvedConfidence, ) from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.managers.scheduler_components.single_node_dp2_sync import ( + exchange_single_node_dp2_verify_tier, +) from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 @@ -331,14 +334,13 @@ def _maybe_gather_dp_verify_tier( batch.global_spec_verify_tier_num_tokens = None return cpu_group = get_tp_group().cpu_group - local_tensor = torch.tensor([local_tier_num_tokens], dtype=torch.int64) - gathered = torch.empty( - (torch.distributed.get_world_size(group=cpu_group),), dtype=torch.int64 - ) - torch.distributed.all_gather_into_tensor( - gathered, local_tensor, group=cpu_group + batch.global_spec_verify_tier_num_tokens = exchange_single_node_dp2_verify_tier( + local_tier_num_tokens, + group=cpu_group, + dp_size=2, + tp_size=get_parallel().attn_tp_size, + cp_size=get_parallel().attn_cp_size, ) - batch.global_spec_verify_tier_num_tokens = gathered.tolist() def note_non_decode_step(self) -> None: if self._budget_planner is not None: @@ -429,11 +431,20 @@ def schedule_layout( ) -> Optional[RaggedVerifyLayout]: if self._ragged_verify_mode is RaggedVerifyMode.STATIC: return None - if self._is_verify_all and self._ragged_verify_mode is RaggedVerifyMode.COMPACT: - # Verify-all: the uniform layout (or None, past the captured grid) - # is constant per (bs, tier); serve it from cache instead of paying - # the per-step schedule and its host<->device round-trips. - key = (int(req_pool_indices.shape[0]), global_num_reqs) + bs = int(req_pool_indices.shape[0]) + full_budget = bs * ( + self.verify_num_draft_tokens - max(self._schedule_cfg.min_verify_len, 1) + ) + if self._ragged_verify_mode is RaggedVerifyMode.COMPACT and ( + self._is_verify_all + or (bs > 0 and budget is not None and budget >= full_budget) + ): + # A locally full budget does not imply the peer is full. Keep the + # fast uniform local layout, but key its graph padding from the + # already-agreed DP token tier. Otherwise one rank can select the + # full global-width graph while its peer selects the smaller shared + # tier, deadlocking their in-graph collectives. + key = (bs, global_num_reqs, dp_tier_num_tokens) if key not in self._uniform_layout_cache: self._uniform_layout_cache[key] = uniform_ragged_layout( bs=key[0], @@ -442,6 +453,7 @@ def schedule_layout( ragged_verify_mode=self._ragged_verify_mode, model_runner=self.model_runner, tier_num_reqs=global_num_reqs, + tier_num_tokens=dp_tier_num_tokens, ) return self._uniform_layout_cache[key] verify_lens = self._schedule_verify_lens( @@ -739,12 +751,14 @@ def uniform_ragged_layout( ragged_verify_mode: RaggedVerifyMode, model_runner, tier_num_reqs: Optional[int] = None, + tier_num_tokens: Optional[int] = None, ) -> Optional[RaggedVerifyLayout]: tier_num_reqs = bs if tier_num_reqs is None else tier_num_reqs if ragged_layout_exceeds_captured_grid( num_reqs=tier_num_reqs, verify_num_draft_tokens=verify_num_draft_tokens, model_runner=model_runner, + tier_tokens_hint=tier_num_tokens, ): return None verify_lens_cpu = [verify_num_draft_tokens] * bs @@ -758,6 +772,7 @@ def uniform_ragged_layout( ragged_verify_mode=ragged_verify_mode, verify_num_draft_tokens=verify_num_draft_tokens, model_runner=model_runner, + tier_num_tokens=tier_num_tokens, ) return RaggedVerifyLayout.from_verify_lens( verify_lens_cpu=verify_lens_cpu, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 4d5478cd53c7..2e2ee52cc76d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -202,9 +202,26 @@ def run_idle_participation( batch.seq_lens_cpu = torch.ones((num_dummy_slots,), dtype=torch.int64) batch.seq_lens_sum = num_dummy_slots batch.forward_mode = ForwardMode.TARGET_VERIFY - verify_forward_batch, _ = verify_input.prepare_for_verify( + verify_forward_batch, can_run_cuda_graph = verify_input.prepare_for_verify( batch, self.target_worker ) + if not can_run_cuda_graph: + graph_runner = self.target_worker.model_runner.decode_cuda_graph_runner + graph_diagnostics = ( + f"batch_size={verify_forward_batch.batch_size}, " + f"batch_can_dp_graph={verify_forward_batch.can_run_dp_cuda_graph}, " + f"global_tokens={verify_forward_batch.global_num_tokens_cpu}, " + f"runner_hidden={getattr(graph_runner, 'capture_hidden_mode', None)}, " + f"runner_max_bs={getattr(graph_runner, 'max_bs', None)}, " + f"runner_tokens={getattr(graph_runner, 'capture_num_tokens', None)}, " + f"attn_ragged={getattr(graph_runner.attn_backend, 'supports_ragged_verify_graph', None)}" + ) + raise RuntimeError( + "DSpark DP idle target-verify participation missed the " + "required ragged CUDA graph; an eager idle verify cannot " + "safely share collectives with an active graph " + f"(layout_tokens={num_dummy_tokens}; {graph_diagnostics})." + ) self.target_worker.forward_batch_generation( batch=None, forward_batch=verify_forward_batch, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index ae778ee1dbff..de4eb1889730 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -1,6 +1,9 @@ import logging +import os +import time +from concurrent.futures import Future, ThreadPoolExecutor from contextlib import nullcontext -from dataclasses import replace +from dataclasses import dataclass, replace from typing import Optional import torch @@ -10,6 +13,7 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker +from sglang.srt.managers.utils import PreparedGrammarResult from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, compute_position, @@ -34,6 +38,9 @@ make_next_draft_input, maybe_build_draft_sampler, ) +from sglang.srt.speculative.dspark_components.dspark_grammar_pipeline import ( + DSparkGrammarPipeline, +) from sglang.srt.speculative.dspark_components.dspark_kv_inject import ( TargetHiddenKvInjector, ) @@ -53,16 +60,20 @@ TargetVerifyExecutor, verify_logits_adjustments_are_noop, ) -from sglang.srt.speculative.spec_utils import ( - GrammarTree, - build_grammar_vocab_mask, - draft_tp_context, -) +from sglang.srt.speculative.spec_utils import draft_tp_context from sglang.srt.utils import get_available_gpu_memory, is_cuda logger = logging.getLogger(__name__) +@dataclass(frozen=True, slots=True) +class GrammarResultCopy: + next_token_ids: torch.Tensor + accept_lens: Optional[torch.Tensor] + done: torch.cuda.Event + future: Future[PreparedGrammarResult] + + class DSparkWorkerV2(BaseSpecWorker): def __init__( @@ -72,6 +83,7 @@ def __init__( ps: ParallelState, nccl_port: int, target_worker: TpModelWorker, + metrics_collector=None, ): self.server_args = server_args self.gpu_id = gpu_id @@ -254,6 +266,97 @@ def __init__( device=self.device, simulate_acc_len=self._simulate_acc_len, ) + self._grammar_pipeline = DSparkGrammarPipeline( + device=self.device, + max_batch_size=max(server_args.cuda_graph_config.decode.bs), + chain_length=self.verify_num_draft_tokens, + vocab_size=int(self.target_worker.model_runner.model_config.vocab_size), + traversal_threads=int( + os.environ.get("SGLANG_GRAMMAR_TRAVERSAL_THREADS", "12") + ), + metrics_collector=metrics_collector, + ) + self._grammar_result_copy_stream = torch.get_device_module(self.device).Stream() + self._grammar_result_executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix=f"dspark-grammar-result-dp{self.ps.dp_rank}", + ) + + @staticmethod + def _prepare_grammar_result( + *, + next_token_ids: torch.Tensor, + accept_lens: Optional[torch.Tensor], + done: torch.cuda.Event, + submitted_at: float, + ) -> PreparedGrammarResult: + """Wait and materialize the early result without blocking the scheduler.""" + worker_started = time.perf_counter() + done.synchronize() + copy_ready = time.perf_counter() + next_token_ids_list = next_token_ids.tolist() + accept_lens_list = accept_lens.tolist() if accept_lens is not None else None + prepared_at = time.perf_counter() + return PreparedGrammarResult( + next_token_ids=next_token_ids_list, + accept_lens=accept_lens_list, + queue_delay_seconds=worker_started - submitted_at, + copy_wait_seconds=copy_ready - worker_started, + tensor_to_list_seconds=prepared_at - copy_ready, + submit_to_ready_seconds=prepared_at - submitted_at, + ) + + def _start_grammar_result_copy( + self, + *, + next_token_ids: torch.Tensor, + accept_lens: Optional[torch.Tensor], + ) -> GrammarResultCopy: + """Publish the minimal next-step grammar dependency before forward end.""" + if not next_token_ids.is_cuda: + raise RuntimeError("DSpark early grammar result must be CUDA-resident") + if accept_lens is not None and not accept_lens.is_cuda: + raise RuntimeError( + "DSpark early grammar accept lengths must be CUDA-resident" + ) + + device_module = torch.get_device_module(self.device) + current_stream = device_module.current_stream() + copy_stream = self._grammar_result_copy_stream + copy_stream.wait_stream(current_stream) + with device_module.stream(copy_stream): + next_token_ids_cpu = torch.empty( + next_token_ids.shape, + dtype=next_token_ids.dtype, + pin_memory=True, + ) + next_token_ids_cpu.copy_(next_token_ids, non_blocking=True) + next_token_ids.record_stream(copy_stream) + accept_lens_cpu = None + if accept_lens is not None: + accept_lens_cpu = torch.empty( + accept_lens.shape, + dtype=accept_lens.dtype, + pin_memory=True, + ) + accept_lens_cpu.copy_(accept_lens, non_blocking=True) + accept_lens.record_stream(copy_stream) + done = device_module.Event() + done.record(copy_stream) + submitted_at = time.perf_counter() + future = self._grammar_result_executor.submit( + self._prepare_grammar_result, + next_token_ids=next_token_ids_cpu, + accept_lens=accept_lens_cpu, + done=done, + submitted_at=submitted_at, + ) + return GrammarResultCopy( + next_token_ids_cpu, + accept_lens_cpu, + done, + future, + ) def _resolve_target_embed_tokens(self, target_model): if hasattr(target_model, "get_input_embeddings"): @@ -364,9 +467,27 @@ def forward_batch_generation( on_publish=None, grammar_barrier=None, ) -> GenerationBatchResult: - if getattr(batch, "return_logprob", False): + # Pure prefill already runs the target model, including its eager + # logits tail, and returns that LogitsProcessorOutput unchanged. It + # therefore supports prompt/output logprobs without involving the + # speculative verifier. Decode and mixed batches still fail closed: + # DSpark does not yet reconstruct their accepted-token logprobs. + # + # DP attention can represent a local decode or idle peer of a global + # prefill with is_extend_in_batch set. The dispatch immediately below + # routes exactly those batches through _forward_prefill so every rank + # participates in the target-model collectives. The request validator + # limits DSpark logprobs to one output token, so a local decode form is + # only the overlap scheduler's already-issued retirement step; it + # cannot expose speculative accepted-token logprobs. + supports_prefill_logprob = ( + batch.forward_mode.is_extend() or batch.is_extend_in_batch + ) + if getattr(batch, "return_logprob", False) and not supports_prefill_logprob: raise ValueError( - "DSpark speculative decoding does not support return_logprob yet." + "DSpark speculative decode does not support return_logprob yet: " + f"forward_mode={batch.forward_mode.name}, " + f"is_extend_in_batch={batch.is_extend_in_batch}." ) if batch.forward_mode.is_extend() or batch.is_extend_in_batch: @@ -392,6 +513,15 @@ def _forward_prefill( logits_output = batch_output.logits_output next_token_ids = batch_output.next_token_ids batch_output.new_seq_lens = batch.seq_lens + if batch.has_grammar: + grammar_copy = self._start_grammar_result_copy( + next_token_ids=next_token_ids, + accept_lens=None, + ) + batch_output.grammar_next_token_ids = grammar_copy.next_token_ids + batch_output.grammar_accept_lens = grammar_copy.accept_lens + batch_output.grammar_copy_done = grammar_copy.done + batch_output.grammar_result_future = grammar_copy.future if on_publish is not None: on_publish(batch_output.new_seq_lens) @@ -558,6 +688,7 @@ def _forward_decode( and batch.global_num_tokens is not None else None ) + dp_tier_num_tokens = self._dp_verify_tier_num_tokens(batch) layout = self._verify_planner.schedule_layout( req_pool_indices=batch.req_pool_indices, prefix_lens=prefix_lens, @@ -565,21 +696,35 @@ def _forward_decode( confidence=confidence, budget=verify_token_budget, global_num_reqs=global_num_reqs, - dp_tier_num_tokens=self._dp_verify_tier_num_tokens(batch), + dp_tier_num_tokens=dp_tier_num_tokens, ) + if envs.SGLANG_LOG_DECODE_GRAPH_KEY.get(): + logger.info( + "DSpark verify graph tier: iter=%d local_bs=%d global_bs=%s " + "local_tier=%d global_tiers=%s dp_tier=%s graph_tokens=%s", + int(batch.forward_iter), + bs, + global_num_reqs, + int(batch.spec_verify_tier_num_tokens), + batch.global_spec_verify_tier_num_tokens, + dp_tier_num_tokens, + None if layout is None else layout.graph_num_tokens, + ) run_compact = self._verify_planner.should_run_compact(layout=layout) verify_ids_2d = torch.cat( [draft_block_ids[:, :1], draft_tokens], dim=1 ).contiguous() - # Must stay ahead of the target verify launch below. - grammar_tree = ( - GrammarTree.from_linear_chain(verify_ids_2d) if batch.has_grammar else None + grammar_step = ( + self._grammar_pipeline.begin( + verify_ids_2d=verify_ids_2d, + requests=batch.reqs, + draft_input=draft_input, + ) + if batch.has_grammar + else None ) - - # A live grammar forces the eager path: the folded epilogue accepts inside - # the cuda graph off its own buffers, where the mask below never lands. fold_eligible = ( self._verify_executor.verify_epilogue is not None and proposal.folded @@ -611,18 +756,18 @@ def _forward_decode( logits_output = target_verify.logits_output can_run_cuda_graph = target_verify.can_run_cuda_graph - if batch.has_grammar: - # run_compact scatters its rows back to (bs * chain_len), so the mask - # lines up with the logits on both verify paths. - grammar_mask = build_grammar_vocab_mask( - reqs=batch.reqs, - tree=grammar_tree, - sampling_info=sampling_info, - device=logits_output.next_token_logits.device, - barrier=grammar_barrier, + if grammar_step is not None: + self._grammar_pipeline.mark_target_verify_enqueued(grammar_step) + grammar_result = self._grammar_pipeline.finish( + grammar_step, + grammar_barrier=grammar_barrier, ) - if grammar_mask is not None: - grammar_mask.apply(logits_output.next_token_logits) + if grammar_result is not None: + grammar_result.grammar.apply_vocab_mask( + logits=logits_output.next_token_logits, + vocab_mask=grammar_result.vocab_mask, + ) + self._grammar_pipeline.mark_consumed(grammar_result) epilogue = self._verify_executor.verify_epilogue folded_accept = fold_eligible and run_compact and can_run_cuda_graph @@ -638,6 +783,14 @@ def _forward_decode( prefix_lens=prefix_lens, draft_tokens=draft_tokens, ) + grammar_copy = ( + self._start_grammar_result_copy( + next_token_ids=accept.out_tokens.reshape(-1), + accept_lens=accept.commit_lens, + ) + if batch.has_grammar + else None + ) if on_publish is not None: if confidence is not None: on_publish(accept.new_seq_lens, confidence=confidence) @@ -697,6 +850,16 @@ def _forward_decode( next_draft_input=next_draft_input, speculative_num_draft_tokens=int(self.verify_num_draft_tokens), new_seq_lens=accept.new_seq_lens, + grammar_next_token_ids=( + grammar_copy.next_token_ids if grammar_copy is not None else None + ), + grammar_accept_lens=( + grammar_copy.accept_lens if grammar_copy is not None else None + ), + grammar_copy_done=(grammar_copy.done if grammar_copy is not None else None), + grammar_result_future=( + grammar_copy.future if grammar_copy is not None else None + ), ) def get_confidence_budget_prepare(self): diff --git a/python/sglang/srt/utils/watchdog.py b/python/sglang/srt/utils/watchdog.py index 7c774e758bb5..b3d01f1aa787 100644 --- a/python/sglang/srt/utils/watchdog.py +++ b/python/sglang/srt/utils/watchdog.py @@ -200,6 +200,32 @@ def stop(self) -> None: self._thread.join(timeout=self._interval * 2) self._thread = None + def wait_for_clean_exit(self, timeout: float) -> None: + """Wait for every tracked child to exit successfully. + + The watchdog must be stopped before expected shutdown exits occur, + otherwise it correctly interprets the first child exit as a crash. + All children run concurrently, so they share one absolute deadline. + """ + if timeout <= 0: + raise ValueError(f"timeout must be positive, got {timeout}") + + self.stop() + deadline = time.monotonic() + timeout + failures = [] + for proc, name in zip(self._processes, self._names): + remaining = max(0.0, deadline - time.monotonic()) + proc.join(remaining) + if proc.is_alive(): + failures.append(f"{name}(pid={proc.pid}, timed out)") + elif proc.exitcode != 0: + failures.append(f"{name}(pid={proc.pid}, exitcode={proc.exitcode})") + + if failures: + raise RuntimeError( + "SGLang subprocess shutdown was not clean: " + ", ".join(failures) + ) + def _monitor_loop(self) -> None: try: while not self._stop_event.wait(self._interval): diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index d76c00880368..e0756f8ec8c7 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -1,5 +1,5 @@ cmake_minimum_required(VERSION 3.26 FATAL_ERROR) -project(sgl-kernel LANGUAGES CXX CUDA) +project(sgl-kernel LANGUAGES C CXX CUDA) # utils include(${CMAKE_CURRENT_LIST_DIR}/cmake/utils.cmake) @@ -23,6 +23,18 @@ find_package(Python COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COM set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3") +# Low-latency, fail-closed metadata exchange for co-located DP2 schedulers. +# This has no CUDA or PyTorch dependency and is loaded by SGLang through +# ctypes only when SGLANG_DSPARK_DP2_SHM_MLP_SYNC=1. +add_library(sglang_dp2_sync SHARED csrc/cpu/dp2_sync.c) +target_compile_features(sglang_dp2_sync PRIVATE c_std_11) +target_compile_options( + sglang_dp2_sync + PRIVATE -O3 -Wall -Wextra -Werror +) +target_link_libraries(sglang_dp2_sync PRIVATE rt) +install(TARGETS sglang_dp2_sync LIBRARY DESTINATION "sgl_kernel") + # CUDA enable_language(CUDA) find_package(CUDAToolkit REQUIRED) diff --git a/sgl-kernel/cmake/flashmla-sparse-prefill-out.patch b/sgl-kernel/cmake/flashmla-sparse-prefill-out.patch new file mode 100644 index 000000000000..c0748ee57665 --- /dev/null +++ b/sgl-kernel/cmake/flashmla-sparse-prefill-out.patch @@ -0,0 +1,39 @@ +diff --git a/csrc/api/sparse_fwd.h b/csrc/api/sparse_fwd.h +index 66d7111..9f7e8e9 100644 +--- a/csrc/api/sparse_fwd.h ++++ b/csrc/api/sparse_fwd.h +@@ -105,5 +105,8 @@ + float sm_scale, + int d_v, + const std::optional &attn_sink, +- const std::optional &topk_length ++ const std::optional &topk_length, ++ const std::optional &out_buffer = std::nullopt, ++ const std::optional &max_logits_buffer = std::nullopt, ++ const std::optional &lse_buffer = std::nullopt + ) { +@@ -162,4 +165,19 @@ +- at::Tensor out = torch::empty({s_q, h_q, d_v}, opts); +- at::Tensor lse = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); +- at::Tensor max_logits = torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); ++ at::Tensor out = out_buffer.has_value() ++ ? out_buffer.value() ++ : torch::empty({s_q, h_q, d_v}, opts); ++ at::Tensor lse = lse_buffer.has_value() ++ ? lse_buffer.value() ++ : torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); ++ at::Tensor max_logits = max_logits_buffer.has_value() ++ ? max_logits_buffer.value() ++ : torch::empty({s_q, h_q}, opts.dtype(torch::kFloat)); ++ KU_CHECK_DEVICE(out); ++ KU_CHECK_DEVICE(lse); ++ KU_CHECK_DEVICE(max_logits); ++ KU_CHECK_DTYPE(out, torch::kBFloat16); ++ KU_CHECK_DTYPE(lse, torch::kFloat32); ++ KU_CHECK_DTYPE(max_logits, torch::kFloat32); ++ KU_CHECK_SHAPE(out, s_q, h_q, d_v); ++ KU_CHECK_SHAPE(lse, s_q, h_q); ++ KU_CHECK_SHAPE(max_logits, s_q, h_q); + KU_CHECK_CONTIGUOUS(out); + KU_CHECK_CONTIGUOUS(lse); + KU_CHECK_CONTIGUOUS(max_logits); diff --git a/sgl-kernel/cmake/flashmla.cmake b/sgl-kernel/cmake/flashmla.cmake index 7387f3e5d21b..58fd139fbf1d 100644 --- a/sgl-kernel/cmake/flashmla.cmake +++ b/sgl-kernel/cmake/flashmla.cmake @@ -7,6 +7,52 @@ FetchContent_Declare( ) FetchContent_Populate(repo-flashmla) +# The DSV4 sparse-prefill path owns persistent output buffers. Extend the +# pinned FlashMLA interface so its kernel can write directly into them instead +# of allocating a ~1 GiB output at runtime. Exact pre/post hashes make +# reconfiguration idempotent while every unknown source state fails closed. +set(FLASHMLA_SPARSE_PREFILL_HEADER + "${repo-flashmla_SOURCE_DIR}/csrc/api/sparse_fwd.h") +set(FLASHMLA_SPARSE_PREFILL_HEADER_UPSTREAM_SHA256 + "87800f0f73f055a1bcbe13a9a0c1e388e36b106e0405c06afb5fa86d11221cc8") +set(FLASHMLA_SPARSE_PREFILL_HEADER_PATCHED_SHA256 + "006aa754f22360e02547c967788e5a43d5beb092f626b0f5a07a1766e0ef71a8") +file(SHA256 "${FLASHMLA_SPARSE_PREFILL_HEADER}" + FLASHMLA_SPARSE_PREFILL_HEADER_SHA256) +if(FLASHMLA_SPARSE_PREFILL_HEADER_SHA256 STREQUAL + FLASHMLA_SPARSE_PREFILL_HEADER_UPSTREAM_SHA256) + find_program(FLASHMLA_PATCH_EXECUTABLE patch REQUIRED) + execute_process( + COMMAND + "${FLASHMLA_PATCH_EXECUTABLE}" + --batch + --forward + -p1 + --input=${CMAKE_CURRENT_LIST_DIR}/flashmla-sparse-prefill-out.patch + WORKING_DIRECTORY "${repo-flashmla_SOURCE_DIR}" + RESULT_VARIABLE FLASHMLA_SPARSE_PREFILL_PATCH_RESULT + OUTPUT_VARIABLE FLASHMLA_SPARSE_PREFILL_PATCH_STDOUT + ERROR_VARIABLE FLASHMLA_SPARSE_PREFILL_PATCH_STDERR + ) + if(NOT FLASHMLA_SPARSE_PREFILL_PATCH_RESULT EQUAL 0) + message(FATAL_ERROR + "Failed to apply the pinned FlashMLA sparse-prefill output patch:\n" + "${FLASHMLA_SPARSE_PREFILL_PATCH_STDOUT}\n" + "${FLASHMLA_SPARSE_PREFILL_PATCH_STDERR}" + ) + endif() + file(SHA256 "${FLASHMLA_SPARSE_PREFILL_HEADER}" + FLASHMLA_SPARSE_PREFILL_HEADER_SHA256) +endif() +if(NOT FLASHMLA_SPARSE_PREFILL_HEADER_SHA256 STREQUAL + FLASHMLA_SPARSE_PREFILL_HEADER_PATCHED_SHA256) + message(FATAL_ERROR + "FlashMLA sparse-prefill header is neither the exact pinned source " + "nor the exact optimized source: " + "${FLASHMLA_SPARSE_PREFILL_HEADER_SHA256}" + ) +endif() + # flashmla submodule pin: NVIDIA/cutlass @ 147f5673d0c1c3dcf66f78d677fd647e4a020219 FetchContent_Declare( repo-flashmla-cutlass diff --git a/sgl-kernel/csrc/cpu/dp2_sync.c b/sgl-kernel/csrc/cpu/dp2_sync.c new file mode 100644 index 000000000000..267299cd990e --- /dev/null +++ b/sgl-kernel/csrc/cpu/dp2_sync.c @@ -0,0 +1,827 @@ +#define _GNU_SOURCE + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define SGLANG_DP2_SYNC_ABI_VERSION 2U +#define SGLANG_DP2_SYNC_LAYOUT_VERSION 1U +#define SGLANG_DP2_SYNC_MAGIC UINT64_C(0x53484450324d4c50) +#define SGLANG_DP2_SYNC_WORLD_SIZE 2U +#define SGLANG_DP2_SYNC_WIDTH 7U +#define SGLANG_DP2_SYNC_SLOTS 2U +#define SGLANG_DP2_SYNC_ERROR_BUFFER_MIN 128U +#define SGLANG_DP2_SYNC_WAIT_SLICE_NS UINT64_C(100000000) + +_Static_assert(ATOMIC_INT_LOCK_FREE == 2, "32-bit atomics must be lock-free"); +_Static_assert( + ATOMIC_LLONG_LOCK_FREE == 2, "64-bit atomics must be lock-free" +); + +struct sglang_dp2_sync_stats { + uint64_t sequence; + uint64_t total_ns; + uint64_t peer_wait_ns; + uint64_t arrival_skew_ns; + uint64_t post_latest_arrival_ns; +}; + +struct sglang_dp2_rank_state { + alignas(64) _Atomic uint64_t published_sequence; + _Atomic uint32_t futex_epoch; + _Atomic int32_t pid; + uint32_t reserved; + uint64_t arrival_ns[SGLANG_DP2_SYNC_SLOTS]; + int64_t + payload[SGLANG_DP2_SYNC_SLOTS][SGLANG_DP2_SYNC_WIDTH]; +}; + +struct sglang_dp2_shared_state { + alignas(64) uint64_t magic; + uint32_t layout_version; + uint32_t world_size; + uint32_t width; + uint32_t state_size; + _Atomic uint32_t error_code; + _Atomic uint32_t unlinked; + uint8_t header_padding[32]; + struct sglang_dp2_rank_state rank[SGLANG_DP2_SYNC_WORLD_SIZE]; +}; + +struct sglang_dp2_sync_handle { + struct sglang_dp2_shared_state *shared; + size_t mapping_size; + uint64_t timeout_ns; + uint64_t local_sequence; + int rank; + int fd; + pid_t pid; + char shm_name[64]; +}; + +_Static_assert( + offsetof(struct sglang_dp2_shared_state, rank) % 64 == 0, + "rank state must start on a cache-line boundary" +); +_Static_assert( + sizeof(struct sglang_dp2_rank_state) % 64 == 0, + "rank states must not share cache lines" +); + +static void set_error( + char *error_buffer, + size_t error_buffer_size, + const char *format, + ... +) { + if (error_buffer == NULL || error_buffer_size == 0) { + return; + } + va_list args; + va_start(args, format); + (void)vsnprintf(error_buffer, error_buffer_size, format, args); + va_end(args); + error_buffer[error_buffer_size - 1] = '\0'; +} + +static uint64_t monotonic_ns(void) { + struct timespec timestamp; + if (clock_gettime(CLOCK_MONOTONIC_RAW, ×tamp) != 0) { + return 0; + } + return (uint64_t)timestamp.tv_sec * UINT64_C(1000000000) + + (uint64_t)timestamp.tv_nsec; +} + +static uint64_t fnv1a_64(const char *value) { + uint64_t hash = UINT64_C(14695981039346656037); + const unsigned char *cursor = (const unsigned char *)value; + while (*cursor != '\0') { + hash ^= (uint64_t)*cursor++; + hash *= UINT64_C(1099511628211); + } + return hash; +} + +static int futex_wait_shared( + _Atomic uint32_t *address, + uint32_t expected, + uint64_t timeout_ns +) { + struct timespec timeout = { + .tv_sec = (time_t)(timeout_ns / UINT64_C(1000000000)), + .tv_nsec = (long)(timeout_ns % UINT64_C(1000000000)), + }; + return (int)syscall( + SYS_futex, + (uint32_t *)address, + FUTEX_WAIT, + expected, + &timeout, + NULL, + 0 + ); +} + +static void futex_wake_all(_Atomic uint32_t *address) { + (void)syscall( + SYS_futex, + (uint32_t *)address, + FUTEX_WAKE, + INT32_MAX, + NULL, + NULL, + 0 + ); +} + +static void wake_rank(struct sglang_dp2_rank_state *rank) { + (void)atomic_fetch_add_explicit( + &rank->futex_epoch, 1U, memory_order_release + ); + futex_wake_all(&rank->futex_epoch); +} + +static void publish_shared_error( + struct sglang_dp2_sync_handle *handle, + uint32_t error_code +) { + uint32_t expected = 0; + (void)atomic_compare_exchange_strong_explicit( + &handle->shared->error_code, + &expected, + error_code, + memory_order_acq_rel, + memory_order_acquire + ); + for (size_t rank = 0; rank < SGLANG_DP2_SYNC_WORLD_SIZE; ++rank) { + wake_rank(&handle->shared->rank[rank]); + } +} + +static int validate_shared_state( + const struct sglang_dp2_shared_state *shared, + char *error_buffer, + size_t error_buffer_size +) { + if (shared->magic != SGLANG_DP2_SYNC_MAGIC) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync has invalid magic 0x%016llx", + (unsigned long long)shared->magic + ); + return -1; + } + if (shared->layout_version != SGLANG_DP2_SYNC_LAYOUT_VERSION) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync layout mismatch: expected %u, got %u", + SGLANG_DP2_SYNC_LAYOUT_VERSION, + shared->layout_version + ); + return -1; + } + if ( + shared->world_size != SGLANG_DP2_SYNC_WORLD_SIZE || + shared->width != SGLANG_DP2_SYNC_WIDTH || + shared->state_size != sizeof(*shared) + ) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync geometry mismatch: world=%u width=%u " + "state_size=%u expected=%zu", + shared->world_size, + shared->width, + shared->state_size, + sizeof(*shared) + ); + return -1; + } + return 0; +} + +uint32_t sglang_dp2_sync_abi_version(void) { + return SGLANG_DP2_SYNC_ABI_VERSION; +} + +int sglang_dp2_sync_open( + const char *session_id, + int rank, + uint64_t timeout_ns, + void **output_handle, + char *error_buffer, + size_t error_buffer_size +) { + if (error_buffer != NULL && error_buffer_size > 0) { + error_buffer[0] = '\0'; + } + if ( + session_id == NULL || session_id[0] == '\0' || + output_handle == NULL || + error_buffer == NULL || + error_buffer_size < SGLANG_DP2_SYNC_ERROR_BUFFER_MIN + ) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync received invalid open arguments" + ); + return -1; + } + if (rank < 0 || rank >= (int)SGLANG_DP2_SYNC_WORLD_SIZE) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync rank must be 0 or 1, got %d", + rank + ); + return -1; + } + if (timeout_ns == 0) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync timeout must be positive" + ); + return -1; + } + + struct sglang_dp2_sync_handle *handle = + calloc(1, sizeof(struct sglang_dp2_sync_handle)); + if (handle == NULL) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync handle allocation failed" + ); + return -1; + } + handle->fd = -1; + handle->mapping_size = sizeof(struct sglang_dp2_shared_state); + handle->timeout_ns = timeout_ns; + handle->rank = rank; + handle->pid = getpid(); + (void)snprintf( + handle->shm_name, + sizeof(handle->shm_name), + "/sglang_dp2_%016llx", + (unsigned long long)fnv1a_64(session_id) + ); + + handle->fd = shm_open( + handle->shm_name, + O_RDWR | O_CREAT | O_CLOEXEC, + S_IRUSR | S_IWUSR + ); + if (handle->fd < 0) { + set_error( + error_buffer, + error_buffer_size, + "shm_open(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + free(handle); + return -1; + } + if (flock(handle->fd, LOCK_EX) != 0) { + set_error( + error_buffer, + error_buffer_size, + "flock(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + close(handle->fd); + free(handle); + return -1; + } + + struct stat file_stat; + if (fstat(handle->fd, &file_stat) != 0) { + set_error( + error_buffer, + error_buffer_size, + "fstat(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + free(handle); + return -1; + } + + bool initialize = file_stat.st_size == 0; + if ( + !initialize && + (uint64_t)file_stat.st_size != (uint64_t)handle->mapping_size + ) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync segment has size %lld, expected %zu", + (long long)file_stat.st_size, + handle->mapping_size + ); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + free(handle); + return -1; + } + if (initialize && ftruncate(handle->fd, (off_t)handle->mapping_size) != 0) { + set_error( + error_buffer, + error_buffer_size, + "ftruncate(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + (void)shm_unlink(handle->shm_name); + free(handle); + return -1; + } + + void *mapping = mmap( + NULL, + handle->mapping_size, + PROT_READ | PROT_WRITE, + MAP_SHARED, + handle->fd, + 0 + ); + if (mapping == MAP_FAILED) { + set_error( + error_buffer, + error_buffer_size, + "mmap(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + if (initialize) { + (void)shm_unlink(handle->shm_name); + } + free(handle); + return -1; + } + handle->shared = mapping; + + if (initialize) { + memset(handle->shared, 0, handle->mapping_size); + handle->shared->layout_version = SGLANG_DP2_SYNC_LAYOUT_VERSION; + handle->shared->world_size = SGLANG_DP2_SYNC_WORLD_SIZE; + handle->shared->width = SGLANG_DP2_SYNC_WIDTH; + handle->shared->state_size = (uint32_t)sizeof(*handle->shared); + atomic_thread_fence(memory_order_release); + handle->shared->magic = SGLANG_DP2_SYNC_MAGIC; + } else if ( + validate_shared_state( + handle->shared, error_buffer, error_buffer_size + ) != 0 + ) { + (void)munmap(handle->shared, handle->mapping_size); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + free(handle); + return -1; + } + + if (mlock(handle->shared, handle->mapping_size) != 0) { + set_error( + error_buffer, + error_buffer_size, + "mlock(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + (void)munmap(handle->shared, handle->mapping_size); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + if (initialize) { + (void)shm_unlink(handle->shm_name); + } + free(handle); + return -1; + } + + _Atomic int32_t *pid_slot = &handle->shared->rank[rank].pid; + int32_t expected_pid = 0; + if (!atomic_compare_exchange_strong_explicit( + pid_slot, + &expected_pid, + (int32_t)handle->pid, + memory_order_acq_rel, + memory_order_acquire + )) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync rank %d is already owned by pid %d", + rank, + expected_pid + ); + (void)munlock(handle->shared, handle->mapping_size); + (void)munmap(handle->shared, handle->mapping_size); + (void)flock(handle->fd, LOCK_UN); + close(handle->fd); + free(handle); + return -1; + } + + if (flock(handle->fd, LOCK_UN) != 0) { + set_error( + error_buffer, + error_buffer_size, + "flock unlock(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + atomic_store_explicit(pid_slot, 0, memory_order_release); + (void)munlock(handle->shared, handle->mapping_size); + (void)munmap(handle->shared, handle->mapping_size); + close(handle->fd); + free(handle); + return -1; + } + + int32_t peer_pid = atomic_load_explicit( + &handle->shared->rank[1 - rank].pid, memory_order_acquire + ); + if (peer_pid > 0) { + uint32_t expected_unlinked = 0; + if (atomic_compare_exchange_strong_explicit( + &handle->shared->unlinked, + &expected_unlinked, + 1U, + memory_order_acq_rel, + memory_order_acquire + )) { + if (shm_unlink(handle->shm_name) != 0) { + set_error( + error_buffer, + error_buffer_size, + "shm_unlink(%s) failed: %s", + handle->shm_name, + strerror(errno) + ); + publish_shared_error(handle, 1U); + atomic_store_explicit(pid_slot, 0, memory_order_release); + (void)munlock(handle->shared, handle->mapping_size); + (void)munmap(handle->shared, handle->mapping_size); + close(handle->fd); + free(handle); + return -1; + } + } + } + + *output_handle = handle; + return 0; +} + +static int peer_is_dead( + struct sglang_dp2_sync_handle *handle, + int peer_rank +) { + int32_t peer_pid = atomic_load_explicit( + &handle->shared->rank[peer_rank].pid, memory_order_acquire + ); + if (peer_pid <= 0) { + return 0; + } + if (kill((pid_t)peer_pid, 0) == 0 || errno == EPERM) { + return 0; + } + return errno == ESRCH; +} + +int sglang_dp2_sync_exchange( + void *opaque_handle, + const int64_t *local_payload, + int64_t *global_payload, + struct sglang_dp2_sync_stats *stats, + char *error_buffer, + size_t error_buffer_size +) { + if (error_buffer != NULL && error_buffer_size > 0) { + error_buffer[0] = '\0'; + } + if ( + opaque_handle == NULL || + local_payload == NULL || + global_payload == NULL || + stats == NULL || + error_buffer == NULL || + error_buffer_size < SGLANG_DP2_SYNC_ERROR_BUFFER_MIN + ) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync received invalid exchange arguments" + ); + return -1; + } + + struct sglang_dp2_sync_handle *handle = opaque_handle; + struct sglang_dp2_shared_state *shared = handle->shared; + const int rank = handle->rank; + const int peer_rank = 1 - rank; + struct sglang_dp2_rank_state *local_state = &shared->rank[rank]; + struct sglang_dp2_rank_state *peer_state = &shared->rank[peer_rank]; + + uint32_t shared_error = atomic_load_explicit( + &shared->error_code, memory_order_acquire + ); + if (shared_error != 0) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync peer previously published error %u", + shared_error + ); + return -1; + } + if ( + atomic_load_explicit(&local_state->pid, memory_order_acquire) != + (int32_t)handle->pid + ) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync rank %d ownership changed", + rank + ); + publish_shared_error(handle, 2U); + return -1; + } + + const uint64_t started_ns = monotonic_ns(); + if (started_ns == 0) { + set_error( + error_buffer, + error_buffer_size, + "clock_gettime(CLOCK_MONOTONIC_RAW) failed" + ); + publish_shared_error(handle, 3U); + return -1; + } + const uint64_t sequence = ++handle->local_sequence; + const uint64_t prior_sequence = atomic_load_explicit( + &local_state->published_sequence, memory_order_acquire + ); + if (prior_sequence != sequence - 1U) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync local sequence mismatch: expected %llu, " + "found %llu", + (unsigned long long)(sequence - 1U), + (unsigned long long)prior_sequence + ); + publish_shared_error(handle, 4U); + return -1; + } + + const size_t slot = (size_t)(sequence % SGLANG_DP2_SYNC_SLOTS); + memcpy( + local_state->payload[slot], + local_payload, + sizeof(local_state->payload[slot]) + ); + local_state->arrival_ns[slot] = started_ns; + atomic_store_explicit( + &local_state->published_sequence, sequence, memory_order_release + ); + wake_rank(local_state); + const uint64_t wait_started_ns = monotonic_ns(); + + uint64_t peer_sequence = atomic_load_explicit( + &peer_state->published_sequence, memory_order_acquire + ); + while (peer_sequence < sequence) { + shared_error = atomic_load_explicit( + &shared->error_code, memory_order_acquire + ); + if (shared_error != 0) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync peer published error %u while " + "waiting for sequence %llu", + shared_error, + (unsigned long long)sequence + ); + return -1; + } + if (peer_is_dead(handle, peer_rank)) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync peer rank %d died while waiting for " + "sequence %llu", + peer_rank, + (unsigned long long)sequence + ); + publish_shared_error(handle, 5U); + return -1; + } + + const uint64_t now_ns = monotonic_ns(); + const uint64_t elapsed_ns = now_ns - started_ns; + if (elapsed_ns >= handle->timeout_ns) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync timed out after %.3f ms waiting for " + "rank %d sequence %llu (peer sequence %llu)", + (double)elapsed_ns / 1000000.0, + peer_rank, + (unsigned long long)sequence, + (unsigned long long)peer_sequence + ); + publish_shared_error(handle, 6U); + return -1; + } + + const uint32_t epoch = atomic_load_explicit( + &peer_state->futex_epoch, memory_order_acquire + ); + peer_sequence = atomic_load_explicit( + &peer_state->published_sequence, memory_order_acquire + ); + if (peer_sequence >= sequence) { + break; + } + const uint64_t remaining_ns = handle->timeout_ns - elapsed_ns; + const uint64_t slice_ns = + remaining_ns < SGLANG_DP2_SYNC_WAIT_SLICE_NS + ? remaining_ns + : SGLANG_DP2_SYNC_WAIT_SLICE_NS; + if (futex_wait_shared(&peer_state->futex_epoch, epoch, slice_ns) != 0) { + if ( + errno != EAGAIN && + errno != EINTR && + errno != ETIMEDOUT + ) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync futex wait failed: %s", + strerror(errno) + ); + publish_shared_error(handle, 7U); + return -1; + } + } + peer_sequence = atomic_load_explicit( + &peer_state->published_sequence, memory_order_acquire + ); + } + + if (peer_sequence > sequence + 1U) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync peer lapped local rank: local sequence " + "%llu, peer sequence %llu", + (unsigned long long)sequence, + (unsigned long long)peer_sequence + ); + publish_shared_error(handle, 8U); + return -1; + } + + const uint64_t peer_arrival_ns = peer_state->arrival_ns[slot]; + if (peer_arrival_ns == 0) { + set_error( + error_buffer, + error_buffer_size, + "DP2 shared-memory sync peer published sequence %llu without an " + "arrival timestamp", + (unsigned long long)sequence + ); + publish_shared_error(handle, 9U); + return -1; + } + + memcpy( + global_payload, + shared->rank[0].payload[slot], + sizeof(shared->rank[0].payload[slot]) + ); + memcpy( + global_payload + SGLANG_DP2_SYNC_WIDTH, + shared->rank[1].payload[slot], + sizeof(shared->rank[1].payload[slot]) + ); + const uint64_t finished_ns = monotonic_ns(); + const uint64_t latest_arrival_ns = + started_ns > peer_arrival_ns ? started_ns : peer_arrival_ns; + const uint64_t earliest_arrival_ns = + started_ns < peer_arrival_ns ? started_ns : peer_arrival_ns; + + stats->sequence = sequence; + stats->total_ns = finished_ns - started_ns; + stats->peer_wait_ns = finished_ns - wait_started_ns; + stats->arrival_skew_ns = latest_arrival_ns - earliest_arrival_ns; + stats->post_latest_arrival_ns = finished_ns - latest_arrival_ns; + return 0; +} + +int sglang_dp2_sync_exchange_values( + void *opaque_handle, + int64_t num_tokens, + int64_t num_tokens_for_logprob, + int64_t can_cuda_graph, + int64_t is_extend_in_batch, + int64_t local_can_run_tbo, + int64_t local_forward_mode, + int64_t can_run_breakable_cuda_graph, + int64_t *global_payload, + struct sglang_dp2_sync_stats *stats, + char *error_buffer, + size_t error_buffer_size +) { + const int64_t local_payload[SGLANG_DP2_SYNC_WIDTH] = { + num_tokens, + num_tokens_for_logprob, + can_cuda_graph, + is_extend_in_batch, + local_can_run_tbo, + local_forward_mode, + can_run_breakable_cuda_graph, + }; + return sglang_dp2_sync_exchange( + opaque_handle, + local_payload, + global_payload, + stats, + error_buffer, + error_buffer_size + ); +} + +void sglang_dp2_sync_close(void *opaque_handle) { + if (opaque_handle == NULL) { + return; + } + struct sglang_dp2_sync_handle *handle = opaque_handle; + if (handle->shared != NULL) { + _Atomic int32_t *pid_slot = + &handle->shared->rank[handle->rank].pid; + int32_t expected_pid = (int32_t)handle->pid; + (void)atomic_compare_exchange_strong_explicit( + pid_slot, + &expected_pid, + 0, + memory_order_acq_rel, + memory_order_acquire + ); + wake_rank(&handle->shared->rank[handle->rank]); + const int peer_rank = 1 - handle->rank; + const int32_t peer_pid = atomic_load_explicit( + &handle->shared->rank[peer_rank].pid, memory_order_acquire + ); + if (peer_pid <= 0) { + uint32_t expected_unlinked = 0; + if (atomic_compare_exchange_strong_explicit( + &handle->shared->unlinked, + &expected_unlinked, + 1U, + memory_order_acq_rel, + memory_order_acquire + )) { + (void)shm_unlink(handle->shm_name); + } + } + (void)munlock(handle->shared, handle->mapping_size); + (void)munmap(handle->shared, handle->mapping_size); + } + if (handle->fd >= 0) { + close(handle->fd); + } + free(handle); +} diff --git a/sgl-kernel/csrc/flashmla_extension.cc b/sgl-kernel/csrc/flashmla_extension.cc index b9f2fe00357c..85c9d971f44d 100644 --- a/sgl-kernel/csrc/flashmla_extension.cc +++ b/sgl-kernel/csrc/flashmla_extension.cc @@ -71,6 +71,30 @@ static std::tuple, std::option num_splits); } +static void sgl_sparse_prefill_fwd_into( + const at::Tensor& q, + const at::Tensor& kv, + const at::Tensor& indices, + double sm_scale, + int64_t d_v, + const std::optional& attn_sink, + const std::optional& topk_length, + const at::Tensor& out, + const at::Tensor& max_logits, + const at::Tensor& lse) { + sparse_attn_prefill_interface( + q, + kv, + indices, + static_cast(sm_scale), + static_cast(d_v), + attn_sink, + topk_length, + out, + max_logits, + lse); +} + TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { /* * From FlashMLA @@ -117,6 +141,11 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "Tensor? topk_length=None) -> Tensor[]"); m.impl("sparse_prefill_fwd", torch::kCUDA, &sparse_prefill_fwd); + m.def( + "sparse_prefill_fwd_into(Tensor q, Tensor kv, Tensor indices, float sm_scale, int d_v, Tensor? attn_sink, " + "Tensor? topk_length, Tensor(a!) out, Tensor(b!) max_logits, Tensor(c!) lse) -> ()"); + m.impl("sparse_prefill_fwd_into", torch::kCUDA, &sgl_sparse_prefill_fwd_into); + m.def( "fwd_kvcache_mla_fp8(Tensor q, Tensor kcache, int head_size_v, Tensor seqlens_k, Tensor block_table, float " "softmax_scale, bool is_causal, Tensor tile_scheduler_metadata, Tensor num_splits, Tensor? descale_q, Tensor? " diff --git a/sgl-kernel/python/sgl_kernel/flash_mla.py b/sgl-kernel/python/sgl_kernel/flash_mla.py index bd8d558a6826..8fbd57b369f7 100644 --- a/sgl-kernel/python/sgl_kernel/flash_mla.py +++ b/sgl-kernel/python/sgl_kernel/flash_mla.py @@ -315,6 +315,9 @@ def flash_mla_sparse_fwd( d_v: int = 512, attn_sink: Optional[torch.Tensor] = None, topk_length: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor] = None, + max_logits: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Sparse attention prefill kernel @@ -325,6 +328,9 @@ def flash_mla_sparse_fwd( indices: [s_q, h_kv, topk], int32. Invalid indices should be set to -1 or numbers >= s_kv sm_scale: float d_v: The dimension of value vectors. Can only be 512 + out: Optional preallocated [s_q, h_q, d_v] bfloat16 output. + max_logits: Optional preallocated [s_q, h_q] float32 scratch output. + lse: Optional preallocated [s_q, h_q] float32 scratch output. Returns: (output, max_logits, lse) @@ -336,6 +342,25 @@ def flash_mla_sparse_fwd( if _flashmla_import_error is not None: raise _IMPORT_ERROR from _flashmla_import_error + preallocated = (out, max_logits, lse) + if any(tensor is not None for tensor in preallocated): + if not all(tensor is not None for tensor in preallocated): + raise ValueError("out, max_logits, and lse must be provided together") + assert out is not None and max_logits is not None and lse is not None + torch.ops.sgl_kernel.sparse_prefill_fwd_into.default( + q, + kv, + indices, + sm_scale, + d_v, + attn_sink, + topk_length, + out, + max_logits, + lse, + ) + return out, max_logits, lse + results = torch.ops.sgl_kernel.sparse_prefill_fwd.default( q, kv, indices, sm_scale, d_v, attn_sink, topk_length ) diff --git a/test/registered/unit/batch_overlap/test_deepseek_v4_prefill_tbo.py b/test/registered/unit/batch_overlap/test_deepseek_v4_prefill_tbo.py new file mode 100644 index 000000000000..480cce8a36fd --- /dev/null +++ b/test/registered/unit/batch_overlap/test_deepseek_v4_prefill_tbo.py @@ -0,0 +1,191 @@ +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.batch_overlap.two_batch_overlap import ( + TboDPAttentionPreparer, + TboForwardBatchPreparer, +) +from sglang.srt.layers.attention.deepseek_v4_backend import DeepseekV4AttnBackend +from sglang.srt.layers.attention.tbo_backend import ( + _build_tbo_prefill_child_replay_fb_view, + _split_prefill_replay_layout, +) +from sglang.srt.managers.overlap_utils import decide_needs_cpu_seq_lens +from sglang.srt.model_executor.forward_batch_info import ForwardMode +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class TestDeepSeekV4PrefillTBO(CustomTestCase): + def _assert_layout( + self, + extend_lens: list[int], + prefix_lens: list[int], + boundary: int, + expected_left: list[int], + expected_right: list[int], + ) -> None: + left, right = _split_prefill_replay_layout( + extend_seq_lens=extend_lens, + extend_prefix_lens=prefix_lens, + split_token_index=boundary, + ) + self.assertEqual(left.extend_seq_lens, expected_left) + self.assertEqual(right.extend_seq_lens, expected_right) + self.assertEqual(left.extend_num_tokens, sum(expected_left)) + self.assertEqual(right.extend_num_tokens, sum(expected_right)) + self.assertEqual( + left.extend_num_tokens + right.extend_num_tokens, + sum(extend_lens), + ) + + for index, (extend_len, prefix_len) in enumerate( + zip(extend_lens, prefix_lens, strict=True) + ): + left_len = expected_left[index] + right_len = expected_right[index] + if left_len: + self.assertEqual(left.extend_prefix_lens[index], prefix_len) + self.assertEqual(left.seq_lens[index], prefix_len + left_len) + else: + self.assertEqual(left.extend_prefix_lens[index], 0) + self.assertEqual(left.seq_lens[index], 0) + if right_len: + self.assertEqual( + right.extend_prefix_lens[index], + prefix_len + left_len, + ) + self.assertEqual( + right.seq_lens[index], + prefix_len + extend_len, + ) + else: + self.assertEqual(right.extend_prefix_lens[index], 0) + self.assertEqual(right.seq_lens[index], 0) + + def test_fixed_boundary_layouts(self) -> None: + self._assert_layout([100], [7], 50, [50], [50]) + self._assert_layout( + [10, 20, 30, 0], + [1, 2, 3, 0], + 25, + [10, 15, 0, 0], + [0, 5, 30, 0], + ) + self._assert_layout([8, 8, 8], [2, 4, 6], 16, [8, 8, 0], [0, 0, 8]) + self._assert_layout([7, 3, 2], [0, 10, 20], 32, [7, 3, 2], [0, 0, 0]) + + def test_child_replay_view_preserves_request_geometry(self) -> None: + parent = SimpleNamespace( + batch_size=4, + forward_mode=ForwardMode.EXTEND, + actual_forward_mode=ForwardMode.EXTEND, + input_ids=torch.arange(64, dtype=torch.int64), + positions=torch.arange(64, dtype=torch.int64), + out_cache_loc=torch.arange(64, dtype=torch.int64), + req_pool_indices=torch.tensor([11, 12, 13, 0], dtype=torch.int64), + seq_lens=torch.tensor([11, 22, 33, 0], dtype=torch.int64), + extend_seq_lens=torch.tensor([10, 20, 30, 0], dtype=torch.int64), + extend_prefix_lens=torch.tensor([1, 2, 3, 0], dtype=torch.int64), + extend_start_loc=torch.tensor([0, 10, 30, 60], dtype=torch.int64), + ) + left, right = _split_prefill_replay_layout( + extend_seq_lens=[10, 20, 30, 0], + extend_prefix_lens=[1, 2, 3, 0], + split_token_index=25, + ) + left_view = _build_tbo_prefill_child_replay_fb_view( + parent, + layout=left, + tok_slice=slice(None, 32), + ) + right_view = _build_tbo_prefill_child_replay_fb_view( + parent, + layout=right, + tok_slice=slice(32, None), + ) + self.assertEqual(left_view.req_pool_indices.tolist(), [11, 12, 0, 0]) + self.assertEqual(right_view.req_pool_indices.tolist(), [0, 12, 13, 0]) + self.assertEqual(left_view.extend_start_loc.tolist(), [0, 10, 25, 25]) + self.assertEqual(right_view.extend_start_loc.tolist(), [0, 0, 5, 35]) + self.assertEqual(left_view.extend_num_tokens, 25) + self.assertEqual(right_view.extend_num_tokens, 35) + + def test_capture_child_padding_uses_zero_sentinels(self) -> None: + child = SimpleNamespace( + batch_size=1, + input_ids=torch.arange(32, dtype=torch.int64), + req_pool_indices=torch.tensor([9], dtype=torch.int64), + seq_lens=torch.tensor([37], dtype=torch.int64), + seq_lens_cpu=torch.tensor([37], dtype=torch.int64), + orig_seq_lens=torch.tensor([37], dtype=torch.int64), + extend_seq_lens=torch.tensor([32], dtype=torch.int64), + extend_prefix_lens=torch.tensor([5], dtype=torch.int64), + extend_start_loc=torch.tensor([0], dtype=torch.int64), + extend_seq_lens_cpu=[32], + extend_prefix_lens_cpu=[5], + extend_logprob_start_lens_cpu=[32], + ) + TboForwardBatchPreparer.pad_sequence_axis_for_cuda_graph( + child, + target_batch_size=4, + ) + self.assertEqual(child.batch_size, 4) + self.assertEqual(child.req_pool_indices.tolist(), [9, 0, 0, 0]) + self.assertEqual(child.seq_lens.tolist(), [37, 0, 0, 0]) + self.assertEqual(child.extend_seq_lens.tolist(), [32, 0, 0, 0]) + self.assertEqual(child.extend_prefix_lens.tolist(), [5, 0, 0, 0]) + self.assertEqual(child.extend_start_loc.tolist(), [0, 32, 32, 32]) + + def test_phase_and_dp_sync_gates(self) -> None: + self.assertTrue( + DeepseekV4AttnBackend.tbo_supports_cuda_graph_for(ForwardMode.EXTEND) + ) + self.assertFalse( + DeepseekV4AttnBackend.tbo_supports_cuda_graph_for(ForwardMode.DECODE) + ) + self.assertFalse( + DeepseekV4AttnBackend.tbo_supports_cuda_graph_for(ForwardMode.TARGET_VERIFY) + ) + self.assertFalse(DeepseekV4AttnBackend.tbo_supports_decode_cuda_graph) + self.assertFalse(DeepseekV4AttnBackend.tbo_requires_global_cpu_seq_lens) + + server_args = SimpleNamespace( + enable_two_batch_overlap=True, + speculative_algorithm="DSPARK", + ) + self.assertFalse( + decide_needs_cpu_seq_lens( + server_args, + [ + SimpleNamespace( + tbo_requires_global_cpu_seq_lens=False, + needs_cpu_seq_lens=False, + ) + ], + ) + ) + + preparer = TboDPAttentionPreparer() + preparer.enable_two_batch_overlap = True + preparer.local_tbo_split_seq_index = 7 + split, mode = preparer.compute_output_from_values( + local_can_run_tbo=[1, 1], + forward_modes=[ForwardMode.EXTEND.value, ForwardMode.IDLE.value], + ) + self.assertEqual(split, 7) + self.assertEqual(mode, ForwardMode.EXTEND) + split, mode = preparer.compute_output_from_values( + local_can_run_tbo=[1, 0], + forward_modes=[ForwardMode.EXTEND.value, ForwardMode.EXTEND.value], + ) + self.assertIsNone(split) + self.assertIsNone(mode) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/constrained/test_base_grammar_backend.py b/test/registered/unit/constrained/test_base_grammar_backend.py index 4c6265ed6930..039c935fe4bb 100644 --- a/test/registered/unit/constrained/test_base_grammar_backend.py +++ b/test/registered/unit/constrained/test_base_grammar_backend.py @@ -15,6 +15,8 @@ python -m pytest test_base_grammar_backend.py -v """ +import os +import threading import unittest from concurrent.futures import Future from unittest.mock import MagicMock, patch @@ -33,6 +35,35 @@ register_cpu_ci(2.0, "base-a-test-cpu") +class _SizedGrammar(BaseGrammarObject): + def __init__(self, value: str, size: int): + super().__init__() + self.value = value + self.size = size + self.grammar_stats = GrammarStats() + + def copy(self): + return _SizedGrammar(self.value, self.size) + + def cache_memory_bytes(self): + return self.size + + +class _BlockingGrammarBackend(BaseGrammarBackend): + def __init__(self): + self.started = threading.Event() + self.release = threading.Event() + self.compile_count = 0 + super().__init__() + + def dispatch_json(self, key_string): + self.compile_count += 1 + self.started.set() + if not self.release.wait(timeout=5): + raise TimeoutError("test did not release grammar compilation") + return _SizedGrammar(key_string, 8) + + class TestGrammarStats(unittest.TestCase): """Test GrammarStats dataclass.""" @@ -180,27 +211,67 @@ def test_init_value_dispatch_none_grammar(self): result = self.backend._init_value_dispatch(("json", "schema"), False) self.assertIsNone(result) - def test_cache_miss_duplicate_key_submits_separate_futures(self): - """Two cache misses for the same key each get their own Future. - - The backend does not deduplicate in-flight compilations — that is - handled at the GrammarManager level via grammar_queue. Each call - to get_cached_or_future_value with an uncached key submits a new - task to the executor.""" - key = ("json", "schema") - result1, hit1 = self.backend.get_cached_or_future_value(key, False) - result2, hit2 = self.backend.get_cached_or_future_value(key, False) + def test_cache_miss_duplicate_key_shares_compilation(self): + backend = _BlockingGrammarBackend() + self.addCleanup(backend.executor.shutdown, wait=True) + pretty = '{"type": "object", "properties": {"x": {"type": "string"}}}' + compact = '{"type":"object","properties":{"x":{"type":"string"}}}' + result1, hit1 = backend.get_cached_or_future_value( + ("json", pretty), require_reasoning=False + ) + self.assertTrue(backend.started.wait(timeout=5)) + result2, hit2 = backend.get_cached_or_future_value( + ("json", compact), require_reasoning=False + ) self.assertFalse(hit1) self.assertFalse(hit2) self.assertIsInstance(result1, Future) self.assertIsInstance(result2, Future) - # They are independent futures, not shared + # Each request gets an independently cancellable wrapper. self.assertIsNot(result1, result2) - # Both should complete successfully - self.assertIsInstance(result1.result(timeout=5), InvalidGrammarObject) - self.assertIsInstance(result2.result(timeout=5), InvalidGrammarObject) + backend.release.set() + first = result1.result(timeout=5) + second = result2.result(timeout=5) + self.assertEqual(backend.compile_count, 1) + self.assertEqual(first.grammar_stats.cache_source, "compile") + self.assertEqual(second.grammar_stats.cache_source, "inflight") + + def test_immediate_compile_completion_does_not_deadlock(self): + completed = Future() + completed.set_result(InvalidGrammarObject("done")) + self.backend.executor.shutdown(wait=True) + executor = MagicMock() + executor.submit.return_value = completed + executor.shutdown.return_value = None + self.backend.executor = executor + + result, cache_hit = self.backend.get_cached_or_future_value( + ("json", "schema"), require_reasoning=False + ) + + self.assertFalse(cache_hit) + self.assertEqual(result.result(timeout=1).error_message, "done") + self.assertIn(("json", "schema"), self.backend.cache) + + def test_cache_is_bounded_by_memory_and_lru_order(self): + with patch.dict( + os.environ, + { + "SGLANG_GRAMMAR_CACHE_BYTES": "10", + "SGLANG_GRAMMAR_CACHE_ENTRIES": "100", + }, + ): + backend = BaseGrammarBackend() + self.addCleanup(backend.executor.shutdown, wait=True) + + backend.set_cache(("regex", "first"), _SizedGrammar("first", 8)) + backend.set_cache(("regex", "second"), _SizedGrammar("second", 8)) + + self.assertNotIn(("regex", "first"), backend.cache) + self.assertIn(("regex", "second"), backend.cache) + self.assertEqual(backend._cache_bytes, 8) class TestRegisterGrammarBackend(unittest.TestCase): diff --git a/test/registered/unit/constrained/test_xgrammar_persistent_cache.py b/test/registered/unit/constrained/test_xgrammar_persistent_cache.py new file mode 100644 index 000000000000..f6cf0804e81e --- /dev/null +++ b/test/registered/unit/constrained/test_xgrammar_persistent_cache.py @@ -0,0 +1,216 @@ +import json +import os +import tempfile +import unittest +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import patch + +from sglang.srt.constrained import xgrammar_persistent_cache as cache_module +from sglang.srt.constrained.xgrammar_persistent_cache import ( + PersistentXGrammarCache, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class _FakeTokenizerInfo: + def serialize_json(self) -> str: + return '{"vocab":["a","b"]}' + + +class _FakeCompiledGrammar: + def __init__(self, value: str) -> None: + self.value = value + + def serialize_json(self) -> str: + return json.dumps({"value": self.value}) + + @classmethod + def deserialize_json( + cls, + serialized: str, + tokenizer_info: _FakeTokenizerInfo, + ) -> "_FakeCompiledGrammar": + del tokenizer_info + payload = json.loads(serialized) + if set(payload) != {"value"} or not isinstance(payload["value"], str): + raise ValueError("invalid compiled grammar") + return cls(payload["value"]) + + +class TestPersistentXGrammarCache(CustomTestCase): + def _make_cache( + self, + directory: str, + *, + deserialize_bytes_per_second: int = 10**18, + max_bytes: int = 1024**2, + ) -> PersistentXGrammarCache: + return PersistentXGrammarCache( + tokenizer_info=_FakeTokenizerInfo(), + cache_directory=directory, + max_bytes=max_bytes, + deserialize_bytes_per_second=deserialize_bytes_per_second, + local_compile_speedup=2, + compiler_identity={ + "any_whitespace": True, + "compiler_threads": 2, + "override_stop_tokens": [], + "vocab_size": 2, + }, + ) + + def test_serialized_round_trip_and_corruption_are_fail_closed(self) -> None: + compile_calls = 0 + + def compile_grammar() -> _FakeCompiledGrammar: + nonlocal compile_calls + compile_calls += 1 + return _FakeCompiledGrammar("root") + + with ( + tempfile.TemporaryDirectory(prefix="xgrammar-cache-test-") as cache_dir, + patch.object(cache_module, "CompiledGrammar", _FakeCompiledGrammar), + patch.dict( + os.environ, + {"SGLANG_XGRAMMAR_CACHE_SESSION_ID": "session-1"}, + clear=False, + ), + ): + cache = self._make_cache(cache_dir) + first = cache.get_or_compile( + key_type="ebnf", + key_string='root ::= "ab"', + compile_fn=compile_grammar, + ) + self.assertEqual(first.source, "compile") + self.assertEqual(first.grammar.value, "root") + self.assertEqual( + set(first.phase_seconds), + { + "native_compile", + "serialize", + "entry_write", + "account_prune", + "policy_serialized", + }, + ) + + second_cache = self._make_cache(cache_dir) + second = second_cache.get_or_compile( + key_type="ebnf", + key_string='root ::= "ab"', + compile_fn=compile_grammar, + ) + self.assertEqual(second.source, "disk") + self.assertEqual(second.grammar.value, "root") + self.assertEqual(compile_calls, 1) + + entry = next(second_cache.entries.glob("*.json")) + entry.write_text("corrupt", encoding="utf-8") + with self.assertRaisesRegex(RuntimeError, "cache entry is invalid"): + second_cache.get_or_compile( + key_type="ebnf", + key_string='root ::= "ab"', + compile_fn=compile_grammar, + ) + self.assertEqual(compile_calls, 1) + + def test_adaptive_local_compile_marker_and_checksum(self) -> None: + compile_calls = 0 + + def compile_grammar() -> _FakeCompiledGrammar: + nonlocal compile_calls + compile_calls += 1 + return _FakeCompiledGrammar("root") + + with ( + tempfile.TemporaryDirectory( + prefix="xgrammar-local-compile-test-" + ) as cache_dir, + patch.object(cache_module, "CompiledGrammar", _FakeCompiledGrammar), + ): + cache = self._make_cache( + cache_dir, + deserialize_bytes_per_second=1, + ) + first = cache.get_or_compile( + key_type="ebnf", + key_string='root ::= "ab"', + compile_fn=compile_grammar, + ) + self.assertEqual(first.source, "local_compile") + entry = next(cache.entries.glob("*.json")) + marker = entry.read_bytes() + self.assertTrue(marker.startswith(cache._LOCAL_COMPILE_MAGIC)) + self.assertEqual(len(marker), cache._LOCAL_COMPILE_MARKER_BYTES) + + peer = self._make_cache( + cache_dir, + deserialize_bytes_per_second=1, + ) + second = peer.get_or_compile( + key_type="ebnf", + key_string='root ::= "ab"', + compile_fn=compile_grammar, + ) + self.assertEqual(second.source, "local_compile") + self.assertEqual( + set(second.phase_seconds), + {"adaptive_marker_read", "native_compile"}, + ) + self.assertEqual(compile_calls, 2) + + corrupted_marker = bytearray(marker) + corrupted_marker[-1] ^= 1 + entry.write_bytes(corrupted_marker) + with self.assertRaisesRegex(RuntimeError, "cache entry is invalid"): + peer.get_or_compile( + key_type="ebnf", + key_string='root ::= "ab"', + compile_fn=compile_grammar, + ) + self.assertEqual(compile_calls, 2) + + def test_concurrent_accounting_and_session_recovery(self) -> None: + with ( + tempfile.TemporaryDirectory( + prefix="xgrammar-accounting-test-" + ) as cache_dir, + patch.object(cache_module, "CompiledGrammar", _FakeCompiledGrammar), + patch.dict( + os.environ, + {"SGLANG_XGRAMMAR_CACHE_SESSION_ID": "session-1"}, + clear=False, + ), + ): + cache = self._make_cache(cache_dir) + + def compile_unique(index: int) -> None: + cache.get_or_compile( + key_type="ebnf", + key_string=f'root ::= "ab" /* {index} */', + compile_fn=lambda: _FakeCompiledGrammar(str(index)), + ) + + with ThreadPoolExecutor(max_workers=8) as executor: + list(executor.map(compile_unique, range(64))) + + entry_bytes = sum( + entry.stat().st_size for entry in cache.entries.glob("*.json") + ) + self.assertEqual(cache._read_size_ledger(), entry_bytes) + + cache._size_ledger.write_bytes(b"invalid") + with self.assertRaisesRegex(RuntimeError, "size ledger is invalid"): + self._make_cache(cache_dir) + + os.environ["SGLANG_XGRAMMAR_CACHE_SESSION_ID"] = "session-2" + recovered = self._make_cache(cache_dir) + self.assertEqual(recovered._read_size_ledger(), entry_bytes) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/entrypoints/test_scheduler_health.py b/test/registered/unit/entrypoints/test_scheduler_health.py new file mode 100644 index 000000000000..ce6b85b269c1 --- /dev/null +++ b/test/registered/unit/entrypoints/test_scheduler_health.py @@ -0,0 +1,122 @@ +import time +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.entrypoints import http_server +from sglang.srt.managers.tokenizer_manager import ServerStatus +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +class _FakeTokenizerManager: + def __init__( + self, + snapshots: list[SimpleNamespace], + *, + gracefully_exit: bool = False, + server_status: ServerStatus = ServerStatus.Up, + error: Exception | None = None, + ) -> None: + self.snapshots = snapshots + self.gracefully_exit = gracefully_exit + self.server_status = server_status + self.elastic_worker_count = 2 + self.error = error + self.get_loads_calls = 0 + + async def get_loads( + self, + include: list[str] | None = None, + ) -> list[SimpleNamespace]: + self.get_loads_calls += 1 + if include != ["core"]: + raise AssertionError(f"unexpected load fields: {include}") + if self.error is not None: + raise self.error + return self.snapshots + + +def _snapshot(dp_rank: int, timestamp: float) -> SimpleNamespace: + return SimpleNamespace(dp_rank=dp_rank, timestamp=timestamp) + + +class TestSchedulerHealth(unittest.IsolatedAsyncioTestCase): + async def _status(self, manager: _FakeTokenizerManager | None) -> int: + global_state = ( + None if manager is None else SimpleNamespace(tokenizer_manager=manager) + ) + with patch.object(http_server, "_global_state", global_state): + return (await http_server.health_scheduler()).status_code + + async def test_healthy_when_all_dp_snapshots_are_fresh(self) -> None: + now = time.time() + manager = _FakeTokenizerManager( + [_snapshot(0, now), _snapshot(1, now)], + ) + self.assertEqual(await self._status(manager), 200) + self.assertEqual(manager.get_loads_calls, 1) + + async def test_fails_closed_for_bad_snapshot_sets(self) -> None: + now = time.time() + cases = { + "missing rank": [_snapshot(0, now)], + "duplicate rank": [_snapshot(0, now), _snapshot(0, now)], + "stale timestamp": [ + _snapshot( + 0, + now - http_server.SCHEDULER_HEALTH_MAX_STALENESS_SECONDS - 1, + ), + _snapshot(1, now), + ], + "future timestamp": [ + _snapshot( + 0, + now + http_server.SCHEDULER_HEALTH_MAX_FUTURE_SKEW_SECONDS + 1, + ), + _snapshot(1, now), + ], + "invalid timestamp": [_snapshot(0, 0), _snapshot(1, now)], + } + for name, snapshots in cases.items(): + with self.subTest(name=name): + self.assertEqual( + await self._status(_FakeTokenizerManager(snapshots)), + 503, + ) + + async def test_fails_closed_when_state_is_unavailable(self) -> None: + now = time.time() + cases = { + "load read error": _FakeTokenizerManager( + [_snapshot(0, now), _snapshot(1, now)], + error=RuntimeError("synthetic load snapshot failure"), + ), + "starting": _FakeTokenizerManager( + [_snapshot(0, now), _snapshot(1, now)], + server_status=ServerStatus.Starting, + ), + "shutting down": _FakeTokenizerManager( + [_snapshot(0, now), _snapshot(1, now)], + gracefully_exit=True, + ), + } + for name, manager in cases.items(): + with self.subTest(name=name): + self.assertEqual(await self._status(manager), 503) + + self.assertEqual(await self._status(None), 503) + + def test_route_is_registered_once(self) -> None: + routes = [ + route + for route in http_server.app.routes + if getattr(route, "path", None) == "/health_scheduler" + ] + self.assertEqual(len(routes), 1) + self.assertEqual(routes[0].methods, {"GET"}) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_atomic_dp_batch_routing.py b/test/registered/unit/managers/test_atomic_dp_batch_routing.py new file mode 100644 index 000000000000..cace9ea35a1f --- /dev/null +++ b/test/registered/unit/managers/test_atomic_dp_batch_routing.py @@ -0,0 +1,105 @@ +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +from sglang.srt.managers import data_parallel_controller +from sglang.srt.managers.data_parallel_controller import DataParallelController +from sglang.srt.managers.io_struct import ( + BatchTokenizedGenerateReqInput, + wrap_as_pickle, +) +from sglang.srt.managers.tokenizer_manager import TokenizerManager +from sglang.srt.observability.req_time_stats import ( + APIServerReqTimeStats, + DPControllerReqTimeStats, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _request(dp_rank: int | None) -> SimpleNamespace: + return SimpleNamespace( + routed_dp_rank=dp_rank, + time_stats=wrap_as_pickle(APIServerReqTimeStats()), + ) + + +class TestAtomicDPBatchRouting(CustomTestCase): + def test_tokenizer_batches_only_a_single_explicit_dp_route(self) -> None: + manager = object.__new__(TokenizerManager) + manager.server_args = SimpleNamespace( + enable_tokenizer_batch_encode=False, + enable_dp_attention=True, + ) + manager._batch_has_text = lambda _batch_size, _requests: False + + self.assertTrue( + manager._should_use_batch_tokenization( + 4, + [_request(1) for _ in range(4)], + ) + ) + self.assertFalse( + manager._should_use_batch_tokenization( + 4, + [_request(0), _request(1), _request(0), _request(1)], + ) + ) + self.assertFalse( + manager._should_use_batch_tokenization( + 4, + [_request(None) for _ in range(4)], + ) + ) + + manager.server_args.enable_dp_attention = False + self.assertTrue( + manager._should_use_batch_tokenization( + 4, + [_request(None) for _ in range(4)], + ) + ) + + def test_controller_sends_single_rank_batch_as_one_message(self) -> None: + controller = object.__new__(DataParallelController) + target_worker = object() + controller.workers = [object(), target_worker] + controller._active_workers = [0, 1] + controller.refresh_load_budget_on_dispatch = False + fallback_requests = [] + controller.dispatching = fallback_requests.append + batch = BatchTokenizedGenerateReqInput(batch=[_request(1) for _ in range(4)]) + + with patch.object(data_parallel_controller, "sock_send") as sock_send: + controller.dispatch_batch_generate(batch) + + sock_send.assert_called_once_with(target_worker, batch) + self.assertEqual(fallback_requests, []) + self.assertTrue( + all(isinstance(req.time_stats, DPControllerReqTimeStats) for req in batch) + ) + self.assertTrue(all(req.time_stats.dpc_dispatch_time > 0 for req in batch)) + self.assertTrue( + all(req.time_stats.dpc_dispatch_finish_time > 0 for req in batch) + ) + + def test_controller_falls_back_for_mixed_rank_batch(self) -> None: + controller = object.__new__(DataParallelController) + controller.workers = [object(), object()] + controller._active_workers = [0, 1] + controller.refresh_load_budget_on_dispatch = False + fallback_requests = [] + controller.dispatching = fallback_requests.append + batch = BatchTokenizedGenerateReqInput(batch=[_request(0), _request(1)]) + + with patch.object(data_parallel_controller, "sock_send") as sock_send: + controller.dispatch_batch_generate(batch) + + sock_send.assert_not_called() + self.assertEqual(fallback_requests, list(batch)) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py index ef5e2056aa22..39ab33d84530 100644 --- a/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py +++ b/test/registered/unit/managers/test_batch_result_processor_spec_grammar.py @@ -42,6 +42,9 @@ def is_none(self) -> bool: def is_dflash(self) -> bool: return False + def is_dspark(self) -> bool: + return False + class _FakeForwardMode: def is_decode(self) -> bool: diff --git a/test/registered/unit/managers/test_single_node_dp2_sync.py b/test/registered/unit/managers/test_single_node_dp2_sync.py new file mode 100644 index 000000000000..28d946d814c4 --- /dev/null +++ b/test/registered/unit/managers/test_single_node_dp2_sync.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import ctypes +import multiprocessing +import os +import shutil +import subprocess +import tempfile +import time +import unittest +import uuid +from pathlib import Path +from unittest.mock import patch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=8, suite="base-a-test-cpu") + +_WORLD_SIZE = 2 +_PAYLOAD_WIDTH = 7 +_ERROR_BUFFER_SIZE = 512 +_EXPECTED_ABI_VERSION = 2 + + +class _SyncStats(ctypes.Structure): + _fields_ = [ + ("sequence", ctypes.c_uint64), + ("total_ns", ctypes.c_uint64), + ("peer_wait_ns", ctypes.c_uint64), + ("arrival_skew_ns", ctypes.c_uint64), + ("post_latest_arrival_ns", ctypes.c_uint64), + ] + + +def _load_library(path: str) -> ctypes.CDLL: + library = ctypes.CDLL(path) + library.sglang_dp2_sync_abi_version.argtypes = [] + library.sglang_dp2_sync_abi_version.restype = ctypes.c_uint32 + library.sglang_dp2_sync_open.argtypes = [ + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_uint64, + ctypes.POINTER(ctypes.c_void_p), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.sglang_dp2_sync_open.restype = ctypes.c_int + library.sglang_dp2_sync_exchange.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(ctypes.c_int64), + ctypes.POINTER(_SyncStats), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.sglang_dp2_sync_exchange.restype = ctypes.c_int + library.sglang_dp2_sync_close.argtypes = [ctypes.c_void_p] + library.sglang_dp2_sync_close.restype = None + if library.sglang_dp2_sync_abi_version() != _EXPECTED_ABI_VERSION: + raise RuntimeError("unexpected native DP2 ABI") + return library + + +def _open_handle( + library: ctypes.CDLL, + session_id: str, + rank: int, + timeout_ns: int, +) -> ctypes.c_void_p: + handle = ctypes.c_void_p() + error = ctypes.create_string_buffer(_ERROR_BUFFER_SIZE) + result = library.sglang_dp2_sync_open( + session_id.encode(), + rank, + timeout_ns, + ctypes.byref(handle), + error, + len(error), + ) + if result != 0: + raise RuntimeError(error.value.decode()) + if handle.value is None: + raise RuntimeError("native DP2 open returned a null handle") + return handle + + +def _exchange_worker( + library_path: str, + session_id: str, + rank: int, + iterations: int, + ready, + result_queue, +) -> None: + try: + library = _load_library(library_path) + handle = _open_handle( + library, + session_id, + rank, + timeout_ns=5_000_000_000, + ) + ready.wait(timeout=10) + local = (ctypes.c_int64 * _PAYLOAD_WIDTH)() + gathered = (ctypes.c_int64 * (_WORLD_SIZE * _PAYLOAD_WIDTH))() + stats = _SyncStats() + error = ctypes.create_string_buffer(_ERROR_BUFFER_SIZE) + total_ns = [] + try: + for sequence in range(1, iterations + 1): + for index in range(_PAYLOAD_WIDTH): + local[index] = rank * 1_000_000_000 + sequence * 10 + index + result = library.sglang_dp2_sync_exchange( + handle, + local, + gathered, + ctypes.byref(stats), + error, + len(error), + ) + if result != 0: + raise RuntimeError(error.value.decode()) + if stats.sequence != sequence: + raise AssertionError( + f"sequence mismatch: {stats.sequence} != {sequence}" + ) + for gathered_rank in range(_WORLD_SIZE): + for index in range(_PAYLOAD_WIDTH): + expected = gathered_rank * 1_000_000_000 + sequence * 10 + index + actual = gathered[gathered_rank * _PAYLOAD_WIDTH + index] + if actual != expected: + raise AssertionError(f"{actual} != {expected}") + if sequence > 100: + total_ns.append(stats.total_ns) + finally: + library.sglang_dp2_sync_close(handle) + result_queue.put((rank, total_ns, None)) + except BaseException as error: + result_queue.put((rank, [], repr(error))) + + +class TestSingleNodeDP2Sync(CustomTestCase): + @classmethod + def setUpClass(cls) -> None: + super().setUpClass() + compiler = shutil.which("cc") + if compiler is None: + raise unittest.SkipTest("a C compiler is required") + cls._temp_dir = tempfile.TemporaryDirectory(prefix="sglang-dp2-sync-test-") + cls.library_path = str(Path(cls._temp_dir.name) / "sglang_dp2_sync.so") + source = ( + Path(__file__).resolve().parents[4] + / "sgl-kernel" + / "csrc" + / "cpu" + / "dp2_sync.c" + ) + subprocess.run( + [ + compiler, + "-O3", + "-std=c11", + "-fPIC", + "-shared", + "-Wall", + "-Wextra", + "-Werror", + "-o", + cls.library_path, + str(source), + "-lrt", + ], + check=True, + ) + + @classmethod + def tearDownClass(cls) -> None: + cls._temp_dir.cleanup() + super().tearDownClass() + + def test_invalid_rank_is_rejected(self) -> None: + library = _load_library(self.library_path) + handle = ctypes.c_void_p() + error = ctypes.create_string_buffer(_ERROR_BUFFER_SIZE) + result = library.sglang_dp2_sync_open( + f"invalid-rank-{uuid.uuid4()}".encode(), + 2, + 1_000_000_000, + ctypes.byref(handle), + error, + len(error), + ) + self.assertNotEqual(result, 0) + self.assertIn(b"rank must be 0 or 1", error.value) + self.assertIsNone(handle.value) + + def test_timeout_is_fail_closed(self) -> None: + library = _load_library(self.library_path) + handle = _open_handle( + library, + f"timeout-{uuid.uuid4()}", + rank=0, + timeout_ns=50_000_000, + ) + local = (ctypes.c_int64 * _PAYLOAD_WIDTH)(*range(_PAYLOAD_WIDTH)) + gathered = (ctypes.c_int64 * (_WORLD_SIZE * _PAYLOAD_WIDTH))() + stats = _SyncStats() + error = ctypes.create_string_buffer(_ERROR_BUFFER_SIZE) + started = time.monotonic() + try: + result = library.sglang_dp2_sync_exchange( + handle, + local, + gathered, + ctypes.byref(stats), + error, + len(error), + ) + finally: + library.sglang_dp2_sync_close(handle) + self.assertNotEqual(result, 0) + self.assertIn(b"timed out", error.value) + self.assertGreaterEqual(time.monotonic() - started, 0.04) + + def test_exchange_correctness_and_latency(self) -> None: + context = multiprocessing.get_context("spawn") + ready = context.Barrier(_WORLD_SIZE) + result_queue = context.Queue() + session_id = f"correctness-{uuid.uuid4()}" + processes = [ + context.Process( + target=_exchange_worker, + args=( + self.library_path, + session_id, + rank, + 2_000, + ready, + result_queue, + ), + ) + for rank in range(_WORLD_SIZE) + ] + for process in processes: + process.start() + + durations = [] + for _ in processes: + rank, total_ns, error = result_queue.get(timeout=30) + self.assertIsNone(error, f"rank {rank}: {error}") + durations.extend(total_ns) + for process in processes: + process.join(timeout=30) + self.assertEqual(process.exitcode, 0) + + durations.sort() + p50_us = durations[len(durations) // 2] / 1_000 + self.assertLess(p50_us, 1_000) + + def test_python_runtime_validation_loads_packaged_abi(self) -> None: + from sglang.srt.managers.scheduler_components import single_node_dp2_sync + + environment = { + "SGLANG_DSPARK_DP2_SHM_MLP_SYNC": "1", + "SGLANG_DSPARK_DP2_SHM_SESSION_ID": f"validate-{uuid.uuid4()}", + "SGLANG_DSPARK_DP2_SHM_TIMEOUT_MS": "5000", + "SGLANG_DSPARK_DP2_SHM_METRICS": "1", + "SGLANG_DSPARK_DP2_SHM_LIBRARY": self.library_path, + "SGLANG_SCHEDULER_SKIP_ALL_GATHER": "0", + "SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH": "0", + } + with ( + patch.dict(os.environ, environment, clear=False), + patch.object(single_node_dp2_sync, "_enabled", None), + ): + single_node_dp2_sync.validate_single_node_dp2_sync_runtime() + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py index 8c79c3565752..6a2f6fbecc82 100644 --- a/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py +++ b/test/registered/unit/mem_cache/test_hicache_staged_write_back_dispatch.py @@ -672,6 +672,22 @@ def test_logical_host_pool_preserves_page_first_group_layout(self): self.assertEqual(group.layout, "page_first") self.assertTrue(group.can_use_write_back_jit) + def test_logical_host_pool_supports_group_shutdown(self): + logical_host_pool = LogicalHostPool(8, 2, layout="page_first") + group = HostPoolGroup( + [ + PoolEntry( + name=PoolName.KV, + host_pool=logical_host_pool, + device_pool=None, + layer_mapper=lambda _: 0, + is_primary_index_anchor=True, + ) + ] + ) + + group.destroy() + def test_write_back_jit_hybrid_write_keeps_extra_host_indices_on_cpu(self): captured = {} diff --git a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py index 2e0baefd4a8f..6c0c489d9b16 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py +++ b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py @@ -1,6 +1,8 @@ import unittest from types import SimpleNamespace +import torch + from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, @@ -22,8 +24,9 @@ def _make_runner(self): runner.prefill_backend_name = Backend.TC_PIECEWISE runner.has_mha_companion_layers = False runner.capture_hidden_mode = CaptureHiddenMode.NULL - runner.capture_num_tokens = [4, 16] - runner.max_num_tokens = 16 + runner.capture_num_tokens = [4, 32] + runner.max_num_tokens = 32 + runner.enable_tbo_prefill_graph = False return runner def _make_forward_batch(self, num_tokens): @@ -35,20 +38,56 @@ def _make_forward_batch(self, num_tokens): forward_mode=ForwardMode.EXTEND, capture_hidden_mode=CaptureHiddenMode.NULL, global_num_tokens_cpu=None, + can_run_dp_breakable_cuda_graph=False, + can_run_tbo=False, return_logprob=False, input_ids=list(range(num_tokens)), ) - def test_rejects_more_than_two_x_token_padding(self): + def test_rejects_more_than_four_x_token_padding(self): runner = self._make_runner() self.assertFalse(runner.can_run_graph(self._make_forward_batch(5))) - def test_accepts_two_x_token_padding(self): + def test_accepts_four_x_token_padding(self): runner = self._make_runner() self.assertTrue(runner.can_run_graph(self._make_forward_batch(8))) + def test_full_graph_request_slots_shrink_with_token_bucket(self): + runner = self._make_runner() + runner._is_full_backend = True + runner._capture_req_slots = 64 + + self.assertEqual(runner._capture_req_slots_for_tokens(4), 4) + self.assertEqual(runner._capture_req_slots_for_tokens(32), 32) + self.assertEqual(runner._capture_req_slots_for_tokens(128), 64) + + def test_nested_body_output_trims_only_token_axis_tensors(self): + output = ( + ( + torch.arange(32), + torch.tensor(7), + ), + [ + torch.arange(64).reshape(32, 2), + torch.arange(5), + ], + ) + + trimmed = PrefillCudaGraphRunner._trim_replayed_body_output( + output, + raw_num_tokens=5, + static_num_tokens=32, + ) + + self.assertIsInstance(trimmed, tuple) + self.assertEqual(trimmed[0][0].shape, (5,)) + self.assertEqual(trimmed[0][1].shape, ()) + self.assertIsInstance(trimmed[1], list) + self.assertEqual(trimmed[1][0].shape, (5, 2)) + self.assertEqual(trimmed[1][1].shape, (5,)) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/observability/test_runtime_path_metrics.py b/test/registered/unit/observability/test_runtime_path_metrics.py new file mode 100644 index 000000000000..a55f2bb902a3 --- /dev/null +++ b/test/registered/unit/observability/test_runtime_path_metrics.py @@ -0,0 +1,152 @@ +import unittest +from functools import partial +from types import SimpleNamespace +from unittest.mock import patch + +from prometheus_client import ( + CollectorRegistry, + Counter, + Gauge, + Histogram, + Summary, + generate_latest, +) + +from sglang.srt.observability.metrics_collector import ( + RadixCacheMetricsCollector, + SchedulerMetricsCollector, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + + +def _metric_line(metrics: str, name: str, *required_labels: str) -> str: + return next( + line + for line in metrics.splitlines() + if line.startswith(name + "{") + and all(label in line for label in required_labels) + ) + + +class TestRuntimePathMetrics(CustomTestCase): + def test_scheduler_prefill_grammar_and_hicache_metrics(self) -> None: + registry = CollectorRegistry() + counter = partial(Counter, registry=registry) + gauge = partial(Gauge, registry=registry) + histogram = partial(Histogram, registry=registry) + summary = partial(Summary, registry=registry) + + with ( + patch.object(SchedulerMetricsCollector, "_counter_cls", counter), + patch.object(SchedulerMetricsCollector, "_gauge_cls", gauge), + patch.object(SchedulerMetricsCollector, "_histogram_cls", histogram), + patch.object(SchedulerMetricsCollector, "_summary_cls", summary), + patch.object(RadixCacheMetricsCollector, "_counter_cls", counter), + patch.object(RadixCacheMetricsCollector, "_gauge_cls", gauge), + patch.object(RadixCacheMetricsCollector, "_histogram_cls", histogram), + ): + radix = RadixCacheMetricsCollector( + labels={"cache_type": "test", "rank": "0"} + ) + radix.observe_hicache_scheduler_phase( + "write_completion_sync", + duration_seconds=0.125, + calls=4, + max_seconds=0.05, + ) + radix.set_hicache_pending_operations("write_acks", 3) + radix.observe_hicache_backup(512, 0.02) + + scheduler = SchedulerMetricsCollector( + labels={"model_name": "test", "dp_rank": "0", "moe_ep_rank": 0}, + server_args=SimpleNamespace( + enable_metrics=True, + prefill_delayer_forward_passes_buckets=None, + prefill_delayer_max_delay_passes=30, + prefill_delayer_wait_seconds_buckets=None, + ), + ) + scheduler.add_scheduler_phase( + "schedule_plan", + duration_seconds=0.25, + calls=8, + max_seconds=0.075, + ) + scheduler.set_runtime_gc_frozen() + scheduler.observe_prefill_execution( + outcome="cuda_graph", + scheduled_tokens=129, + executed_tokens=256, + requests=7, + bucket_tokens=256, + ) + + metrics = generate_latest(registry).decode() + expected = { + ( + "sglang:hicache_scheduler_phase_seconds_total", + ('cache_type="test"', 'phase="write_completion_sync"', 'rank="0"'), + ): "0.125", + ( + "sglang:hicache_scheduler_phase_calls_total", + ('cache_type="test"', 'phase="write_completion_sync"', 'rank="0"'), + ): "4.0", + ( + "sglang:hicache_scheduler_phase_max_seconds", + ('cache_type="test"', 'phase="write_completion_sync"', 'rank="0"'), + ): "0.05", + ( + "sglang:hicache_pending_operations", + ('cache_type="test"', 'kind="write_acks"', 'rank="0"'), + ): "3.0", + ( + "sglang:scheduler_phase_seconds_total", + ('dp_rank="0"', 'model_name="test"', 'phase="schedule_plan"'), + ): "0.25", + ( + "sglang:scheduler_phase_calls_total", + ('dp_rank="0"', 'model_name="test"', 'phase="schedule_plan"'), + ): "8.0", + ( + "sglang:scheduler_phase_max_seconds", + ('dp_rank="0"', 'model_name="test"', 'phase="schedule_plan"'), + ): "0.075", + ( + "sglang:runtime_gc_frozen", + ('dp_rank="0"', 'model_name="test"'), + ): "1.0", + ( + "sglang:prefill_graph_admissions_total", + ('bucket="256"', 'outcome="cuda_graph"'), + ): "1.0", + ( + "sglang:prefill_graph_shapes_total", + ('bucket="256"', 'request_slots="other"'), + ): "1.0", + ( + "sglang:prefill_execution_tokens_total", + ('bucket="256"', 'kind="scheduled"', 'path="cuda_graph"'), + ): "129.0", + ( + "sglang:prefill_execution_tokens_total", + ('bucket="256"', 'kind="executed"', 'path="cuda_graph"'), + ): "256.0", + } + for (name, labels), value in expected.items(): + with self.subTest(metric=name, labels=labels): + self.assertEqual( + _metric_line(metrics, name, *labels).split()[-1], + value, + ) + + self.assertIn( + 'sglang:hicache_backup_tokens_total{cache_type="test",rank="0"} 512.0', + metrics, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/utils/test_subprocess_watchdog.py b/test/registered/unit/utils/test_subprocess_watchdog.py index d486a2d48cb6..3791ac11666d 100644 --- a/test/registered/unit/utils/test_subprocess_watchdog.py +++ b/test/registered/unit/utils/test_subprocess_watchdog.py @@ -134,6 +134,38 @@ def test_normal_exit_no_sigquit(self): "SIGQUIT should not be triggered for normal exit (exitcode=0)", ) + def test_wait_for_clean_exit_accepts_successful_children(self): + proc = self._spawn(noop_worker) + monitor = self._watch(proc, names=["scheduler"]) + + monitor.wait_for_clean_exit(timeout=2) + + self.assertEqual(proc.exitcode, 0) + self.assertFalse(self.sigquit_triggered.is_set()) + + def test_wait_for_clean_exit_reports_failed_children(self): + proc = self._spawn(crashing_worker) + proc.join(timeout=2) + monitor = self._watch(proc, names=["scheduler"]) + + with self.assertRaisesRegex( + RuntimeError, + r"scheduler\(pid=\d+, exitcode=1\)", + ): + monitor.wait_for_clean_exit(timeout=2) + + def test_wait_for_clean_exit_reports_timeout(self): + proc = self._spawn(healthy_worker) + monitor = self._watch(proc, names=["scheduler"]) + + with self.assertRaisesRegex(RuntimeError, "scheduler.*timed out"): + monitor.wait_for_clean_exit(timeout=0.01) + + def test_wait_for_clean_exit_rejects_nonpositive_timeout(self): + monitor = self._watch([]) + with self.assertRaisesRegex(ValueError, "timeout must be positive"): + monitor.wait_for_clean_exit(timeout=0) + if __name__ == "__main__": import unittest