[Minor] Code style improvements (#2355)

python/sglang/srt/layers/torchao_utils.py

@@ -2,12 +2,10 @@
 Common utilities for torchao.
 """
 
-from typing import Dict, Set
-
 import torch
 
 
-def apply_torchao_config_to_model_(
+def apply_torchao_config_to_model(
     model: torch.nn.Module, torchao_config: str, filter_fn=None
 ):
     """Quantize a modelwith torchao quantization specified by torchao_config
@@ -21,6 +19,7 @@ def apply_torchao_config_to_model_(
     # Lazy import to suppress some warnings
     from torchao.quantization import (
         float8_dynamic_activation_float8_weight,
+        float8_weight_only,
         int4_weight_only,
         int8_dynamic_activation_int8_weight,
         int8_weight_only,
@@ -28,6 +27,11 @@ def apply_torchao_config_to_model_(
     )
     from torchao.quantization.observer import PerRow, PerTensor
 
+    if filter_fn is None:
+
+        def filter_fn(module, fqn):
+            return "proj" in fqn
+
     if torchao_config == "" or torchao_config is None:
         return model
     elif "int8wo" in torchao_config:
@@ -44,8 +48,6 @@ def apply_torchao_config_to_model_(
         ], f"int4wo groupsize needs to be one of [32, 64, 128, 256] but got {group_size}"
         quantize_(model, int4_weight_only(group_size=group_size), filter_fn=filter_fn)
     elif "fp8wo" in torchao_config:
-        from torchao.quantization import float8_weight_only
-
         # this requires newer hardware
         # [rank0]: AssertionError: fp8e4nv data type is not supported on CUDA arch < 89
         quantize_(model, float8_weight_only(), filter_fn=filter_fn)