refactor

Author: yundai424

python/sglang/srt/layers/attention/flashinfer_backend.py

@@ -6,7 +6,7 @@
 FlashInfer is faster and Triton is easier to customize.
 Each backend supports two operators: extend (i.e. prefill with cached prefix) and decode.
 """
-import logging
+
 import os
 from dataclasses import dataclass
 from enum import Enum, auto
@@ -37,7 +37,7 @@
     from flashinfer.cascade import merge_state
     from flashinfer.decode import _get_range_buf, get_seq_lens
 
-logger = logging.getLogger(__name__)
+
 class WrapperDispatch(Enum):
     SLIDING_WINDOW = auto()
     CROSS_ATTENTION = auto()
@@ -83,6 +83,7 @@ def __init__(
         self.max_context_len = model_runner.model_config.context_len
         self.skip_prefill = skip_prefill
         self.is_multimodal = model_runner.model_config.is_multimodal
+        self.kv_cache_dtype = model_runner.kv_cache_dtype
 
         assert not (
             model_runner.sliding_window_size is not None
@@ -392,8 +393,8 @@ def forward_extend(
         forward_batch: ForwardBatch,
         save_kv_cache=True,
     ):
-        k_scale = layer.k_scale_float if layer.kv_cache_dtype != "auto" else None
-        v_scale = layer.v_scale_float if layer.kv_cache_dtype != "auto" else None
+        k_scale = layer.k_scale_float if self.kv_cache_dtype != "auto" else None
+        v_scale = layer.v_scale_float if self.kv_cache_dtype != "auto" else None
         prefill_wrapper_paged = self.forward_metadata.prefill_wrappers[
             self._get_wrapper_idx(layer)
         ]
@@ -462,8 +463,8 @@ def forward_decode(
         forward_batch: ForwardBatch,
         save_kv_cache=True,
     ):
-        k_scale = layer.k_scale_float if layer.kv_cache_dtype != "auto" else None
-        v_scale = layer.v_scale_float if layer.kv_cache_dtype != "auto" else None
+        k_scale = layer.k_scale_float if self.kv_cache_dtype != "auto" else None
+        v_scale = layer.v_scale_float if self.kv_cache_dtype != "auto" else None
         decode_wrapper = self.forward_metadata.decode_wrappers[
             self._get_wrapper_idx(layer)
         ]

python/sglang/srt/layers/radix_attention.py

@@ -56,7 +56,6 @@ def __init__(
         self.v_scale = None
         self.k_scale_float = None
         self.v_scale_float = None
-        self.kv_cache_dtype = "auto"
         if quant_config is not None:
             self.quant_method = quant_config.get_quant_method(self, prefix=prefix)
 

python/sglang/srt/model_executor/model_runner.py

@@ -388,12 +388,6 @@ def load_model(self):
         monkey_patch_vllm_parallel_state(reverse=True)
         monkey_patch_isinstance_for_vllm_base_layer(reverse=True)
 
-        # Set KV cache dtype for RadixAttention if the model uses it
-        if hasattr(self.model, "set_kv_cache_dtype"):
-            self.model.set_kv_cache_dtype(self.server_args.kv_cache_dtype)
-            logger.info(
-                f"Set KV cache dtype to {self.server_args.kv_cache_dtype} for {type(self.model).__name__}"
-            )
         if self.server_args.kv_cache_dtype == "fp8_e4m3":
             if self.server_args.quantization_param_path is not None:
                 if callable(getattr(self.model, "load_kv_cache_scales", None)):

python/sglang/srt/model_loader/loader.py

@@ -373,7 +373,6 @@ def load_model(
             for _, module in model.named_modules():
                 quant_method = getattr(module, "quant_method", None)
                 if quant_method is not None:
-                    logger.warning(f"{module.__class__.__name__}, {quant_method}")
                     # When quant methods need to process weights after loading
                     # (for repacking, quantizing, etc), they expect parameters
                     # to be on the global target device. This scope is for the

python/sglang/srt/models/llama.py

@@ -185,8 +185,6 @@ def forward(
         output, _ = self.o_proj(attn_output)
         return output
 
-    def set_kv_cache_dtype(self, kv_cache_dtype: str):
-        self.attn.kv_cache_dtype = kv_cache_dtype
 
 class LlamaDecoderLayer(nn.Module):
     def __init__(
@@ -238,9 +236,6 @@ def __init__(
         self.post_attention_layernorm = RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
-    
-    def set_kv_cache_dtype(self, kv_cache_dtype: str):
-        self.self_attn.set_kv_cache_dtype(kv_cache_dtype)
 
     def forward(
         self,
@@ -628,11 +623,6 @@ def set_eagle3_layers_to_capture(self):
         self.capture_aux_hidden_states = True
         num_layers = self.config.num_hidden_layers
         self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
-    
-    def set_kv_cache_dtype(self, kv_cache_dtype: str):
-        for layer in self.model.layers:
-            if hasattr(layer, "set_kv_cache_dtype"):
-                layer.set_kv_cache_dtype(kv_cache_dtype)
 
 
 class Phi3ForCausalLM(LlamaForCausalLM):

test/srt/test_eval_fp8_accuracy.py

@@ -137,6 +137,9 @@ def _run_test(self, model, other_args, expected_score):
         finally:
             kill_process_tree(process.pid)
 
+    @unittest.skipIf(
+        torch.version.hip is not None, "modelopt quantization unsupported on ROCm"
+    )
     def test_mmlu_offline_only(self):
         """Test with offline quantization only."""
         self._run_test(