[CB] Reduce wastage in prefill compute and pad blocks in homogeneous continuous batching

[yannicks1]

[CB] Reduce wastage in prefill compute and pad blocks in homogeneous continuous batching

implement optimization idea by @JRosenkranz: do prefill only on next multiple of block size and then during decode pad with (valid) block id. Reduces computes for prefill and does not waist any valid blocks ids if whole blocks are padded to make tkv homogeneous.

solves #255

[github-actions]

👋 Hi! Thank you for contributing to vLLM support on Spyre.
Just a reminder: Make sure that your code passes all the linting checks, otherwise your PR won't be able to be merged. To do so, first install the linting requirements, then run format.sh and commit the changes. This can be done with uv directly:

uv sync --frozen --group lint --active --inexact

Or this can be done with pip:

uv pip compile --group lint > requirements-lint.txt
pip install -r requirements-lint.txt
bash format.sh

Now you are good to go 🚀

[yannicks1]

great news: This runs on Spyre 🎉

I just ran cb_spyre_inference.py which (with the parameters on this branch) exploits all functionality:

• prefill seq 1 of size 128 (left pads)
• prefill seq 2 of size 64 (left pads) (see [CB] Reduce wastage in prefill compute and pad blocks in homogeneous continuous batching #255)
• decode of batch 2
• strips fully padded blocks (see [CB] strip repeated left padding on batch level #131 )
• prefill seq 3 by left padding to 66 (align tkv) and right padding to 128 (pad to block boundary)
• decode of batch 2
• decode of batch 1 (see [CB] add min batch size of 2 in decode #182)

cc: @tdoublep @JRosenkranz @joerunde @nikolaospapandreou @sducouedic

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[yannicks1]

6/7 tests passed on the Spyre card! looks like the failure is a known issue unrelated to this PR. 🥳

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb.py MARKERS="spyre"

[yannicks1]

bot:test
MARKERS="spyre"

[yannicks1]

bot:test
TEST_FILE=tests/e2e/test_spyre_cb_scheduler_step.py MARKERS="spyre"

[prashantgupta24]

nit (should revert changes to this file)

Suggested change

	print("\nAll results match!\n")
	print("\nAll results match!\n")

[prashantgupta24]

worth adding some debug logs to the optimizations?

[maxdebayser]

Perhaps this is an incredibly stupid question, but why is it ok to use block id 0? Does it make a difference if it free (i.e. it's in self.block_pool) or not?

[prashantgupta24]

Based on Josh's comment in an internal issue,

"The only requirement when padding this is that we choose a block ID that exists in the pool of allotted block ids at server start (this way it will map to a real location in the memory space). In this case, when performing paged attention compute, the placeholder block will be part of compute (we will still take a performance hit during decode - at least until heterogeneous tkv is available), but will not take up any extra space as part of the block allotment at the time of server start."

IIUC I think that's why block 0 makes sense?

[prashantgupta24]

Although yeah I'm not sure what will happen if block 0 is actually being used by another request?

[prashantgupta24]

I managed to create a scenario where the block table looks like this for 2 requests:

block table: [deque([0, 2]), deque([0, 4])]

From what I can see the output still is correct, although anyone looking at the block table could be confused

[prashantgupta24]

Still processing the changes, but wondering if the comment above needs rewording...