[Kernel][Model] Improve continuous batching for Jamba and Mamba

[mzusman]

• Mamba kernels now avoid writing/reading to/from padded cache entries
• Jamba now uses the previously implemented functionality for continuous batching in [Kernel] Change interface to Mamba causal_conv1d_update for continuous batching #8012 [Kernel] Change interface to Mamba selective_state_update for continuous batching #8039 [Kernel][Model] Varlen prefill + Prefill chunking support for mamba kernels and Jamba model #8533

CC @tlrmchlsmth

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

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:

• Add ready label to the PR
• Enable auto-merge.

🚀

[tlrmchlsmth]

One suggestion: To mirror the PAD_SLOT_ID that's in backends/utils.py, we could #define PAD_SLOT_ID -1 in csrc/attention_generic.cuh, so that we name the constant.

And then we can add comments letting the reader know to keep them in sync

[mzusman]

One suggestion: To mirror the PAD_SLOT_ID that's in backends/utils.py, we could #define PAD_SLOT_ID -1 in csrc/attention_generic.cuh, so that we name the constant.

And then we can add comments letting the reader know to keep them in sync

I've taken another approach, passing the pad_slot_id through the kernel params, keeps it simpler since we don't need to keep the cpp and python variables in sync and keeps the mamba kernels as standalone entities. WDYT?

[mzusman]

Pushed adaptations for #6484 , The PR is ready again.

[tlrmchlsmth]

Changing the title since Jamba and Mamba already support continuous batching -- this just makes it better

[tlrmchlsmth]

One suggestion: To mirror the PAD_SLOT_ID that's in backends/utils.py, we could #define PAD_SLOT_ID -1 in csrc/attention_generic.cuh, so that we name the constant.
And then we can add comments letting the reader know to keep them in sync

I've taken another approach, passing the pad_slot_id through the kernel params, keeps it simpler since we don't need to keep the cpp and python variables in sync and keeps the mamba kernels as standalone entities. WDYT?

Sounds good to me

[tlrmchlsmth]

I tried out the PR and am currently seeing the following error:

E           RuntimeError: Error in model execution (input dumped to /tmp/err_execute_model_input_20241015-011501.pkl): _C::causal_conv1d_fwd() expected at most 8 argument(s) but received 9 argument(s). Declaration: _C::causal_conv1d_fwd(Tensor($0! -> ) x, Tensor($1! -> ) weight, Tensor? bias_, Tensor($2! -> )? conv_states, Tensor? query_start_loc, Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation) -> Tensor

If you're seeing the same issue, could you LMK when it's fixed? Looks good otherwise, thanks!

[mzusman]

I tried out the PR and am currently seeing the following error:

E           RuntimeError: Error in model execution (input dumped to /tmp/err_execute_model_input_20241015-011501.pkl): _C::causal_conv1d_fwd() expected at most 8 argument(s) but received 9 argument(s). Declaration: _C::causal_conv1d_fwd(Tensor($0! -> ) x, Tensor($1! -> ) weight, Tensor? bias_, Tensor($2! -> )? conv_states, Tensor? query_start_loc, Tensor? cache_indices, Tensor? has_initial_state, bool silu_activation) -> Tensor

If you're seeing the same issue, could you LMK when it's fixed? Looks good otherwise, thanks!

@tlrmchlsmth I cannot reproduce this error, could you let me know how did you run into it?

[tlrmchlsmth]

@tlrmchlsmth I cannot reproduce this error, could you let me know how did you run into it?

Possibly user error, let's see if it goes through the CI ;)

[tlrmchlsmth]

Thanks for the great work, much nicer management of the Mamba Mixer state.

Now it's my turn to extend mamba_chunk_scan_combined to support this for Mamba2 :)

[tlrmchlsmth]

Merging now! from @simon-mo the readthedocs failure looks transient