-
-
Notifications
You must be signed in to change notification settings - Fork 11k
[Model] Clean up and simplify Mamba2 Metadata Usage in both V0 and V1 #24331
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Model] Clean up and simplify Mamba2 Metadata Usage in both V0 and V1 #24331
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request refactors the Mamba2 metadata handling to unify the logic for V0 and V1 APIs. The changes primarily involve renaming metadata fields to be consistent across both versions, which simplifies the code in mamba_mixer2 and plamo2 by removing duplicated logic. The change to how has_initial_states_p is calculated is also a good improvement for correctness and clarity. I've found one critical issue where attribute assignments were not updated after renaming, which could lead to a runtime error.
…code Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
ed9781f to
1c9fc08
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
…vllm-project#24331) Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
…vllm-project#24331) Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>
…vllm-project#24331) Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
…vllm-project#24331) Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com> Signed-off-by: xuebwang-amd <xuebwang@amd.com>
Purpose
This PR cleans up some naming inconsistencies in mamba2 metadata between V0 and V1. This affects models that use mamba_mixer2 or plamo2
Test Plan
Run e2e lm_eval on
ibm-granite/granite-4.0-tiny-previewandpfnet/plamo-2.1-2b-cptfor both V0 and V1Test Result
They passed and no significant accuracy differences observed.
granite-4-tiny-preview V1 (main 006e7a3)
granite-4-tiny-preview V1 (this PR)
granite-4-tiny-preview V0 (this PR)
Plamo2 V1 (main 006e7a3)
Plamo2 V1 (this PR)
Plamo2 V0 (this PR)
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.