[Bugfix] Mamba cache Cuda Graph padding#6214

Merged

simon-mo merged 3 commits intovllm-project:mainfrom

tomeras91:mamba-cg-cache-padding

Jul 8, 2024

Commits on Jul 8, 2024

bugfix: when working in CG mode, batch size should be by input_ids shape and not by number of sequences. This is so we pad the mamba cache to the captured CG batch sizes
tomeras91
committed
Add relevant test
tomeras91
committed
rename batch_szoe -> cg_batch_size for clarity
tomeras91
committed