Skip to content

feat(rdna2): opt-in routed-expert DRAM residency cache - #37

Draft
BlivionIaG wants to merge 2 commits into
rdna_extra/v0.28.0from
cursor/moe-expert-residency-17a5
Draft

BlivionIaG wants to merge 2 commits into
rdna_extra/v0.28.0from
cursor/moe-expert-residency-17a5

Conversation

@BlivionIaG

@BlivionIaG BlivionIaG commented Oct 4, 2026 •

Copy link
Copy Markdown

Purpose

Opt-in VRAM residency cache for routed MoE experts, in the spirit of Strata's expert cache. A hot subset stays resident; a miss is copied from a pinned host buffer. An optional popularity profile chooses which experts start hot.

This is a vLLM path, not a Strata port.

Taken from Strata: the residency idea (hot experts in VRAM, miss filled from host RAM) and an optional hottest-first popularity ranking.

Left behind: Strata's runtime, GGUF packs, CPU expert kernels, the STRP profile container, and any tok/s numbers. The profile here is JSON {"ranked": [[layer, expert], ...]} or text lines layer expert, hottest first.

VLLM_RDNA_MOE_RESIDENT / VLLM_RDNA_MOE_RESIDENT_SKINNY were investigated and not reused. Those flags keep the full W4A16 set in the native device layout and dispatch the resident kernel. They are not a hot-expert DRAM cache. If either flag is on, this cache stays off so the two policies do not overlap.

Integration point:

  • When the knob is on, RoutedExperts.create_weights places new routed-expert parameters on CPU so the full set is not committed to VRAM first. Meta init, an already-CPU default device, monolithic kernels, and EPLB are not redirected.
  • finalize_routed_expert_residency runs at the end of process_weights_after_loading, after the existing quant packing (fdot2 / sdot4).
  • RoutedExperts.forward_modular calls cache.prepare, which fetches misses and remaps local ids onto slot indices. The existing quant_method.apply path is unchanged. After that GEMM is launched, after_resident_gemm prefetches the next MoE layer.

Prefetch overlap

The old copier called current.wait_stream(copy_stream) before returning, so every fetch stalled the compute stream before the GEMM.

Now:

  • A hit on an already-resident expert does not record a compute-stream wait.
  • The miss set is the next routed MoE layer's local experts that are not in the bound bank (last routed set when that layer has run, otherwise the profile order). Successor lookup uses next_moe_index, so dense, linear, GDN, KDA, and QSA-only blocks are skipped. It is not layer_index + 1.
  • Shared experts and MoVA V-experts are not in that set. This tree has no MoVA module; the name filter (mova / v_expert / shared_expert) is what keeps them out.
  • The copy is pinned host memory on a side stream (copy_(non_blocking=True) there, the torch equivalent of hipMemcpyAsync). An event is recorded after the copy. The copier does not wait_stream in either direction.
  • The copy destination is the staging bank, a second allocation. It is not a slot whose storage the in-flight GEMM is still reading. If that staging bank is itself live, the prefetch is skipped rather than aliasing those bytes.
  • The next MoE GEMM waits on the copy event before it reads the staged row, then publishes into the bank it will run. A miss that was not prefetched is still fetched before use. A half-copied slot is not consumed.
  • Under TP the ids are this rank's local experts only (expert_map / local range). Another rank's experts are not pulled.

No speedup is claimed. CPU tests check the wait and destination rules, not overlap timing.

Board constraints verified against the tree

  • Routed experts only. Shared experts and leftover dense tensors (embed, head, norms, QSA/GDN/PLE, router, mHC, MTP) stay on their existing resident or host contracts. The cache only sees RoutedExperts parameters whose leading dim is the local expert count. Families without a shared expert (MiMo, Naive) are the same bind minus that tensor.
  • Engram, PLE, and n-gram tables are not expert pages. There is no Engram module in this tree. PLE / n-gram tensors are excluded by name, as are shared-expert and MoVA V-expert names.
  • Host copy is a pinned CPU buffer on a side stream. It is not a fatbin and not a produce-quant change. Host pages live on the cache object. prepare and prefetch refuse while a HIP/CUDA graph is capturing. Serve this offload eagerly.
  • Heaps. This tree has no named cumem tags for GDN, QSA, PLE, or hc_combine (only "weights" and "kv_cache"). Those four stay separate because this path never allocates into those modules.
  • TP / EP: the residency table is per rank. Prefetch does not pull another rank's experts.
  • Once an expert is on device, fdot2 / sdot4 is the existing kernel. No WMMA. ROCm 7.14.0 pin is untouched.
  • Knob is default-off and serve-side. VLLM_RDNA_MOE_EXPERT_OFFLOAD defaults to 0. Optional VLLM_RDNA_MOE_EXPERT_PROFILE. Clients, gateways, and the tools/rdna2/serve_*.sh scripts do not export it. It engages only when routed-expert bytes exceed the physical VRAM budget. If they fit, no host pages or cache are kept.
  • Not in this PR: dest tip, UNC bump, Linear, wiki, FA, AR, produce (-cb 3inst / fdot2), W4A8 sdot4 knobs, fatbins. rdna_extras is not pushed. Draft only. Do not merge.

Duplicate check

No other open PR on opengfx1030/vllm-rdna covers expert DRAM offload or this residency cache. This follow-up stays on draft #37.

AI assistance was used (Cursor cloud agent). A human submitter must review every changed line before this leaves draft.

Test Plan

CPU policy and cache tests (no gfx1030 in this environment):

.venv/bin/python -m pytest \
  tests/model_executor/layers/fused_moe/test_rdna_expert_residency.py \
  -v --tb=short --noconftest

--noconftest avoids the repo root tests/conftest.py import (tblib), which this unit file does not need.

Covered: default off; resident flags block a second path; engage only when enabled and over budget; slot assignment; profile parse; routed-parameter filter (shared / n-gram / MoVA / router skipped); per-rank local ids; alias rebind; disabled install and fitting budget stay dead; profile seed, hit, and host miss; colder eviction; capture refuses the fetch; next MoE index skips dense / GDN / KDA / QSA / MoVA; a hot hit records no compute wait and the prefetch H2D targets the next MoE layer's staging bank; the consumer waits before bind; a live GEMM slot is not the copy destination; prefetch ids stay inside this rank's local range; meta init is not redirected to CPU.

Test Result

19 passed, 2 warnings in 4.34s

Warnings: missing vllm._version commit hash, and a requires_grad scalar conversion on the excluded n-gram tensor. Neither is a failure.

No model evaluation was run. This VM has no gfx1030. Produce quant, dest tip, and the on-device fdot2 / sdot4 kernels are unchanged. No tok/s number is claimed.

Open in Web Open in Cursor 

Hot routed experts stay in VRAM and a miss is copied from a pinned host
buffer on a side stream. The cache is default-off, serve-side, and stays
dead when the routed experts already fit. Shared experts, dense leftovers,
and PLE / n-gram tables are not expert pages. The existing fdot2 / sdot4
path is unchanged.

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
Prefetch the next routed MoE layer's local misses with an async copy
of pinned host pages on a side stream, into a staging bank the in-flight
GEMM is not reading. The resident GEMM does not wait. The next MoE GEMM
waits on the copy event before it uses that slot. Dense, GDN, KDA, and
QSA-only blocks are not the next miss set.

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
cursor Bot pushed a commit that referenced this pull request Oct 11, 2026
Keep hot routed experts and primary KV on the fast tier, send only
cold pairs to the gfx1030 pool, and add an LRU KV block tier that
copies blocks back before attention. With VLLM_HETERO_MOE unset the
MoE forward is unchanged.

Replayed onto rdna_extras after the v0.31 port. Engine startup still
publishes hetero env only when a knob is set, and keeps the v0.31
logging setup. Cold W4A16 stays on moe_gptq_gemm_rdna2 (fdot2). RDNA3
WMMA objects are refused on gfx1030. A non-tensor quant_method.apply
result is refused. The PR #37 residency bridge still fails closed.

Co-authored-by: Cursor Agent <cursoragent@cursor.com>
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.

2 participants