[Bugfix] Clean up some cruft in mamba.py #9343
Conversation
|
👋 Hi! Thank you for contributing to the vLLM project. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging. To run CI, PR reviewers can do one of these:
🚀 |
github-actions
bot
added
the
ready
|
Ok, this was extremely boneheaded -- just removed an entire feed forwards section that I had copy-pasted from Jamba, that isn't present in the Mamba model. |
vllm/model_executor/models/mamba.py
Outdated
| @@ -332,40 +282,20 @@ def forward( | |||
| current_ssm_state = ssm_state[i] | |||
| current_conv_state = conv_state[i] | |||
|
|
|||
| hidden_states, residual = layer( | |||
| hidden_states = layer( | |||
| positions=positions, | |||
| hidden_states=hidden_states, | |||
| attn_metadata=attn_metadata, | |||
| residual=residual, | |||
There was a problem hiding this comment.
Residual is still passed in here. I personally think it is fine to keep the previous structure of passing residual into rmsnorm, but up to you
There was a problem hiding this comment.
I saw modeling_mamba.py do it this way and like that it's super clean, but you're right -- going to revert to the previous way so that we fuse. Sidenote: this is the kind of optimization we should be relying on torch.compile to do for us IMO
Delete some unnecessary/useless code I noticed while reviewing #7478