Repository navigation
[ROCm][RDNA2] Add TP2 AWQ MoE expert offloading with graph replay - #146
Merged
Merged
Conversation
Keep symmetric W4A16 experts in pinned RAM with bounded GPU slots, decayed LFU/LRU replacement, and consumer-fenced transfers. Support TP2 MonoKernel execution, segmented graph replay, MTP2, and the native CPU KV tier. Add opt-in deterministic dense reduction and canonical QSA ordering for reproducible restore validation. Exclude graph padding from cache demand, and add focused tests, launchers, CI coverage, and 16k/32k validation receipts. Co-authored-by: Codex <noreply@openai.com> Signed-off-by: BlivionIaG <kev29lt@gmail.com>
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.
Flash-Next's expert weights do not fit the tested TP2 deployment when fully resident. This adds immutable pinned-RAM expert banks, bounded GPU slots with decayed LFU/LRU replacement, protected demand leases, consumer-fenced transfers, and overflow passes.
MonoKernel execution is split into captured prefix, expert-pass, and finish stages around host cache planning. Segmented FULL_AND_PIECEWISE graphs, MTP2, prefix caching, and the existing native CPU KV tier are supported. Padding rows do not enter residency statistics.
Dense W4A16 split-K atomics and unordered QSA block selection caused repeat variation. The offload launcher enables guarded Q32 dense accumulation; RDNA QSA canonicalizes selected block order. The existing dense resident default is preserved.
Validation
Target: two Radeon PRO V620 on par1-cs25 (HIP 6–7), TP2/EP2, FP16, MTP2, 96 expert slots per layer, 16 GiB native CPU KV tier. Model: wtdcode/Qwen3.8-Flash-Next-AWQ-W4A16, compressed-tensors symmetric INT4, group 128.
Limits and review scope
Forced-length random outputs have whitespace/repetition warnings and are performance evidence only. Long cells span two named isolated runs; KV parity uses an approximately 4k prompt separately from 32k generation/retrieval. No controlled resident-versus-offload speedup or global optimum is claimed.
Q32 has a measured cost: isolated qkvz prefill at M4096 rises from 8.84 to 14.17 ms, with a 224 MiB accumulator. Predictive prefetch, adaptive expert/KV rebudgeting, and arbitrary active-context paging are not implemented. Native KV offload covers eligible prefix/hybrid state.
Draft for review: the commit includes diagnostics, benchmark receipts, and native formatting changes alongside implementation. Details and evidence: docs/rdna2/moe-offloading-target.md and bench_results/2026-10-28-moe-offload/qualification-summary.json.
Related: #37 (earlier draft), #40.