SIGN IN SIGN UP

nixl_ep: round max tokens up to a multiple of 4 for TMA in elastic.py (#2056)

Low-latency dispatch aborts with "TMA requires the number of tokens to
be
multiple of 4" whenever active_rank_bound *
num_max_dispatch_tokens_per_rank is not a multiple of 4. For example:
two ranks with an odd max_tokens could not
run. Four nightly cases fail on this.
 
Round the maximum tokens up in elastic.py instead of hitting the assert.
Note
this rounds the declared capacity only, the actual token count is
untouched.
 
Applied where elastic.py hands the value to NIXL EP:
get_rdma_size_hint() and
the two dispatch() calls. 
 
The constraint comes from the FP8 scale tensor. Scales are stored
channel-major: all M token slots for one channel block, then the next,
where M = num_ranks * max_tokens. Each slot is 4 bytes in both
quantization modes, so the gap between blocks is 4M bytes, and TMA needs
that to be a multiple of 16 -- i.e. M a multiple of 4.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

* **Bug Fixes**
* Improved token-capacity handling for TMA alignment by rounding
capacities up to multiples of four.
* Ensured aligned capacities are used for buffer allocation and dispatch
while preserving the actual workload token count.

* **Documentation**
* Documented the requirement for low-latency buffer sizes and dispatch
token capacities to be multiples of four.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->
L
lishapira committed
b475f75b8357a3b68ca898b6647b0d1962c624a9
Parent: 4d030b9
Committed by GitHub <noreply@github.com> on 8/9/2026, 9:59:53 AM