[Bug_fix] fix torchair o_proj forward parameter #4166
Conversation
Signed-off-by: zzhx1 <[email protected]>
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![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 fix a parameter mismatch for the o_proj forward pass by adding the is_force_scatter parameter to AscendRowParallelLinear.forward. While the parameter is correctly added to the method signature, it is not passed down to the custom operator's apply method. This makes the fix incomplete and the new parameter will have no effect. I've left a comment with a suggested code change to properly pass the parameter. Please note that this will also require further changes in vllm_ascend/ops/linear_op.py to update the signatures of the apply and apply_impl methods.
| is_force_scatter: bool = False, | ||
| ) -> Union[torch.Tensor, tuple[torch.Tensor, Optional[Parameter]]]: | ||
| if self.custom_op is not None: | ||
| return self.custom_op.apply(input_) |
There was a problem hiding this comment.
The newly added is_force_scatter parameter and the existing is_prefill parameter are not being passed to self.custom_op.apply(). This makes the fix incomplete as the custom op will not receive these values. You should pass them to the apply method. This will also require updating the signature of apply and apply_impl methods in CustomLinearOp and its subclasses in vllm_ascend/ops/linear_op.py.
| return self.custom_op.apply(input_) | |
| return self.custom_op.apply(input_, is_prefill=is_prefill, is_force_scatter=is_force_scatter) |
1 participant
What this PR does / why we need it?
In
torchair_mla.py, theself.oprojfunction includes an additional parameteris_force_scatter, while theAscendRowParallelLinearfunction inlinear.pydoes not add this parameter.