Skip to content

MoE offload: decode overhead from padding masks (#83) and padding counted as expert 0 #89

Closed

Description

@gspark-etri

Problem

After #83 (merged through #73), MoE offload decode throughput is about 5% lower. Padding ids are also counted as routing hits on expert 0.

Evidence

Qwen3-30B-A3B-Instruct-2507, moe_offload_generate.py --gpu-budget-gb 20 --kv-cache --max-new-tokens 128, 1× A100 80GB + host CPU, torch 2.13.0+cu130, transformers 5.15.0. Interleaved runs, 3 each:

tok/s
#73 as opened (41cfb2d) 11.22 / 11.02 / 11.60 (mean 11.28)
with #83 (11b3f37) 10.88 / 10.03 / 11.00 (mean 10.64)
  • −5.7% mean, −3.0% median. Generated text is identical.
  • A one-off adaptive run (--adapt-every 48) shows the same gap: 11.07 → 10.21.
  • Likely cause: OffloadedLayer.forward now runs two compares, an and, a clamp and a mul on every call (about 5 kernel launches × 48 layers per token).
  • _counts_d and the tier counts use the clamped ids, so padded rows count as expert 0.

Scope

  1. Skip the mask/clamp when padding cannot occur, or fold it into the existing routing copy.
  2. Count routing with the unclamped ids masked by valid.
  3. In opt_grouped_mm, start from torch.zeros instead of hidden*0, so an inf in hidden does not become NaN.
  4. Make tests/ skip, instead of fail, model families that the installed transformers does not have (on 4.46.2: solar_open, qwen3_moe, qwen3_next; rope_parameters).

Done when

  • Throughput is back within noise of 41cfb2d on the setup above, with identical text.
  • Padding no longer changes routing counts (test).
  • pytest opt_prime/tests on torch 2.5.1 + transformers 4.46.2 has no failures, only skips.

Owner

@gspark-etri

Activity

  1. gspark-etri commented on Oct 30, 2026

    @gspark-etri
    CollaboratorAuthor

    Fixed by #288, squash-merged into optimus-rm (A/B: +6.5% mean tok/s on Qwen3-30B-A3B, see the PR).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions