feat(rdna): default-off hetero MoE for gfx1100 hot experts and gfx1030 cold experts

[BlivionIaG]

Purpose

Serve-side, default-off attention/expert disaggregation, replayed onto current rdna_extras after the v0.31 port.

Pool A (gfx1100 / W7800, pluggable so a Spark GB10 class can replace it later) keeps embeddings, attention and hybrid state, the router, shared experts, hot routed experts, and primary KV. Pool B (8x gfx1030 / V620) runs routed cold experts, expert-parallel, on the existing RDNA2W4A16MoEExperts / moe_gptq_gemm_rdna2 (fdot2) path, plus an LRU second-tier KV block store.

With VLLM_HETERO_MOE unset, RoutedExperts.forward_modular does not import the new package and still calls quant_method.apply. Clients stay on one OpenAI endpoint.

Design notes and the bring-up checklist are in docs/rdna2/hetero-moe-rdna3-rdna2.md. That doc cites MegaScale-Infer ping-pong (arXiv:2504.02263), Step-3 AFD / StepMesh (arXiv:2507.19427), the vLLM AFD plugin, vLLM EP/EPLB, --cpu-offload-params / --offload-params, and the KV connector API (LMCacheConnectorV1, OffloadingConnector).

Base

Retargeted from archived rdna_extra/v0.28.0 (c6d99ca3) onto rdna_extras.

The requested tip was 8abf4c97 (fix(rocm): moe_w4a8_rdna2 builds for every RDNA arch, the v0.31 port merge from PR #47). Live rdna_extras when this was replayed is 1dc3bac4cd (Merge feat/ple-int4: INT4 PLE table on the upstream PLE path, PR #49). 8abf4c97 is an ancestor of 1dc3bac4cd. The branch is based on the live tip so the draft is not immediately behind #49.

Head: 4144f2c642353e4bfb773c16fe684dcdfafd164c.

rdna_extras was not pushed. Nothing was opened upstream.

Duplicate check

No other open PR on opengfx1030/vllm-rdna implements mixed gfx1100/gfx1030 expert placement. Draft #37 is the routed-expert DRAM residency cache this change calls and does not reimplement. #37 is not being moved onto this base, so the bridge still fails closed.

AI assistance was used (Cursor cloud agent). This stays a draft until a human reviews every changed line.

What changed in the v0.31 replay

One conflict, in vllm/engine/arg_utils.py create_engine_config. v0.31 starts that method with logging setup (create_logging_config / configure_logging_if_needed). The hetero commit started it with _publish_hetero_moe_env(). Both are kept. Publish runs first and still writes nothing when every knob is at its default. envs.validate_environ on the old tree is current_platform.validate_environ on v0.31; that call was already on the new side and was left as-is.

Auto-merged, checked, and left in place:

• RoutedExperts.forward_modular still guards on os.environ.get("VLLM_HETERO_MOE", "0") == "1" before quant_method.apply. The v0.31 signature is unchanged: layer, x, topk_weights, topk_ids, shared_experts, shared_experts_input. The return type is now Tensor | UnfinalizedMoEOutput. The hook refuses a non-tensor so a deferred top-k reduction is not combined as hidden states.
• Hetero CLI fields, _publish_hetero_moe_env, and the --hetero-moe* argparse group.
• envs.py hetero entries.
• HeteroKVTierConnector registration.

Judgement calls:

• Did not call vllm.platforms.rdna.on_rdna_family() from the flag path. That helper imports vllm.platforms (and, on ROCm, vllm.platforms.rocm). Reading it at RoutedExperts import would run on every server with the flag off. Hetero stays an explicit two-pool switch, not "any RDNA GPU".
• gfx1100 hot kernel stays triton_wna16. gfx1030 cold kernel stays moe_gptq_gemm_rdna2 (fdot2). v0.31 WMMA units (gptq_gemm_rdna3_wmma, moe_gptq_gemm_rdna3, gptq_gemm_rdna3, __HIP__RDNA3__) are refused on gfx1030. Same-arch W4A8 sdot4 (moe_w4a8_gemm_rdna2) is not selected by this path.
• Peer probe stays at tools/rdna2/probe_hetero_peer.py. Serve launchers moved to tools/rdna/; tools/rdna2/ is still on this branch for machine probes, and the probe is not a launcher.
• PR feat(rdna2): opt-in routed-expert DRAM residency cache #37 residency is not vendored. ResidencyBridge.module() raises RuntimeError naming PR feat(rdna2): opt-in routed-expert DRAM residency cache #37 when vllm.model_executor.layers.fused_moe.rdna_expert_residency is missing.

What is implemented

• vllm/distributed/hetero_moe/: config, gfx1100 fast-tier class (gb10 registered and unimplemented), LFU/EMA hot set under a per-device byte budget, pair partition, float64 combine, loopback and host-staged transport (pinned buffer, shm, TCP length-prefixed frames), ping-pong micro-batch schedule, fatbin guard, contiguous 8-way cold EP map.
• Cold residency goes through ResidencyBridge to dest PR feat(rdna2): opt-in routed-expert DRAM residency cache #37 (ExpertResidencyCache.prepare, after_resident_gemm, next_moe_index, should_engage_expert_offload). There is no second slot table.
• HeteroKVTier plus HeteroKVTierConnector. Full-attention blocks are LRU-evicted by block hash and copied back before attention. remote_read always raises. GDN/PLE/recurrent groups are refused. QSA main KV and the compressed-key heap move only as one qsa_heap (two slots, or neither).
• get_num_new_matched_tokens returns no prefix hits. With APC off after the 2026-09-15 prefix corruption, the tier is for preemption swap and long-context spill, not prefix reuse.
• Serve-side env and CLI knobs. Defaults write nothing into the environment.
• tools/rdna2/probe_hetero_peer.py prints the peer-access matrix and a copy time only when it actually timed one.

What is stubbed or not wired

• split_routed_forward finishes its own receive before returning. run_ping_pong is the overlap schedule; gpu_model_runner does not call it.
• HeteroRuntime.finish_step is not called by the runner. Until it is, the placed hot set stays empty.
• The KV connector's tier is in-process. It does not HIP-copy paged KV onto a V620.
• Peer transport can be constructed only from a passing probe file, and its status stays UNVERIFIED. It does not perform a device peer memcpy and it does not use RCCL.
• The production cold function raises until a gfx1030 runner is installed. The CPU tests inject a reference apply. moe_gptq_gemm_rdna2 is selected by name and was not launched.
• gfx1100 W4A16 is the existing Triton path by name (triton_wna16), not launched here.
• PR feat(rdna2): opt-in routed-expert DRAM residency cache #37 is not on this base, so importing the residency module raises until that cache is present.

Unverified

No ROCm and no GPU in this environment. Not measured here: hipDeviceCanAccessPeer between gfx1100 and gfx1030, peer-copy bandwidth, mixed-arch RCCL (intentionally unused), either fatbin, and numerical match of -1 top-k padding against moe_gptq_gemm_rdna2. The probe script on this machine reported bandwidth not measured. No bandwidth number is hardcoded.

Open questions

• Same host for both pools, or a second host on the TCP frame?
• Fast tier stays 2x W7800 (gfx1100), or moves to Spark GB10?
• Which checkpoint is the first full-model target? The gfx1030 W4A16 kernel in this tree is the one used for Qwen4Exp / Qwen3.8 Flash-Next.

Test plan

.venv/bin/python -m pytest \
  tests/distributed/hetero_moe/test_hetero_moe.py \
  -q --tb=line --noconftest
.venv/bin/ruff check vllm/distributed/hetero_moe \
  vllm/distributed/kv_transfer/kv_connector/v1/hetero_kv_tier_connector.py \
  tests/distributed/hetero_moe/test_hetero_moe.py \
  tools/rdna2/probe_hetero_peer.py

CPU only: hot-set placement under a byte budget, loopback split+combine versus all-local MoE, empty-cold skip, host-staged and TCP round trip, ping-pong order, peer gate, KV LRU eviction and copy-back, QSA heap kept together, recurrent-state refusal, flags-off guard, probe script emits no invented bandwidth, gfx1030 refuses RDNA3 WMMA names, non-tensor quant_method.apply is refused.

Test result

On head 4144f2c642, with torch 2.14.1+cpu, pytest 9.1.1, and ruff 0.14.0 (the pre-commit pin):

24 passed, 1 warning in 2.03s

The warning is Torch failing to initialize NumPy (No module named 'numpy'). It is not a failed assertion. Two of the 24 tests are new on this replay: gfx1030 refuses the RDNA3 WMMA names, and a non-tensor quant_method.apply result is refused.

ruff check on the new modules passed after adding the blank line D413 wants after the last docstring section. ruff format --check reported those files unchanged.

No model evaluation. No tok/s claim. ROCm pin in the doc is 7.14.0; the gfx1030 line is hipcc --offload-arch=gfx1030 -O3 -mno-wavefrontsize64, without -ffp-contract=off.