From a352e833c44115fcacf3955da31e455a750c528e Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Thu, 20 Nov 2025 05:57:01 -0800 Subject: [PATCH] CI: Kill zombie diffusion processes in CI & minor code style fix on rotary embedding fallback (#13637) --- .../runtime/managers/gpu_worker.py | 2 +- python/sglang/srt/layers/rotary_embedding.py | 9 ++++--- .../srt/layers/vocab_parallel_embedding.py | 4 ++- python/sglang/srt/metrics/label_transform.py | 27 +++++++++++++++++++ scripts/killall_sglang.sh | 4 +-- 5 files changed, 39 insertions(+), 7 deletions(-) create mode 100644 python/sglang/srt/metrics/label_transform.py diff --git a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py index ddf9dbc78..8ab64af0f 100644 --- a/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py +++ b/python/sglang/multimodal_gen/runtime/managers/gpu_worker.py @@ -62,7 +62,7 @@ class GPUWorker: def init_device_and_model(self) -> None: """Initialize the device and load the model.""" - setproctitle(f"sgl_diffusion::scheduler:{self.local_rank}") + setproctitle(f"sgl_diffusion::scheduler_TP{self.local_rank}") torch.cuda.set_device(self.local_rank) # Set environment variables for distributed initialization os.environ["MASTER_ADDR"] = "localhost" diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index dfddc1d67..8678aaef3 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -120,7 +120,10 @@ class RotaryEmbedding(CustomOp): ): from vllm._custom_ops import rotary_embedding - self.vllm_rotary_embedding = rotary_embedding + self.use_fallback_kernel = True + self.fallback_rotary_embedding = rotary_embedding + else: + self.use_fallback_kernel = False self.cos_sin_cache: torch.Tensor self.register_buffer("cos_sin_cache", cache, persistent=False) @@ -273,7 +276,7 @@ class RotaryEmbedding(CustomOp): offsets: Optional[torch.Tensor] = None, fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - if _is_cuda and (self.head_size in [64, 128, 256, 512]): + if not self.use_fallback_kernel: apply_rope_with_cos_sin_cache_inplace( positions=positions, query=query, @@ -293,7 +296,7 @@ class RotaryEmbedding(CustomOp): fused_set_kv_buffer_arg is None ), "save kv cache is not supported for vllm_rotary_embedding." self.cos_sin_cache = self.cos_sin_cache.to(query.device, dtype=query.dtype) - self.vllm_rotary_embedding( + self.fallback_rotary_embedding( positions, query, key, diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py index 5912cb86d..f171a15f2 100644 --- a/python/sglang/srt/layers/vocab_parallel_embedding.py +++ b/python/sglang/srt/layers/vocab_parallel_embedding.py @@ -473,11 +473,13 @@ class VocabParallelEmbedding(torch.nn.Module): ) else: masked_input = input_ + # Get the embeddings. with use_symmetric_memory(get_tp_group(), disabled=not self.enable_tp): output_parallel = self.quant_method.embedding(self, masked_input.long()) - # Mask the output embedding. + if self.tp_size > 1: + # Mask the output embedding. output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) if not get_attn_tp_context().input_scattered: # Reduce across all the model parallel GPUs. diff --git a/python/sglang/srt/metrics/label_transform.py b/python/sglang/srt/metrics/label_transform.py new file mode 100644 index 000000000..6d2d9b161 --- /dev/null +++ b/python/sglang/srt/metrics/label_transform.py @@ -0,0 +1,27 @@ +from typing import Optional + +_PRIORITY_MIN = 0 +_PRIORITY_MAX = 31 +_LOW_PRIORITY_VALUE = "LOW" +_HIGH_PRIORITY_VALUE = "HIGH" + +UNKNOWN_PRIORITY_VALUE = "UNKNOWN" + + +def transform_priority(priority: Optional[int]) -> str: + """Transform the priority to a string for metrics reporting. + Limit the range to prevent high cardinality issues. + + Args: + priority: The priority to transform. + Returns: + The transformed priority. + """ + if priority is None: + return UNKNOWN_PRIORITY_VALUE + elif priority < _PRIORITY_MIN: + return _LOW_PRIORITY_VALUE + elif priority >= _PRIORITY_MAX: + return _HIGH_PRIORITY_VALUE + else: + return str(priority) diff --git a/scripts/killall_sglang.sh b/scripts/killall_sglang.sh index 7d0fe8bca..e637aaee9 100755 --- a/scripts/killall_sglang.sh +++ b/scripts/killall_sglang.sh @@ -4,14 +4,14 @@ if [ "$1" = "rocm" ]; then echo "Running in ROCm mode" # Clean SGLang processes - pgrep -f 'sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt' | xargs -r kill -9 + pgrep -f 'sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt|sgl_diffusion::' | xargs -r kill -9 else # Show current GPU status nvidia-smi # Clean SGLang processes - pgrep -f 'sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt' | xargs -r kill -9 + pgrep -f 'sglang::|sglang\.launch_server|sglang\.bench|sglang\.data_parallel|sglang\.srt|sgl_diffusion::' | xargs -r kill -9 # Clean all GPU processes if any argument is provided if [ $# -gt 0 ]; then