#3845·axolotl

[ 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