[ Memory ] Compact then pad for sonicmoe EP sentinel
Author: NanoCode012Created Jul 17, 2026Updated Jul 17, 2026
Labelsenhancement
⚠️ Please check that this feature request hasn't been suggested before.
- I searched previous Ideas in Discussions didn't find any similar feature requests.
- I searched previous Issues didn't find any similar feature requests.
Feature description
This issue tracks Sonicmoe to add similar "memory" improvements (not compute) https://github.com/axolotl-ai-cloud/axolotl/pull/3714/
To satisfy quack's autotuner, we pad the recv to pow2 which adds unnecessary computation, but does speed up computation compared to without. In future, if we want to optimize further, we can try the below.
Cons:
- Needs a gather-in / scatter-out wrapper around sonicmoe_experts_forward_with_lora
✔️ Solution
Proposal: compact-then-pad. Drop -1 rows first, then pad the valid count
to its pow2 bucket instead of padding the full recv. At high EP degree
(ep4-ep8) valid_count ≪ recv_count, so we land in a much smaller pow2
bucket = real activation/input memory savings, while still collapsing to
a handful of cached shape keys (no autotune churn). Compute is unchanged
(already valid-only).❓ Alternatives
No response
Additional Context
No response
Acknowledgements
- My issue title is concise, descriptive, and in title casing.
- I have searched the existing issues to make sure this feature has not been requested yet.
- I have provided enough information for the maintainers to understand and evaluate this request.
Source: axolotl-ai-cloud/axolotl