supports differnt layer setting

Author: gmlwns2000

python/sglang/srt/layers/attention/hip_attention/hip_config.py

@@ -1,5 +1,6 @@
 from dataclasses import dataclass, field, InitVar
 from typing import List, Optional, Union
+import warnings
 
 from hip.models.hip_attention.gen3.attention_metadata import ScanStage
 
@@ -79,6 +80,10 @@ class HiPAttentionConfig:
         HiPAttentionPerLayerConfig(parsed_json={"second_stage_k": 4096, "sliding_window_size": 1024, "sink_token_size": 256}),
         HiPAttentionPerLayerConfig(),
     ])
+    prefill_layers: list[HiPAttentionPerLayerConfig] = field(default_factory=lambda: [
+        HiPAttentionPerLayerConfig(parsed_json={"second_stage_k": 4096, "sliding_window_size": 1024, "sink_token_size": 256}),
+        HiPAttentionPerLayerConfig(),
+    ])
 
     # deprecated
     apply_v_dot: bool = False
@@ -129,9 +134,16 @@ def __post_init__(self, parsed_json: dict | None):
                     HiPAttentionPerLayerConfig(parsed_json=layer)
                     for layer in parsed_json['layers']
                 ]
+                self.prefill_layers = self.layers
                 parsed_json.pop('layers')
+            if 'prefill_layers' in parsed_json:
+                self.prefill_layers = [
+                    HiPAttentionPerLayerConfig(parsed_json=layer)
+                    for layer in parsed_json['prefill_layers']
+                ]
+                parsed_json.pop('prefill_layers')
             if parsed_json:
-                raise ValueError(f'Unknown keys in json: {parsed_json.keys()}')
+                raise Exception(f'Unknown keys in json: {parsed_json.keys()}')
 
         num_stages = len(self.layers[0].stages)
         for layer_config in self.layers:

python/sglang/srt/layers/attention/hip_attention/hip_radix_attention.py

@@ -528,16 +528,22 @@ def forward_paged_hip(
 
         online_update_cache: bool = False,
     ) -> tuple[torch.Tensor, "HiPAttentionOutputMetadata"]:
-        is_dense = layer.layer_id in self.hip_config.dense_layers
-        
-        if len(self.hip_config.layers) == 2:
-            layer_config = self.hip_config.layers[0 if is_dense else 1]
-        else:
-            layer_config = self.hip_config.layers[layer.layer_id]
-
         N, num_heads, hidden_dims = query.shape
         dst_seq_len = N // batch_size
 
+        is_decode = dst_seq_len == 1
+        is_dense = layer.layer_id in self.hip_config.dense_layers
+        if not is_decode:
+            if len(self.hip_config.prefill_layers) == 2:
+                layer_config = self.hip_config.prefill_layers[0 if is_dense else 1]
+            else:
+                layer_config = self.hip_config.prefill_layers[layer.layer_id]
+        else:
+            if len(self.hip_config.layers) == 2:
+                layer_config = self.hip_config.layers[0 if is_dense else 1]
+            else:
+                layer_config = self.hip_config.layers[layer.layer_id]
+
         query = query.view(batch_size, dst_seq_len, num_heads, hidden_dims)
 
         if k_cache is not None: