d3b110803cd67417684abaf7b92d1e97025afcba
fused_experts() is called 64 times per decode step (once per MoE layer). Each call allocated 3 intermediate tensors via torch.empty = 192 mallocs. For decode (M=1, topk=8), all 64 calls use identical shapes. Fix: module-level _moe_intermediate_cache dict that reuses tensors when shapes match. First layer call allocates, subsequent 63 calls reuse. Saves 189 CUDA mallocs per decode step = 74,655 mallocs/second at 395 TPS. Design follows CCCL's dispatch_reduce.cuh pattern: pre-allocate temp_storage once via alias_temporaries, reuse across kernel invocations. No functional change — tensors are .empty() (uninitialized), overwritten before use by ixformer kernels.
project_6
Description
Languages
C++
41.8%
Cuda
31.6%
Python
22.2%
C
2.1%
CMake
1.1%
Other
1.1%