Feature description
Give the bucketed-decode block-count axis its own denser bucket ladder, instead of sharing the power-of-two ladder used for the sequence-count axis.
Motivation and context
Bucketed decode pads every sequence in a batch up to bucket_num_blocks, so a batch whose sequences differ in context length already pays for blocks its shorter sequences never use. A power-of-two ladder adds a second, independent layer of the same waste on top: a batch needing 19 blocks is rounded to 32, so nearly half the padded blocks come from the lattice rather than from the length spread.
The two axes are not symmetric, which is what makes this safe to change on one of them: b_seqs sets the matmul lead dimension and core tiling, so it wants coarse steps, while b_blocks only unrolls the kernel's block loop and can take a finer ladder without changing kernel structure.
Proposed solution
Add a separate ladder for the block axis with intermediate steps between the powers of two, leaving the sequence axis unchanged. Worst-case rounding then drops from half the bucket to about a third, for example, the 19-block case takes 24.
Builder-side only, with no runtime cost, and it cannot pick a worse bucket than the power-of-two ladder because every added step is a tighter fit. The cost is more compiled artifacts: the bucket count grows from 6 to 9 at max_model_len=2048, populated lazily.
Worth a test asserting the rounding bound across all block counts, so a future change to the step list cannot silently widen it.
Checklist
Feature description
Give the bucketed-decode block-count axis its own denser bucket ladder, instead of sharing the power-of-two ladder used for the sequence-count axis.
Motivation and context
Bucketed decode pads every sequence in a batch up to
bucket_num_blocks, so a batch whose sequences differ in context length already pays for blocks its shorter sequences never use. A power-of-two ladder adds a second, independent layer of the same waste on top: a batch needing 19 blocks is rounded to 32, so nearly half the padded blocks come from the lattice rather than from the length spread.The two axes are not symmetric, which is what makes this safe to change on one of them:
b_seqssets the matmul lead dimension and core tiling, so it wants coarse steps, whileb_blocksonly unrolls the kernel's block loop and can take a finer ladder without changing kernel structure.Proposed solution
Add a separate ladder for the block axis with intermediate steps between the powers of two, leaving the sequence axis unchanged. Worst-case rounding then drops from half the bucket to about a third, for example, the 19-block case takes 24.
Builder-side only, with no runtime cost, and it cannot pick a worse bucket than the power-of-two ladder because every added step is a tighter fit. The cost is more compiled artifacts: the bucket count grows from 6 to 9 at
max_model_len=2048, populated lazily.Worth a test asserting the rounding bound across all block counts, so a future change to the step list cannot silently widen it.
Checklist