Skip to content

refactor(moe): extract hybrid decode orchestration - #491

Open
zihaomu wants to merge 1 commit into
FlashML-org:mainfrom
zihaomu:proposal/hybrid-decode-executor-v1
Open

zihaomu wants to merge 1 commit into
FlashML-org:mainfrom
zihaomu:proposal/hybrid-decode-executor-v1

Conversation

@zihaomu

@zihaomu zihaomu commented Sep 16, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Extract the shared Hybrid MoE decode schedule from OffloadMoELayer into
HybridDecodeExecutor. This is now a standalone, single-commit refactor on
main (9b585b7), not a stack containing #378.

  • Preserve the current LRU/fetch policy, in-place route rewriting, optional
    statistics, CPU submit/sync order, and FREETOKEN_HYBRID_OVERLAP=0 behavior.
  • Keep quantization dispatch and supports_inactive_slots handling in the layer:
    capable kernels receive -1; other kernels get zero-weighted slot 0.
  • Keep cache ownership acyclic and retain the executor across cache rebuilds.
  • Add routing/order/numerical/lifetime tests and make an existing offload-config
    test use Triton rather than an unrelated CUDA-only FlashInfer selection.

No changes to CPU kernels, CUDA/HIP synchronization, cache replacement policy,
or expert math. The short design note describes these boundaries.

Related: #350 and #378. This PR does not fix ROCm CPU/Hybrid graph safety;
#378 remains a separate change. Neither its implementation nor its validation
claims are included here.

Validation

One AMD Radeon AI PRO R9700 (gfx1201, 32 GiB); dual EPYC 9334; Ubuntu 24.04;
host kernel/driver 7.0.0-31-generic; PyTorch 2.11.0+rocm7.14.0,
HIP 7.14.60850, Triton 3.7.1. Synthetic expert banks; no model checkpoint.
Only one GPU was exposed to the container.

python setup.py build_ext --inplace
PYTHONPATH=python HIP_VISIBLE_DEVICES=0 CUDA_VISIBLE_DEVICES=0 \
  OMP_NUM_THREADS=4 MKL_NUM_THREADS=4 taskset -c 32-39 python -m pytest -q \
  tests/moe/test_hybrid_decode_executor.py \
  tests/moe/test_hybrid_fetch.py tests/moe/test_offload.py \
  tests/moe/test_cpu_moe.py tests/moe/test_cpu_moe_q4_0.py \
  tests/engine/test_moe_cpu_layers.py \
  tests/engine/test_attention_backend_arch.py \
  tests/engine/test_attention_backend_matrix.py \
  tests/models/test_quant_config.py \
  -m "not slow" -k "not cuda_graph"

181 passed, 42 skipped, 5 deselected. The new file contributes 21 CPU
contract cases and 12 real BF16 Hybrid CPU/GPU cases. The existing CPU-executor
tests additionally cover BF16, NVFP4, MXFP4, DS-FP4 and Q4_0 eager computation;
this is not five-format Hybrid end-to-end coverage.

A supplemental local A/B harness compared the saved main layer implementation
against this branch using identical cache/kernels/CPU executor: 36 combinations
x 12 steps, bitwise-equal outputs and rewritten routing, identical cache state
and statistics
. Batch sizes 1/4/8, fetch caps 0/1/8, overlap on/off and
router-weight-on-input on/off were exercised.

Local single-layer eager microbenchmark (not a shipped benchmark or model-serving
result): BF16, 3 layers, 8 experts, H=1024, I=512, top-k=2, fetch cap=1,
4 CPU workers, CPUs 32-39, warmup then 10 alternating A/B rounds of 128 calls.

Batch Overlap main (us/layer) refactor (us/layer) Change
1 off 313.01 315.75 +0.88%
1 on 325.73 328.44 +0.83%
4 off 434.34 434.37 +0.01%
4 on 370.60 370.98 +0.10%

Limits

  • No NVIDIA hardware, full-model serving/throughput, or CPU/Hybrid graph replay
    validation in this revision. The four graph-replay tests and one slow test
    were explicitly deselected above; graph support is inherited from main.
  • Initial unrestricted-affinity diagnostics intermittently waited a long time
    in HIP synchronization, including while executing the saved main path;
    two runs were interrupted. The completed A/B and pinned-core final suite and
    benchmark passed. This refactor does not claim to diagnose or fix those waits.

Ready for review as a focused orchestration refactor, not as a graph-safety fix
or a claim of zero overhead/full-model performance parity.

@zihaomu

zihaomu commented Sep 16, 2026

Copy link
Copy Markdown
Contributor Author

Rebased this proposal onto the refreshed CPU/Hybrid ROCm graph-safety PR (#378, head 94f7724) and ported the extraction to the current upstream MoE architecture.

The conflict in layers/moe.py was resolved semantically: the current moe_strategy and quant-method dispatch remain intact, while Hybrid route splitting / CPU-GPU overlap / result merging move into HybridDecodeExecutor. OffloadMoELayer now supplies only the narrow cached-GPU-expert callback. No legacy MoE implementation was restored.

Validation:

  • Hybrid executor contracts + ROCm graph-safety tests: 21 passed
  • adjacent Hybrid fetch tests: 3 passed
  • CPU MoE suites: 30 GPU-dependent cases collected/skipped in the no-GPU container
  • Python compile and git diff --check pass

New head: d5fa33231edd2328370fe7a7968b00701147c082.

@zihaomu
zihaomu force-pushed the proposal/hybrid-decode-executor-v1 branch from d5fa332 to 00e4bb1 Compare September 21, 2026 03:02
@zihaomu zihaomu mentioned this pull request Sep 23, 2026
2 of 29 tasks
LZ-QWQ added a commit to LZ-QWQ/FreeToken that referenced this pull request Oct 8, 2026
Integrate FlashML-org#133, FlashML-org#134, FlashML-org#135, FlashML-org#378 and FlashML-org#491 before FlashML-org#587. Preserve the tested FlashML-org#378 handshake and runtime-only linkage, and retain the FlashML-org#134 installation document. Exclude CI and closed GGUF PR FlashML-org#136.
@zihaomu
zihaomu force-pushed the proposal/hybrid-decode-executor-v1 branch from 00e4bb1 to 43fc3f0 Compare October 9, 2026 04:24
@zihaomu
zihaomu marked this pull request as ready for review October 9, 2026 04:25
LZ-QWQ added a commit to LZ-QWQ/FreeToken that referenced this pull request Oct 9, 2026
Use upstream FlashML-org#638 and FlashML-org#134; integrate incremental FlashML-org#135 and FlashML-org#378 plus the updated standalone FlashML-org#491 before FlashML-org#587. Retain the probed FlashML-org#378 handshake and upstream installation document.

This branch has not been deployed

No deployments
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.

1 participant