From 6c7cb903655d4b8523e45838e597c11e10a6600f Mon Sep 17 00:00:00 2001
From: Lianmin Zheng <lianminzheng@gmail.com>
Date: Tue, 10 Sep 2024 11:27:03 -0700
Subject: [PATCH] [Minor] improve kill scripts and torchao import (#1375)

---
 python/sglang/srt/layers/torchao_utils.py | 14 ++++++++------
 test/killall_sglang.sh                    |  2 +-
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/python/sglang/srt/layers/torchao_utils.py b/python/sglang/srt/layers/torchao_utils.py
index 16eb1f2c5..bc7bde86e 100644
--- a/python/sglang/srt/layers/torchao_utils.py
+++ b/python/sglang/srt/layers/torchao_utils.py
@@ -3,15 +3,17 @@ Common utilities for torchao.
 """
 
 import torch
-from torchao.quantization import (
-    int4_weight_only,
-    int8_dynamic_activation_int8_weight,
-    int8_weight_only,
-    quantize_,
-)
 
 
 def torchao_quantize_param_data(param, torchao_config):
+    # Lazy import to suppress some warnings
+    from torchao.quantization import (
+        int4_weight_only,
+        int8_dynamic_activation_int8_weight,
+        int8_weight_only,
+        quantize_,
+    )
+
     dummy_linear = torch.nn.Linear(param.shape[1], param.shape[0], bias=False)
     dummy_linear.weight = param
     if "int8wo" in torchao_config:
diff --git a/test/killall_sglang.sh b/test/killall_sglang.sh
index 0e2cb82a8..c536548d4 100644
--- a/test/killall_sglang.sh
+++ b/test/killall_sglang.sh
@@ -1 +1 @@
-kill -9 $(ps aux | grep 'sglang' | grep -v 'grep' | awk '{print $2}')
+kill -9 $(ps aux | grep 'sglang.launch_server' | grep -v 'grep' | awk '{print $2}')