diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 61d8391d9..6756be8c0 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -594,6 +594,7 @@ class Envs: False, deprecated_name="SGLANG_NSA_HIP_DISABLE_PRESHUFFLE" ) SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2) + SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False) SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True) SGLANG_DSA_TOPK_BROADCAST = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 55f1905e3..3cd24d7a2 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py @@ -12,16 +12,24 @@ from sglang.jit_kernel.fused_store_index_cache import ( can_use_dsa_fused_store, fused_store_index_k_cache, ) +from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.environ import envs from sglang.srt.layers.attention.dsa.utils import ( aiter_can_use_preshuffle_paged_mqa, is_dsa_enable_prefill_cp, is_dsa_prefill_cp_in_seq_split, + is_graph_dsa_split_op_surface, ) from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor from sglang.srt.layers.layernorm import LayerNorm from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype, is_fp8_fnuz from sglang.srt.layers.utils import MultiPlatformOp +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + eager_on_graph, +) +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( get_tc_piecewise_forward_context, is_in_tc_piecewise_cuda_graph, @@ -39,6 +47,7 @@ from sglang.srt.utils import ( is_hip, is_npu, ) +from sglang.srt.utils.custom_op import register_custom_op logger = logging.getLogger(__name__) @@ -97,6 +106,16 @@ if TYPE_CHECKING: DUAL_STREAM_TOKEN_THRESHOLD = 1024 if _is_cuda else 0 +GRAPH_WEIGHTS_PROJ_LORA_ERROR = ( + "DSA indexer weights_proj LoRA is incompatible with " + "piecewise/breakable CUDA graph; remove the explicit " + "prefill cuda-graph backend override or drop " + "indexer.weights_proj from the LoRA target modules." +) + + +def _is_in_piecewise_or_breakable_cuda_graph() -> bool: + return is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph() def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool: @@ -128,52 +147,8 @@ def _uses_dsa_attention_backend(forward_batch: ForwardBatch) -> bool: if _is_cuda: - from sglang.srt.compilation.compilation_config import register_split_op - from sglang.srt.utils.custom_op import register_custom_op - @register_custom_op(mutates_args=["topk_result"]) - @register_split_op() - def k_cache_and_topk_result( - layer_id: int, - key: torch.Tensor, - q_fp8: torch.Tensor, - weights: torch.Tensor, - topk_result: torch.Tensor, - ) -> None: - assert ( - _is_cuda - ), "Internal error: piecewise CUDA graph is only supported on CUDA" - from sglang.srt.layers.attention.dsa.triton_kernel import act_quant - - forward_batch = get_tc_piecewise_forward_context().forward_batch - indexer = get_tc_piecewise_forward_context().dsa_indexers[layer_id] - metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch) - assert metadata is not None, ( - "DSA piecewise CUDA graph requires indexer metadata from the DSA " - "attention backend" - ) - - # slice off padding from piecewise CUDA graph - extend_num_tokens = forward_batch.extend_num_tokens - - indexer._store_index_k_cache( - forward_batch=forward_batch, - layer_id=layer_id, - key=key[:extend_num_tokens], - act_quant=act_quant, - out_cache_loc=forward_batch.out_cache_loc[:extend_num_tokens], - ) - indexer._get_topk_ragged( - False, - forward_batch, - layer_id, - q_fp8[:extend_num_tokens], - weights, - metadata, - topk_result, - ) - - def _logits_head_gate_pcg_fake_impl( + def _logits_head_gate_graph_fake_impl( x: torch.Tensor, weight: torch.Tensor, n_heads_inv_sqrt: float, @@ -186,8 +161,9 @@ if _is_cuda: device=x.device, ) - @register_custom_op(fake_impl=_logits_head_gate_pcg_fake_impl) - def logits_head_gate_pcg( + # In-graph (PCG/BCG) head gate for the NON-prefill path + @register_custom_op(fake_impl=_logits_head_gate_graph_fake_impl) + def logits_head_gate_graph( x: torch.Tensor, weight: torch.Tensor, n_heads_inv_sqrt: float, @@ -467,6 +443,15 @@ class Indexer(MultiPlatformOp): ): return weights.unsqueeze(-1) * q_scale * self.softmax_scale + def _should_skip_logits_computation(self, forward_batch: ForwardBatch) -> bool: + if ( + forward_batch.forward_mode.is_extend_without_speculative() + and forward_batch.seq_lens_cpu is not None + ): + max_kv_len = forward_batch.seq_lens_cpu.max().item() + return max_kv_len <= self.index_topk + return False + def _get_q_k_bf16( self, q_lora: torch.Tensor, @@ -986,21 +971,36 @@ class Indexer(MultiPlatformOp): enable_dual_stream: bool, metadata: BaseIndexerMetadata, return_indices: bool = True, + *, + num_tokens: Optional[int] = None, + topk_result: Optional[torch.Tensor] = None, ) -> Optional[torch.Tensor]: + # Shared by the eager path and the graph DSA split-op dispatch. The two + # keyword args carry the graph contract and default to the eager behavior: + # - num_tokens: slice key/out_cache_loc to the unpadded count (the graph + # runs at a static padded shape). None => full (eager) shape. + # - topk_result: pre-allocated padded buffer to fill in place (a downstream + # captured graph reads it at a fixed address). None => return a fresh, + # naturally-sized tensor. assert forward_batch.forward_mode.is_extend_without_speculative() x_meta = x[0] if isinstance(x, tuple) else x # Fast path: only compute and store k cache, skip all q and weights ops key = self._get_k_bf16(x, positions, enable_dual_stream) - - if not forward_batch.out_cache_loc.is_contiguous(): + out_cache_loc = None + if num_tokens is not None: + assert num_tokens <= key.shape[0] + assert num_tokens <= forward_batch.out_cache_loc.shape[0] + key = key[:num_tokens] + out_cache_loc = forward_batch.out_cache_loc[:num_tokens] + elif not forward_batch.out_cache_loc.is_contiguous(): forward_batch.out_cache_loc = forward_batch.out_cache_loc.contiguous() - self._store_index_k_cache( forward_batch=forward_batch, layer_id=layer_id, key=key, act_quant=act_quant, + out_cache_loc=out_cache_loc, ) # MHA doesn't need topk_indices @@ -1016,7 +1016,13 @@ class Indexer(MultiPlatformOp): dtype=torch.float32, device=x_meta.device, ) - return metadata.topk_transform(dummy_logits, self.index_topk) + raw_topk_result = metadata.topk_transform(dummy_logits, self.index_topk) + if topk_result is not None: + # PCG/BCG: fill the valid prefix of the padded static buffer and + # leave padded rows at the -1 sentinel. + topk_result[: raw_topk_result.shape[0]] = raw_topk_result + return None + return raw_topk_result def _get_topk_ragged_with_cp( self, @@ -1029,9 +1035,10 @@ class Indexer(MultiPlatformOp): actual_seq_q: int, cp_index: List[Tuple[int, int, int]] = None, ) -> torch.Tensor: - assert ( - not is_in_tc_piecewise_cuda_graph() - ), "DSA context parallel (_get_topk_ragged_with_cp) not supported under piecewise CUDA graph" + assert not _is_in_piecewise_or_breakable_cuda_graph(), ( + "DSA context parallel (_get_topk_ragged_with_cp) not supported under " + "piecewise/breakable CUDA graph" + ) if TYPE_CHECKING: assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool) @@ -1178,9 +1185,10 @@ class Indexer(MultiPlatformOp): topk: int, layer_id: int, ) -> Optional[torch.Tensor]: - assert ( - not is_in_tc_piecewise_cuda_graph() - ), "DSA forward_indexer (non-CUDA loop path) not supported under piecewise CUDA graph" + assert not _is_in_piecewise_or_breakable_cuda_graph(), ( + "DSA forward_indexer (non-CUDA loop path) not supported under " + "piecewise/breakable CUDA graph" + ) if not _is_npu: from sglang.srt.layers.attention.dsa.tilelang_kernel import fp8_index @@ -1370,10 +1378,15 @@ class Indexer(MultiPlatformOp): # a tuple like (x_fp8, x_scale[, y]). Use `x_meta` for shape/device queries. x_meta = x[0] if isinstance(x, tuple) else x - # In piecewise CUDA graph mode, metadata is fetched inside custom ops via get_tc_piecewise_forward_context() to - # prevent Dynamo from guarding on forward_metadata identity (which changes each - # replay when init_forward_metadata creates a new ForwardMetadata object). - if not is_in_tc_piecewise_cuda_graph(): + in_piecewise_or_breakable_cuda_graph = ( + _is_in_piecewise_or_breakable_cuda_graph() + ) + + # In piecewise/breakable CUDA graph mode, metadata is fetched inside + # custom ops via get_tc_piecewise_forward_context() to prevent Dynamo + # from guarding on forward_metadata identity, which changes each replay + # when init_forward_metadata creates a new ForwardMetadata object. + if not in_piecewise_or_breakable_cuda_graph: metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch) if metadata is None: return None @@ -1390,13 +1403,10 @@ class Indexer(MultiPlatformOp): # Determine if should skip topk based on sequence length # We can only skip the logits computation if cuda graph is not involved skip_logits_computation = False - if ( - not is_in_tc_piecewise_cuda_graph() - and forward_batch.forward_mode.is_extend_without_speculative() - ): - if forward_batch.seq_lens_cpu is not None: - max_kv_len = forward_batch.seq_lens_cpu.max().item() - skip_logits_computation = max_kv_len <= self.index_topk + if not in_piecewise_or_breakable_cuda_graph: + skip_logits_computation = self._should_skip_logits_computation( + forward_batch + ) # Optimization: fast path when skipping topk computation if skip_logits_computation and (not self.dsa_enable_prefill_cp): @@ -1417,6 +1427,43 @@ class Indexer(MultiPlatformOp): # wrapper owns base+delta and no LoRA kernel runs under torch.compile weights_proj_lora = getattr(self.weights_proj, "set_lora", False) + if ( + is_graph_dsa_split_op_surface(forward_batch) + and not self.dsa_enable_prefill_cp + ): + # Default path for non-CP prefill under PCG/BCG: run the whole indexer + # (q/k proj, head gate, k-cache store, topk) as a single eager split op + # instead of capturing it piecemeal in the graph. + if weights_proj_lora: + raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR) + if return_indices: + topk_result = torch.full( + (x.shape[0], self.index_topk), + -1, + device=x.device, + dtype=torch.int32, + ) + else: + topk_result = torch.empty( + (0, self.index_topk), device=x.device, dtype=torch.int32 + ) + graph_dispatch_fn = ( + bcg_dsa_indexer_prefill_split + if is_in_breakable_cuda_graph() + else pcg_dsa_indexer_prefill_split + ) + graph_dispatch_fn( + layer_id=layer_id, + x=x, + q_lora=q_lora, + positions=positions, + topk_result=topk_result, + ) + result = _broadcast_indexer_topk_from_rank0( + topk_result if return_indices else None + ) + return maybe_capture_indexer_topk(layer_id, result) + if enable_dual_stream and forward_batch.forward_mode.is_decode_or_idle(): current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) @@ -1455,7 +1502,7 @@ class Indexer(MultiPlatformOp): act_quant=act_quant, ) current_stream.wait_stream(self.alt_stream) - elif not is_in_tc_piecewise_cuda_graph(): + elif not in_piecewise_or_breakable_cuda_graph: q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) self._store_index_k_cache( forward_batch=forward_batch, @@ -1464,8 +1511,10 @@ class Indexer(MultiPlatformOp): act_quant=act_quant, ) else: - # piecewise CUDA graph need to split graph on store_k_cache and mqa_logits, - # so delay store_k_cache after weights proj. + # Graph paths not handled by the full DSA indexer split op + # still need q_fp8 for paged topk and q_scale for + # logits_head_gate_graph. K-cache storage is handled by the + # full graph split path when prefill requires it. q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) # aiter (ROCm gfx95): the 3-tuple (fp8, scale, bf16) from @@ -1509,13 +1558,10 @@ class Indexer(MultiPlatformOp): else: x_for_gate = x - if is_in_tc_piecewise_cuda_graph(): + if in_piecewise_or_breakable_cuda_graph: if weights_proj_lora: - raise RuntimeError( - "DSA indexer weights_proj LoRA is incompatible with TC piecewise CUDA graph; remove the explicit" - " prefill cuda-graph backend override or drop indexer.weights_proj from the LoRA target modules." - ) - weights = logits_head_gate_pcg( + raise RuntimeError(GRAPH_WEIGHTS_PROJ_LORA_ERROR) + weights = logits_head_gate_graph( x_for_gate, self.weights_proj.weight, self.n_heads**-0.5, @@ -1529,9 +1575,10 @@ class Indexer(MultiPlatformOp): weights = self._get_logits_head_gate(x_for_gate, q_scale) if _is_cuda or _is_hip: - # In piecewise CUDA graph, any access to seq_lens_cpu creates a Dynamo shape guard. - # Piecewise CUDA graph never has empty batches. - if not is_in_tc_piecewise_cuda_graph(): + # In piecewise/breakable CUDA graph, any access to seq_lens_cpu + # creates a Dynamo shape guard. These graph modes never have empty + # batches. + if not in_piecewise_or_breakable_cuda_graph: assert forward_batch.seq_lens_cpu is not None if len(forward_batch.seq_lens_cpu) == 0: # this seems b/c max-pad, no worries? @@ -1602,28 +1649,14 @@ class Indexer(MultiPlatformOp): topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0) topk_result = _broadcast_indexer_topk_from_rank0(topk_result) return maybe_capture_indexer_topk(layer_id, topk_result) - elif is_in_tc_piecewise_cuda_graph(): - assert ( - not enable_dual_stream - ), "Internal error: piecewise CUDA graph should not be enabled with dual stream" - - if not _uses_dsa_attention_backend(forward_batch): - return None - - topk_result = torch.full( - (q_fp8.shape[0], self.index_topk), - -1, - device=q_fp8.device, - dtype=torch.int32, - ) - k_cache_and_topk_result( - layer_id=layer_id, - key=key, - q_fp8=q_fp8, - weights=weights, - topk_result=topk_result, - ) else: + # In-graph (PCG/BCG) non-CP prefill is handled earlier by the + # graph DSA split-op dispatch, so only the eager path reaches + # here. + assert not in_piecewise_or_breakable_cuda_graph, ( + "Internal error: in-graph DSA prefill must go through the " + "graph DSA split-op dispatch" + ) topk_result = self._get_topk_ragged( enable_dual_stream, forward_batch, @@ -1974,6 +2007,88 @@ class Indexer(MultiPlatformOp): return topk_indices_prev[0], topk_indices_next[0] +@register_custom_op(mutates_args=["topk_result"]) +@register_split_op() +def pcg_dsa_indexer_prefill_split( + layer_id: int, + x: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + topk_result: torch.Tensor, +) -> None: + # Default in-graph indexer path for non-CP prefill: runs the whole indexer + # (q/k proj, head gate, k-cache store, topk) as one eager split op. PCG calls + # this as a split op; BCG uses the explicit eager wrapper below. + # + # Output contract (differs from the eager `forward` path): a split op returns + # None, so results are delivered only by mutating `topk_result` in place. The + # call site pre-allocates it at a static, padded shape and a downstream + # captured graph reads it at a fixed address; eager code instead allocates + # and returns a fresh, naturally-sized tensor each call. + assert _is_cuda, "Internal error: DSA graph dispatch is only supported on CUDA" + from sglang.srt.layers.attention.dsa.triton_kernel import act_quant + + forward_context = get_tc_piecewise_forward_context() + forward_batch = forward_context.forward_batch + indexer = forward_context.dsa_indexers[layer_id] + metadata = get_attn_backend().get_indexer_metadata(layer_id, forward_batch) + + extend_num_tokens = forward_batch.extend_num_tokens + # Empty buffer encodes return_indices=False for graph dispatch. + return_indices = topk_result.numel() != 0 + k_only = not return_indices or ( + indexer._should_skip_logits_computation(forward_batch) + and not indexer.dsa_enable_prefill_cp + ) + if k_only: + indexer._forward_cuda_k_only( + x, + positions, + forward_batch, + layer_id, + act_quant, + enable_dual_stream=False, + metadata=metadata, + return_indices=return_indices, + num_tokens=extend_num_tokens, + topk_result=topk_result, + ) + return + + query, key = indexer._get_q_k_bf16( + q_lora, + x, + positions, + enable_dual_stream=False, + forward_batch=forward_batch, + ) + q_fp8, q_scale = act_quant(query, indexer.block_size, indexer.scale_fmt) + # Reuse the compiled head-gate util shared with the eager path. + weights = indexer._get_logits_head_gate(x, q_scale) + # Store K cache + ragged top-k, sliced to the unpadded count and writing into + # the static padded topk_result buffer (the graph contract). Mirrors the eager + # path's store + _get_topk_ragged. + indexer._store_index_k_cache( + forward_batch=forward_batch, + layer_id=layer_id, + key=key[:extend_num_tokens], + act_quant=act_quant, + out_cache_loc=forward_batch.out_cache_loc[:extend_num_tokens], + ) + indexer._get_topk_ragged( + False, + forward_batch, + layer_id, + q_fp8[:extend_num_tokens], + weights, + metadata, + topk_result, + ) + + +bcg_dsa_indexer_prefill_split = eager_on_graph(True)(pcg_dsa_indexer_prefill_split) + + def scattered_to_tp_attn_full( hidden_states: torch.Tensor, forward_batch, diff --git a/python/sglang/srt/layers/attention/dsa/utils.py b/python/sglang/srt/layers/attention/dsa/utils.py index 506a299ff..8d0e67651 100644 --- a/python/sglang/srt/layers/attention/dsa/utils.py +++ b/python/sglang/srt/layers/attention/dsa/utils.py @@ -9,9 +9,15 @@ from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import ( DpPaddingMode, ) +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + is_in_breakable_cuda_graph, +) +from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + is_in_tc_piecewise_cuda_graph, +) from sglang.srt.runtime_context import get_parallel from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import get_bool_env_var, is_hip +from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip from sglang.srt.utils.common import ceil_align, ceil_div @@ -80,6 +86,19 @@ def is_dsa_prefill_cp_round_robin_split(): ) +# Structural surface where the graph DSA split-op dispatch (DSA indexer) and the +# MLA BMM-into-attention fusion apply: a non-speculative extend (prefill) running +# inside a piecewise/breakable CUDA graph. Both fusions are now on by default on +# this surface (no feature flag); each adds its own extra carve-outs at its call +# site (e.g. the indexer also excludes DSA prefill context parallelism). +def is_graph_dsa_split_op_surface(forward_batch: "ForwardBatch") -> bool: + return ( + is_cuda() + and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph()) + and forward_batch.forward_mode.is_extend_without_speculative() + ) + + def can_dsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"): if not forward_batch.forward_mode.is_context_parallel_extend(): return False diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 4cea05806..eb48fd0ac 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -2471,14 +2471,18 @@ class DeepseekSparseAttnBackend( """ Decide all attention prefill dispatch strategies for this batch. """ + from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, + ) from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) from sglang.srt.utils import get_device_sm, is_blackwell # Decide MHA vs MLA - if is_in_tc_piecewise_cuda_graph(): - # Can't branch on seq_lens_cpu in PCG, force mha off to guarantee correctness. + if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph(): + # Can't branch on seq_lens_cpu in graph replay, force MHA off to + # guarantee correctness. self.use_mha = False elif ( forward_batch and forward_batch.forward_mode.is_extend_without_speculative() diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 840260576..b782c0d0d 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -1,18 +1,24 @@ from __future__ import annotations +from dataclasses import dataclass from typing import TYPE_CHECKING, Optional import torch +from sglang.srt.compilation.compilation_config import register_split_op from sglang.srt.environ import envs from sglang.srt.layers import deep_gemm_wrapper -from sglang.srt.layers.attention.dsa.utils import dsa_use_prefill_cp +from sglang.srt.layers.attention.dsa.utils import ( + dsa_use_prefill_cp, + is_graph_dsa_split_op_surface, +) from sglang.srt.layers.communicator import get_attn_tp_context from sglang.srt.layers.quantization.fp8_kernel import ( fp8_dtype, per_tensor_quant_mla_fp8, per_token_group_quant_mla_deep_gemm_masked_fp8, ) +from sglang.srt.layers.radix_attention import unified_attention_with_output from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp from sglang.srt.lora.deepseek_mla_correction import ( apply_q_correction as apply_kv_b_lora_q_correction, @@ -28,6 +34,12 @@ from sglang.srt.model_executor.forward_context import ( get_attn_backend, get_token_to_kv_pool, ) +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( + eager_on_graph, +) +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( is_in_tc_piecewise_cuda_graph, ) @@ -48,17 +60,25 @@ from sglang.srt.state_capturer.indexer_topk import ( maybe_capture_indexer_topk, ) from sglang.srt.utils import BumpAllocator +from sglang.srt.utils.custom_op import register_custom_op _SGLANG_EXPERIMENTAL_LORA_OPTI = envs.SGLANG_EXPERIMENTAL_LORA_OPTI.get() if TYPE_CHECKING: from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA + +@dataclass(frozen=True) +class MlaBmmFusionPlan: + q_nope_t: torch.Tensor + q_nope_out_buf: torch.Tensor + q_nope_out_view: torch.Tensor + attn_output_buf: torch.Tensor + + if _is_cuda: from sgl_kernel import bmm_fp8 as _raw_bmm_fp8 - from sglang.srt.utils.custom_op import register_custom_op - # TODO(yuwei): remove this wrapper after sgl-kernel registers its own fake/meta impl # Wrap bmm_fp8 as a custom op so torch.compile does not trace into # torch.cuda.current_blas_handle() (which returns a non-Tensor). @@ -140,6 +160,63 @@ class DeepseekMLAForwardMixin: get_global_server_args().flashinfer_mla_disable_ragged ) + def _can_fuse_bmm_into_attention( + self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch + ) -> bool: + # Shared activation surface with the DSA indexer graph dispatch + # (in piecewise/breakable graph + non-speculative extend). Like the indexer + # dispatch, this fusion is on by default on that surface. + if not is_graph_dsa_split_op_surface(forward_batch): + return False + if not self.use_dsa: + return False + if self.use_deep_gemm_bmm or _is_hip: + return False + if is_kv_b_lora_active(self): + return False + # The isolated 1-kernel graph is the bf16 fallback BMM. The fp8 and + # DeepGEMM branches already use different fused paths. + if self.w_kc.dtype == torch.float8_e4m3fn: + return False + if self.current_attention_backend not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS: + return False + return True + + def _split_q_nope_pe( + self: DeepseekV2AttentionMLA, + q: torch.Tensor, + latent_cache: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1) + return q_nope, q_pe, k_pe + + def _make_mla_bmm_fusion_plan( + self: DeepseekV2AttentionMLA, + q: torch.Tensor, + q_nope: torch.Tensor, + ) -> MlaBmmFusionPlan: + q_nope_out_buf = q.new_empty( + ( + self.num_local_heads, + q.shape[0], + self.kv_lora_rank, + ) + ) + q_nope_out_view = q_nope_out_buf.transpose(0, 1) + attn_output_buf = q.new_empty( + ( + q.shape[0], + self.num_local_heads * self.kv_lora_rank, + ) + ) + return MlaBmmFusionPlan( + q_nope_t=q_nope.transpose(0, 1), + q_nope_out_buf=q_nope_out_buf, + q_nope_out_view=q_nope_out_view, + attn_output_buf=attn_output_buf, + ) + def forward_absorb_prepare( self: DeepseekV2AttentionMLA, positions: torch.Tensor, @@ -151,8 +228,16 @@ class DeepseekMLAForwardMixin: ): from sglang.srt.model_executor.runner import get_is_capture_mode + fuse_bmm_attention = ( + self.q_lora_rank is not None + and self._can_fuse_bmm_into_attention(forward_batch) + ) q_lora = None topk_indices = None + q_nope = None + q_pe = None + k_pe = None + fusion_plan: Optional[MlaBmmFusionPlan] = None if self.q_lora_rank is not None: q, latent_cache = ( get_attn_tp_context() @@ -276,6 +361,13 @@ class DeepseekMLAForwardMixin: else: k_nope = k_nope.unsqueeze(1) q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) + + # Hoist these above the DSA indexer split op so the indexer + # and the composite bmm+attention split op are adjacent in FX. + if fuse_bmm_attention: + q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache) + fusion_plan = self._make_mla_bmm_fusion_plan(q, q_nope) + if q_lora is not None: # See the skip_topk note above: shared layers have no # indexer weights, so this gate must not fall back to @@ -302,109 +394,123 @@ class DeepseekMLAForwardMixin: k_nope = latent_cache[..., : self.kv_lora_rank] k_nope = self.kv_a_layernorm(k_nope).unsqueeze(1) - q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) - k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1) + if q_nope is None: + q_nope, q_pe, k_pe = self._split_q_nope_pe(q, latent_cache) _kvb_q = None - if _SGLANG_EXPERIMENTAL_LORA_OPTI: - # Fork the kv_b q-correction A-step onto the LoRA side stream to overlap the bmm. - from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import ( - kv_b_lora_q_prepare, - ) - - _kvb_q = kv_b_lora_q_prepare(self, q_nope) - - if self.use_deep_gemm_bmm: - ( - q_nope_val, - q_nope_scale, - masked_m, - expected_m, - aligned_m, - ) = per_token_group_quant_mla_deep_gemm_masked_fp8(q_nope.transpose(0, 1)) - q_nope_out = q_nope.new_empty( - (self.num_local_heads, aligned_m, self.kv_lora_rank) - ) - deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( - (q_nope_val, q_nope_scale), - (self.w_kc, self.w_scale_k), - q_nope_out, - masked_m, - expected_m, - ) - q_nope_out = q_nope_out[:, :expected_m, :] - elif _is_hip: - # TODO(haishaw): add bmm_fp8 to ROCm - if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8: - x = q_nope.transpose(0, 1) - q_nope_out = torch.empty( - x.shape[0], - x.shape[1], - self.w_kc.shape[2], - device=x.device, - dtype=torch.bfloat16, + if fusion_plan is not None: + # The composite split op fills q_nope_out_buf and attention reads + # this transposed alias directly. + q_nope_out = fusion_plan.q_nope_out_view + else: + if _SGLANG_EXPERIMENTAL_LORA_OPTI: + # Fork the kv_b q-correction A-step onto the LoRA side stream to overlap the bmm. + from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import ( + kv_b_lora_q_prepare, ) - batched_gemm_afp4wfp4_pre_quant( - x, - self.w_kc.transpose(-2, -1), - self.w_scale_k.transpose(-2, -1), - torch.bfloat16, + + _kvb_q = kv_b_lora_q_prepare(self, q_nope) + + if self.use_deep_gemm_bmm: + ( + q_nope_val, + q_nope_scale, + masked_m, + expected_m, + aligned_m, + ) = per_token_group_quant_mla_deep_gemm_masked_fp8( + q_nope.transpose(0, 1) + ) + q_nope_out = q_nope.new_empty( + (self.num_local_heads, aligned_m, self.kv_lora_rank) + ) + deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked( + (q_nope_val, q_nope_scale), + (self.w_kc, self.w_scale_k), q_nope_out, + masked_m, + expected_m, ) - else: - if (_use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn) or ( - get_is_capture_mode() and self.w_kc.dtype == torch.float8_e4m3fnuz - ): - # fp8 Triton kernel: always on gfx950, - # cudagraph-only on gfx942 (hides launch overhead) - q_nope_out = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( - X=q_nope, - WQ=self.w_kc.transpose(-1, -2), - w_scale=self.w_scale, - group_size=128, - YQ=None, # allocate (B, M, N) - transpose_bm=False, # (B, M, N) - transpose_bm_in=True, # (M, B, K) + q_nope_out = q_nope_out[:, :expected_m, :] + elif _is_hip: + # TODO(haishaw): add bmm_fp8 to ROCm + if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8: + x = q_nope.transpose(0, 1) + q_nope_out = torch.empty( + x.shape[0], + x.shape[1], + self.w_kc.shape[2], + device=x.device, dtype=torch.bfloat16, ) - + batched_gemm_afp4wfp4_pre_quant( + x, + self.w_kc.transpose(-2, -1), + self.w_scale_k.transpose(-2, -1), + torch.bfloat16, + q_nope_out, + ) else: + if ( + _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn + ) or ( + get_is_capture_mode() + and self.w_kc.dtype == torch.float8_e4m3fnuz + ): + # fp8 Triton kernel: always on gfx950, + # cudagraph-only on gfx942 (hides launch overhead) + q_nope_out = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant( + X=q_nope, + WQ=self.w_kc.transpose(-1, -2), + w_scale=self.w_scale, + group_size=128, + YQ=None, # allocate (B, M, N) + transpose_bm=False, # (B, M, N) + transpose_bm_in=True, # (M, B, K) + dtype=torch.bfloat16, + ) + + else: + q_nope_out = torch.bmm( + q_nope.to(torch.bfloat16).transpose(0, 1), + self.w_kc.to(torch.bfloat16) * self.w_scale, + ) + + elif self.w_kc.dtype == torch.float8_e4m3fn: + if _is_cpu: q_nope_out = torch.bmm( q_nope.to(torch.bfloat16).transpose(0, 1), self.w_kc.to(torch.bfloat16) * self.w_scale, ) - - elif self.w_kc.dtype == torch.float8_e4m3fn: - if _is_cpu: - q_nope_out = torch.bmm( - q_nope.to(torch.bfloat16).transpose(0, 1), - self.w_kc.to(torch.bfloat16) * self.w_scale, - ) + else: + # fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612 + q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( + q_nope.transpose(0, 1), + ( + torch.zeros((1,), dtype=torch.float32, device=q_nope.device) + if _is_cublas_ge_129 + else zero_allocator.allocate(1) + ), + ) + q_nope_out = bmm_fp8( + q_nope_val, + self.w_kc, + q_nope_scale, + self.w_scale, + torch.bfloat16, + ) else: - # fix bmm_fp8 error under cublas12.9 caused by bumpallocator, detail in pr#11612 - q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8( - q_nope.transpose(0, 1), - ( - torch.zeros((1,), dtype=torch.float32, device=q_nope.device) - if _is_cublas_ge_129 - else zero_allocator.allocate(1) - ), - ) - q_nope_out = bmm_fp8( - q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16 - ) - else: - q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) + q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc) - q_nope_out = q_nope_out.transpose(0, 1) - if _SGLANG_EXPERIMENTAL_LORA_OPTI: - from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import ( - kv_b_lora_q_apply, - ) + q_nope_out = q_nope_out.transpose(0, 1) + if _SGLANG_EXPERIMENTAL_LORA_OPTI: + from sglang.srt.lora.trtllm_lora_temp.deepseek_mla_correction import ( + kv_b_lora_q_apply, + ) - q_nope_out = kv_b_lora_q_apply(self, q_nope, q_nope_out, _kvb_q) - elif is_kv_b_lora_active(self): - q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out) + q_nope_out = kv_b_lora_q_apply(self, q_nope, q_nope_out, _kvb_q) + elif is_kv_b_lora_active(self): + q_nope_out = apply_kv_b_lora_q_correction(self, q_nope, q_nope_out) skip_rope_for_dsa_tilelang_fused = self._skip_rope_for_dsa_tilelang_fused() skip_rope_for_aiter_fused_mla = self._skip_rope_for_aiter_fused_mla() @@ -433,6 +539,7 @@ class DeepseekMLAForwardMixin: positions, topk_indices, llama_4_scaling, + fusion_plan, ) def forward_absorb_core( @@ -446,6 +553,7 @@ class DeepseekMLAForwardMixin: positions, topk_indices, llama_4_scaling, + fusion_plan: Optional[MlaBmmFusionPlan] = None, ): save_kv_cache = True @@ -527,20 +635,44 @@ class DeepseekMLAForwardMixin: "is_neox": self.rotary_emb.is_neox_style, "llama_4_scaling": llama_4_scaling, } - attn_output = self.attn_mqa( - q_nope_out, - k_nope, - k_nope, - forward_batch, - q_rope=q_pe, - k_rope=k_pe, - **extra_args, - **( - dict(topk_indices=topk_indices) - if topk_indices is not None - else {} - ), - ) + if fusion_plan is not None: + bmm_attention_fn = ( + bcg_mla_bmm_then_unified_attention + if is_in_breakable_cuda_graph() + else mla_bmm_then_unified_attention + ) + bmm_attention_fn( + fusion_plan.q_nope_t, + self.w_kc, + fusion_plan.q_nope_out_buf, + q_nope_out, + k_nope, + fusion_plan.attn_output_buf, + save_kv_cache, + self.layer_id, + q_pe, + k_pe, + cos_sin_cache=extra_args.get("cos_sin_cache"), + is_neox=extra_args.get("is_neox"), + llama_4_scaling=extra_args.get("llama_4_scaling"), + topk_indices=topk_indices, + ) + attn_output = fusion_plan.attn_output_buf + else: + attn_output = self.attn_mqa( + q_nope_out, + k_nope, + k_nope, + forward_batch, + q_rope=q_pe, + k_rope=k_pe, + **extra_args, + **( + dict(topk_indices=topk_indices) + if topk_indices is not None + else {} + ), + ) else: if _use_aiter_gfx95: cos = self.rotary_emb.cos_cache @@ -807,3 +939,54 @@ class DeepseekMLAForwardMixin: and self.current_attention_backend not in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS ) + + +# Fuses the absorb BMM (`q_nope @ w_kc`) with `unified_attention_with_output` +# into one eager split op under both PCG and BCG. Without this, the bf16 +# fallback BMM is captured alone in its own single-kernel CUDA graph submodule, +# paying per-submodule host overhead with no fusion benefit. +# +# `q_nope_out_view` aliases `q_nope_out_buf` (transposed). The op writes +# `q_nope_out_buf` via `torch.bmm(..., out=...)` and then reads through +# `q_nope_out_view`, so the alias's storage is mutated too. Declare it in +# `mutates_args` to keep the schema honest. +@register_custom_op( + mutates_args=["q_nope_out_buf", "q_nope_out_view", "attn_output_buf"] +) +@register_split_op() +def mla_bmm_then_unified_attention( + q_nope_t: torch.Tensor, + w_kc: torch.Tensor, + q_nope_out_buf: torch.Tensor, + q_nope_out_view: torch.Tensor, + k_nope: torch.Tensor, + attn_output_buf: torch.Tensor, + save_kv_cache: bool, + layer_id: int, + q_pe: torch.Tensor, + k_pe: torch.Tensor, + cos_sin_cache: Optional[torch.Tensor] = None, + is_neox: Optional[bool] = None, + llama_4_scaling: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, +) -> None: + torch.bmm(q_nope_t, w_kc, out=q_nope_out_buf) + unified_attention_with_output( + q_nope_out_view, + k_nope, + k_nope, + attn_output_buf, + save_kv_cache, + layer_id, + q_rope=q_pe, + k_rope=k_pe, + cos_sin_cache=cos_sin_cache, + is_neox=is_neox, + llama_4_scaling=llama_4_scaling, + topk_indices=topk_indices, + ) + + +bcg_mla_bmm_then_unified_attention = eager_on_graph(True)( + mla_bmm_then_unified_attention +) diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 69814b2bf..46e162e9a 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -94,7 +94,7 @@ from sglang.srt.layers.moe.token_dispatcher.base import ( CombineInput, DispatchOutput, ) -from sglang.srt.layers.moe.topk import TopK, TopKOutputFormat +from sglang.srt.layers.moe.topk import BypassedTopKOutput, TopK, TopKOutputFormat from sglang.srt.layers.moe.utils import ( RoutingMethodType, filter_moe_weight_param_global_expert, @@ -135,6 +135,13 @@ from sglang.srt.model_executor.cuda_graph_config import ( from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_context import get_attn_backend from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import ( + is_in_breakable_cuda_graph, +) +from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import ( + get_tc_piecewise_forward_context, + is_in_tc_piecewise_cuda_graph, +) from sglang.srt.models.deepseek_common.attention_backend_handler import ( AttentionBackendRegistry, ) @@ -207,6 +214,10 @@ else: logger = logging.getLogger(__name__) +_enable_pcg_dsv2_dual_stream = ( + _is_cuda and envs.SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM.get() +) + class DeepseekV2MLP(nn.Module): def __init__( @@ -821,6 +832,26 @@ class DeepseekV2MoE(nn.Module): ) ] + def _can_dual_stream_graph( + self, hidden_states: torch.Tensor, server_args=None + ) -> bool: + if server_args is None: + server_args = get_global_server_args() + return ( + _enable_pcg_dsv2_dual_stream + and (is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph()) + and get_moe_runner_backend().is_flashinfer_trtllm() + and self.alt_stream is not None + and self.num_fused_shared_experts == 0 + and hidden_states.shape[0] > 0 + and hasattr(self, "shared_experts") + and getattr(self.experts, "use_flashinfer_trtllm_moe", False) + and not self._enable_a2a_moe + and not self._fuse_shared_experts_inside_sbo + and not getattr(self, "is_hash", False) + and not server_args.enable_eplb + ) + def forward( self, hidden_states: torch.Tensor, @@ -843,16 +874,24 @@ class DeepseekV2MoE(nn.Module): ) if not self._enable_a2a_moe: - if ( + server_args = get_global_server_args() + if self._can_dual_stream_graph(hidden_states, server_args): + return dsv2_flashinfer_moe_dual_stream_graph( + hidden_states, + self.layer_id, + should_allreduce_fusion, + use_reduce_scatter, + ) + elif ( self.alt_stream is not None and self.num_fused_shared_experts == 0 and hidden_states.shape[0] > 0 and get_is_capture_mode() and not ( - get_global_server_args().enable_torch_compile + server_args.enable_torch_compile and hidden_states.shape[0] - <= get_global_server_args().torch_compile_max_bs - * (get_global_server_args().speculative_num_draft_tokens or 1) + <= server_args.torch_compile_max_bs + * (server_args.speculative_num_draft_tokens or 1) ) ): return self.forward_normal_dual_stream( @@ -886,6 +925,8 @@ class DeepseekV2MoE(nn.Module): gemm_output_zero_allocator: BumpAllocator = None, input_ids: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None, + *, + use_flashinfer_trtllm_bypass: bool = False, ) -> torch.Tensor: current_stream = torch.cuda.current_stream() self.alt_stream.wait_stream(current_stream) @@ -901,17 +942,24 @@ class DeepseekV2MoE(nn.Module): with torch.cuda.stream(self.alt_stream): # router_logits: (num_tokens, n_experts) router_logits = self.gate(hidden_states, gemm_output_zero_allocator) - topk_kwargs = ( - {"input_ids": input_ids_global} - if getattr(self, "is_hash", False) - else {} - ) - topk_output = self.topk( - hidden_states, - router_logits, - expert_location_dispatch_info=dispatch_info, - **topk_kwargs, - ) + if use_flashinfer_trtllm_bypass: + topk_output = BypassedTopKOutput( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk.topk_config, + ) + else: + topk_kwargs = ( + {"input_ids": input_ids_global} + if getattr(self, "is_hash", False) + else {} + ) + topk_output = self.topk( + hidden_states, + router_logits, + expert_location_dispatch_info=dispatch_info, + **topk_kwargs, + ) deferred_finalize = ( shared_output is not None and not self._shared_expert_tp1 @@ -922,6 +970,10 @@ class DeepseekV2MoE(nn.Module): final_hidden_states = self.experts.forward_deferred_finalize( hidden_states, topk_output ) + elif use_flashinfer_trtllm_bypass: + final_hidden_states = self.experts.forward_impl( + hidden_states, topk_output + ) else: final_hidden_states = self.experts(hidden_states, topk_output) if ( @@ -2864,4 +2916,25 @@ def flashinfer_dsv3_router_gemm( ) +@register_custom_op(out_shape="hidden_states") +def dsv2_flashinfer_moe_dual_stream_graph( + hidden_states: torch.Tensor, + layer_id: int, + should_allreduce_fusion: bool, + use_reduce_scatter: bool, +) -> torch.Tensor: + forward_context = get_tc_piecewise_forward_context() + assert forward_context is not None + assert forward_context.moe_fusions is not None + + moe_fusion = forward_context.moe_fusions[layer_id] + assert moe_fusion is not None + return moe_fusion.forward_normal_dual_stream( + hidden_states, + should_allreduce_fusion=should_allreduce_fusion, + use_reduce_scatter=use_reduce_scatter, + use_flashinfer_trtllm_bypass=True, + ) + + EntryClass = [DeepseekV2ForCausalLM, DeepseekV3ForCausalLM, DeepseekV32ForCausalLM] diff --git a/test/registered/cuda_graph/piecewise/test_pcg_glm5_fp8_tp8.py b/test/registered/cuda_graph/piecewise/test_pcg_glm5_fp8_tp8.py new file mode 100644 index 000000000..73c7c79c4 --- /dev/null +++ b/test/registered/cuda_graph/piecewise/test_pcg_glm5_fp8_tp8.py @@ -0,0 +1,75 @@ +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.run_eval import run_eval +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + +register_cuda_ci(est_time=900, stage="base-c", runner_config="8-gpu-h200") + +GLM5_FP8_MODEL = "zai-org/GLM-5-FP8" + + +class TestBCGGlm5Fp8TP8(CustomTestCase): + """Breakable CUDA graph prefill on GLM-5-FP8 (DSA model, TP=8, H200).""" + + @classmethod + def setUpClass(cls): + cls.model = GLM5_FP8_MODEL + cls.base_url = DEFAULT_URL_FOR_TEST + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=[ + "--tp-size", + "8", + "--trust-remote-code", + "--reasoning-parser", + "glm45", + "--tool-call-parser", + "glm47", + "--mem-fraction-static", + "0.8", + "--disable-flashinfer-autotune", + "--cuda-graph-backend-prefill=breakable", + # Small chunks => many prefill iterations, each <= the 2048 + # capture max, so every prefill batch replays the BCG graph and + # exercises the DSA split-op / dual-stream / MLA-fusion paths. + "--chunked-prefill-size", + "512", + "--model-loader-extra-config", + '{"enable_multithread_load": true, "num_threads": 64}', + ], + env={ + "SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM": "1", + }, + ) + + @classmethod + def tearDownClass(cls): + if hasattr(cls, "process") and cls.process: + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + args = SimpleNamespace( + base_url=self.base_url, + model=self.model, + eval_name="gsm8k", + num_examples=200, + num_threads=200, + max_tokens=4096, + ) + metrics = run_eval(args) + print(f"{metrics=}") + self.assertGreater(metrics["score"], 0.92) + + +if __name__ == "__main__": + unittest.main()