Skip to content

Reduce CPU overhead in expert packing - #31

Open
0z5a wants to merge 1 commit into
deepseek-ai:mainfrom
0z5a:codex/eplb-cpu-packing
Open

0z5a wants to merge 1 commit into
deepseek-ai:mainfrom
0z5a:codex/eplb-cpu-packing

Conversation

@0z5a

@0z5a 0z5a commented Nov 4, 2026

Copy link
Copy Markdown

Fixes #30.

rebalance_experts works with CPU float32 weights, but the greedy packing
loop currently performs individual PyTorch tensor indexing and arithmetic for
each item. Convert the sorted indices and weights to Python lists once for
this path, and use an array('f') for the pack sums so each accumulation
rounds to float32 as before. Keep the existing implementation for other dtypes
and devices. The placement policy and tie-breaking are unchanged.

Validation

  • Four focused unit tests pass on the CPU of a B300 host and on a GH200 CPU.
    They compare the placement and rank outputs against the original
    tensor-scalar algorithm across random, tied, very large, and non-contiguous
    weights; exercise the hierarchical and global public policies; and check
    other dtypes, empty batches, and a non-CPU default device.
  • Synthetic 61-layer, 256-logical/288-physical-expert, 16-group, 2-node,
    32-GPU simulated topology; one CPU thread, A/B/B/A order, PyTorch 2.14.1:
    on the B300 host CPU the base runs were 0.652/0.652 s and this change
    0.0439/0.0435 s; on the GH200 CPU the base runs were 0.879/0.884 s and this
    change 0.0341/0.0338 s. All three output tensors matched exactly in these
    runs. These are host-side planning timings, not distributed or model
    inference performance results.
  • git diff --check and Python syntax compilation pass. No GPU execution or
    multi-node communication was required for these tests.

PRs #27 and #28 change the packing policy to address a separate placement
constraint. This change retains the current policy, and may need coordination
if either of those PRs is accepted first.

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

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Reduce scalar tensor overhead in CPU expert packing

1 participant