Fix attention backends for models with per-layer head counts (num_attention_heads_per_layer) (#32625)

python/sglang/test/kits/attention_unittest/attention_methods/dense_attention.py

@@ -295,6 +295,9 @@ class TinyModelConfig:
         assert self.num_attention_heads % tp_size == 0
         return self.num_attention_heads // tp_size
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         assert self.num_key_value_heads % tp_size == 0
         return self.num_key_value_heads // tp_size

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -263,6 +263,9 @@ class TinyDSAModelConfig:
         self.hf_text_config = self.hf_config
         self.linear_attn_registry_result = None
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
 
 class DSAMockModelRunner(ModelRunner):
     def __init__(

python/sglang/test/kits/attention_unittest/attention_methods/dsv4_attention.py

@@ -287,6 +287,9 @@ class TinyDSV4ModelConfig:
         self.hf_text_config = self.hf_config
         self.linear_attn_registry_result = None
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
 
 class MockDSV4ModelRunner:
     """Minimal runner exposing what `DeepseekV4AttnBackend.__init__` reads.

python/sglang/test/kits/attention_unittest/attention_methods/dual_chunk_attention.py

@@ -295,6 +295,9 @@ class TinyDualChunkModelConfig:
         assert self.num_attention_heads % tp_size == 0
         return self.num_attention_heads // tp_size
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         assert self.num_key_value_heads % tp_size == 0
         return self.num_key_value_heads // tp_size

python/sglang/test/kits/attention_unittest/attention_methods/gdn_attention.py

@@ -188,6 +188,9 @@ class TinyGDNModelConfig:
         self.hf_text_config = self.hf_config
         self.linear_attn_registry_result = None
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         assert self.num_key_value_heads % tp_size == 0
         return self.num_key_value_heads // tp_size

python/sglang/test/kits/attention_unittest/attention_methods/kda_attention.py

@@ -193,6 +193,9 @@ class TinyKDAModelConfig:
         self.hf_text_config = self.hf_config
         self.linear_attn_registry_result = None
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         assert self.num_key_value_heads % tp_size == 0
         return self.num_key_value_heads // tp_size

python/sglang/test/kits/attention_unittest/attention_methods/lightning_attention.py

@@ -203,6 +203,9 @@ class TinyLightningModelConfig:
         self.hf_text_config = self.hf_config
         self.linear_attn_registry_result = None
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         assert self.num_key_value_heads % tp_size == 0
         return self.num_key_value_heads // tp_size

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -289,6 +289,9 @@ class TinyMamba2ModelConfig:
         self.hf_text_config = self.hf_config
         self.linear_attn_registry_result = None
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         assert self.num_key_value_heads % tp_size == 0
         return self.num_key_value_heads // tp_size

python/sglang/test/kits/attention_unittest/attention_methods/mla_attention.py

@@ -200,6 +200,9 @@ class TinyMLAModelConfig:
         assert self.num_attention_heads % tp_size == 0
         return self.num_attention_heads // tp_size
 
+    def get_max_num_attention_heads(self) -> int:
+        return self.num_attention_heads
+
     def get_num_kv_heads(self, tp_size: int) -> int:
         return 1