Skip to content

26.09-alpha.rc1

Choose a tag to compare

@mmarcinkiewicz mmarcinkiewicz released this 06 Oct 09:08
· 6 commits to 26.09-alpha since this release
perf(moe): make paged-stash copy/pop kernel launch configurable (#7897)

The paged-stash copy/pop Triton kernels always launched with 1024 elements
per program per inner-loop iteration and at most 2048 programs. For wide
activations this leaves DRAM bandwidth unused: on VR200 an fp8
[32768, 7168] stash reaches ~3.8 TB/s, while a plain copy_ reaches
~9.9 TB/s.

Add TransformerConfig options moe_paged_stash_copy_block_size (power of
two, default 1024) and moe_paged_stash_copy_max_blocks (default 2048),
read by paged_stash_reset() and passed to every copy/pop launch. Defaults
keep the current behavior. With 4096 / 8192 the fp8 [32768, 7168] copy
goes from 119 us to 52 us and bf16 [32768, 4096] from 89 us to 58 us.

Add a bit-exact stash/reload round-trip test through PagedTensor for both
launch configurations, covering masked rows and more rows than programs,
and document the options.

Signed-off-by: Siddhartha Raman <sraman@nvidia.com>