[P/D][Bugfix] Use config to get kv_num_head. #6405
Conversation
Signed-off-by: wangxiaoteng <[email protected]>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request aims to correctly configure num_kv_head for Vision-Language (VL) models. However, there is a critical bug in the implementation where self.num_kv_head is unconditionally set to 1, which negates the intended logic. My review includes a comment with a suggested fix for this issue.
| try: | ||
| if is_vl_model: | ||
| self.num_kv_head = self.vllm_config.model_config.hf_config.text_config.num_key_value_heads | ||
| else: | ||
| self.num_kv_head = self.vllm_config.model_config.hf_config.num_key_value_heads | ||
| except Exception as e: | ||
| logger.error( | ||
| f"MooncakeLayerwise Connector can not get num_key_value_heads from model config, set num_key_value_heads to 1" | ||
| ) | ||
| self.num_kv_head = 1 |
There was a problem hiding this comment.
There is a critical logic error in how self.num_kv_head is being set. The line self.num_kv_head = 1 is outside the try...except block, which means it will always execute and overwrite any value successfully retrieved from the model configuration. This effectively hardcodes self.num_kv_head to 1 for all cases, defeating the purpose of this change.
The assignment to the default value of 1 should only happen within the except block when fetching the configuration fails. I've also updated the error log to include the exception for better debugging.
| try: | |
| if is_vl_model: | |
| self.num_kv_head = self.vllm_config.model_config.hf_config.text_config.num_key_value_heads | |
| else: | |
| self.num_kv_head = self.vllm_config.model_config.hf_config.num_key_value_heads | |
| except Exception as e: | |
| logger.error( | |
| f"MooncakeLayerwise Connector can not get num_key_value_heads from model config, set num_key_value_heads to 1" | |
| ) | |
| self.num_kv_head = 1 | |
| try: | |
| if is_vl_model: | |
| self.num_kv_head = self.vllm_config.model_config.hf_config.text_config.num_key_value_heads | |
| else: | |
| self.num_kv_head = self.vllm_config.model_config.hf_config.num_key_value_heads | |
| except Exception as e: | |
| logger.error( | |
| f"MooncakeLayerwise Connector can not get num_key_value_heads from model config, set num_key_value_heads to 1. Error: {e}" | |
| ) | |
| self.num_kv_head = 1 |
wangxiaoteng888
changed the title
Signed-off-by: wangxiaoteng <[email protected]>
| self.use_mla = self.vllm_config.model_config.use_mla | ||
| try: | ||
| if is_vl_model: | ||
| self.num_kv_head = self.vllm_config.model_config.hf_config.text_config.num_key_value_heads |
There was a problem hiding this comment.
Can we use the hf_config.text_config.num_key_value_heads access method regardless of whether it's a vl model or not?
| logger.error( | ||
| f"MooncakeLayerwise Connector can not get num_key_value_heads from model config, set num_key_value_heads to 1" | ||
| ) | ||
| self.num_kv_head = 1 |
There was a problem hiding this comment.
What does this line of code do?
2 participants
What this PR does / why we need it?
This pull request aims to correctly configure num_kv_head for Vision-Language (VL) models.
Does this PR introduce any user-facing change?
No
How was this patch tested?
By ci