Repository navigation
26.09-alpha.rc1
·
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>