CI: Kill zombie diffusion processes in CI & minor code style fix on rotary embedding fallback (#13637)

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"

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,

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.

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)

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