Conversation
- Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents - Update pinned_tensor.cpp to compile under both nvcc and hipcc - Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family) - Guard NVIDIA arch checks to return None on ROCm - Skip nvcc version check in _toolchain.py when on ROCm - Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch) - Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm - Add is_rocm() and driver_hip_version() in backend.py - Add rocm-smi fallback in __main__.py for clangd generation - Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds, Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache - Add AMD ROCm classifier in pyproject.toml
Fail closed to eager execution when the HIP stream-memory handshake cannot survive capture and replay. Add a ROCm 7.14 graph batch-memop path with executor-owned signal and parameter storage, dynamically size graph flag slots, preserve the existing CUDA module API, and cover the safety and multi-format replay paths.
zihaomu
force-pushed
the
proposal/hybrid-decode-executor-v1
branch
from
September 16, 2026 07:57
7889bb5 to
d5fa332
Compare
Author
|
Rebased this proposal onto the refreshed CPU/Hybrid ROCm graph-safety PR (#378, head The conflict in Validation:
New head: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This PR does one thing:
It does not change LRU policy, expert placement, expert kernels, numerical
behavior, or CUDA/ROCm synchronization.
This is a stacked Draft PR on top of #378. The refactor itself is commit
d5fa332.Depends on #378. Related to #350.
1. What happened before this PR
OffloadMoELayeris a model layer, but its_decode_hybrid()method also owned theentire CPU/GPU execution schedule:
This mixed four different responsibilities in one model class:
CPU worker progress.
Only the fourth responsibility belongs to the model layer.
The practical problem was that supporting another model or device meant touching a
method that knew about all four concerns. In particular, AMD support appeared to need
its own Hybrid implementation even though AMD only differs in the low-level
synchronization mechanism.
2. The refactor principle
The refactor separates policy, orchestration, platform mechanism, and
model-specific computation:
The new executor is deliberately small. It knows how to split routes, overlap CPU and
GPU work, and merge their outputs. It does not contain CUDA checks, ROCm checks, cache
replacement logic, or quantization dispatch.
3. What changed in the code
New:
python/freetoken/moe/hybrid_decode.pyThis file introduces:
HybridDecodeRequest: the four inputs needed by Hybrid decode;HybridDecodeExecutor.decode(): the extracted route/scheduling algorithm.Changed:
python/freetoken/layers/moe.pyOffloadMoELayer._decode_hybrid()is removed.When the decode target is Hybrid, the layer now delegates:
The layer supplies
_run_cached_decode_expertsas the callback. That callback keepsusing the existing
_expert_gemmquantization dispatch and current cache views.In other words, the model layer still decides how GPU experts are computed, but no
longer decides how CPU and GPU execution is scheduled.
Changed:
python/freetoken/moe/offload_cache.pyWhen the Engine attaches a CPU executor for Hybrid mode, the cache creates one
HybridDecodeExecutorand stores it beside the CPU executor.This is temporary ownership chosen to preserve the current Engine initialization and
lifetime order. Moving runtime ownership is intentionally left for a later PR.
New tests and proposal
tests/moe/test_hybrid_decode_executor.pyverifies the extracted contract without aGPU;
docs/proposals/0001-hybrid-decode-executor.mdrecords the longer-term design andfollow-up boundaries.
4. How decode works after this PR
The runtime algorithm is unchanged; it now lives in one reusable executor.
For every Hybrid decode layer:
OffloadMoeCache.ensure_experts_hybrid()to apply the existing LRU andcapped-fetch policy.
-1means that route must run on the CPU.The overlap order remains:
Submitting CPU work before the GPU copy/kernel allows CPU GEMV to overlap with PCIe
traffic and GPU execution. Setting
FREETOKEN_HYBRID_OVERLAP=0keeps the existingdebug/measurement mode by moving the CPU wait before the GPU work.
5. How routes are split
For example, suppose the router selects experts
[11, 42]with weights[0.7, 0.3].The existing cache policy places expert 11 in GPU slot 3 and sends expert 42 to the
CPU:
The CPU receives only CPU routes. The GPU receives only GPU routes. Their masks are
complementary, so each routed contribution is computed exactly once and the two
outputs can be added directly.
This PR also preserves two graph-related invariants:
topk_idsis still rewritten in place;execution structure independent of the cache-hit pattern.
6. Why this works for both CUDA and ROCm
The shared algorithm above does not depend on a GPU vendor.
The vendor-specific boundary is
CpuMoeExecutor.decode_submit()/decode_sync(), which this PR does not change:HybridDecodeExecutorcalls the same two methods on either platform. It neither knowsnor needs to know which mechanism implements them.
For the same reason, AMD does not need a separate LRU.
OffloadMoeCachealready makesthe platform-independent placement decision and produces the same GPU-slot /
-1contract for CUDA and ROCm. Only synchronization differs, and that is below the new
executor boundary.
7. What did not change
-1overflow marker_expert_gemmgpu_partial + cpu_partial)This PR is a responsibility move, not a new Hybrid algorithm.
8. Validation
Run in a PyTorch 2.11 / ROCm 7.14 container without a GPU attached:
The new contract tests cover:
OffloadMoELayer.CUDA and ROCm hardware performance/regression runs are still required before moving
this PR out of Draft. The target is no measurable decode-throughput regression.
9. What reviewers need to decide
The implementation is correct if these four statements hold:
HybridDecodeExecutoris the right owner for shared route splitting and CPU/GPUscheduling.
OffloadMoELayershould expose only the model-specific GPU expert callback.CpuMoeExecutor.OffloadMoeCache.cpu_executoris acceptable until alater
OffloadMoeRuntimeownership refactor.