update mla to obtain the optimal configuration from config

vllm/attention/backends/triton_mla.py

@@ -37,7 +37,18 @@ from vllm.utils import async_tensor_h2d, make_tensor_with_pad
 if TYPE_CHECKING:
     from vllm.worker.model_runner import (ModelInputForGPUBuilder,
                                           ModelInputForGPUWithSamplingMetadata)
-        
+   
+
+def get_config(bs_key, mean_kv_seqlen_key, config):
+    # 转换参数为字符串以匹配字典的键
+    bs_key_str = str(bs_key)
+    mean_kv_seqlen_key_str = str(mean_kv_seqlen_key)
+    
+    # 检查字典中是否存在对应的配置
+    if bs_key_str in config and mean_kv_seqlen_key_str in config[bs_key_str]:
+        return config[bs_key_str][mean_kv_seqlen_key_str]
+    else:
+        raise ValueError(f"No matching configuration found for bs key: {bs_key} and mean kv seq key: {mean_kv_seqlen_key} when init decode attention db")     
                          
 def get_mla_config_file_name(QH: int, KVH: int, QKD: int, VD: int, cache_dtype: Optional[str]) -> str:
     if cache_dtype == "default":
@@ -735,6 +746,8 @@ class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
                                       "encoder/decoder cross-attention "
                                       "are not implemented for "
                                       "TritonMLAImpl")
+            
+        self.attn_configs = get_attention_mla_configs(self.num_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim, self.kv_lora_rank, "fp16")
         
     def _forward_prefill(
         self,
@@ -789,13 +802,16 @@ class TritonMLAImpl(MLACommonImpl[TritonMLAMetadata]):
         kv_c_cache = kv_c_and_k_pe_cache[..., :self.kv_lora_rank]
         PAGE_SIZE = kv_c_and_k_pe_cache.size(1)
         
-        # config = get_attention_mla_configs(self.num_heads, 1, self.kv_lora_rank + self.qk_rope_head_dim, self.kv_lora_rank, "fp16")
-
+        # TODO
+        for bs in self.attn_configs.keys():
+            for mean_seq_len in self.attn_configs[bs].keys():
+                best_config = get_config(bs, mean_seq_len, self.attn_configs)
+                
         # Run MQA
         decode_attention_fwd(q, kv_c_and_k_pe_cache, kv_c_cache, o,
                              decode_meta.block_tables,
                              decode_meta.seq_lens_tensor, attn_logits,
-                             attn_metadata.num_kv_splits, self.scale, # config, 
+                             attn_metadata.num_kv_splits, self.scale, best_config, 
                              PAGE_SIZE)
 
         return self._v_up_proj_and_o_proj(o)

vllm/model_executor/model_loader/utils.py

@@ -80,7 +80,7 @@ def get_model_architecture(
     architectures = getattr(model_config.hf_config, "architectures", [])
     visions = getattr(model_config.hf_config, "visual", []) or getattr(model_config.hf_config, "vision_config", [])
     # TODO: support deepseek distillation series models ( 'LlamaForCausalLM', 'Qwen2ForCausalLM' )
-    support_nn_architectures = ['QWenLMHeadModel', 'Qwen2VLForConditionalGeneration', 'Qwen2_5_VLForConditionalGeneration', 
+    support_nn_architectures = ['LlamaForCausalLM', 'Qwen2ForCausalLM', 'QWenLMHeadModel', 'Qwen2VLForConditionalGeneration', 'Qwen2_5_VLForConditionalGeneration', 
                                 'Qwen2MoeForCausalLM', 'ChatGLMModel', 'ChatGLMForConditionalGeneration', 
                                 'BaichuanForCausalLM', 'BloomForCausalLM', 'MedusaModel', 'MixtralForCausalLM', 
                                 'MLPSpeculatorPreTrainedModel', 'FalconForCausalLM', 'DeepseekV2ForCausalLM',