Skip to content

cuda : pair floating routed gate/up projections - #124

Draft
GenerelSchwerz wants to merge 1 commit into
reference/upstream-kernels-20261001from
kernel/floating-routed-expert-pair-20261003
Draft

GenerelSchwerz wants to merge 1 commit into
reference/upstream-kernels-20261001from
kernel/floating-routed-expert-pair-20261003

Conversation

@GenerelSchwerz

Copy link
Copy Markdown
Owner

Overview

Pair separate floating routed gate/up projections that share activation and expert IDs. One existing MMF launch addresses both weight banks, writes both original projection buffers, then runs the original GLU. Compact batches also share one route-map preparation. No weight repacking, extra projection storage, precision change or model-name rules.

Based solely on official upstream dcd387a412ca54e172a8d60eb71ef6753850c8ca. Independent companion to #122, which handles quantized routed MMQ; this PR handles floating routed MMF. It does not depend on #117, #118, #119 or #123.

Coverage

  • Separate matching F32/F16/BF16 banks, F32 activations/output and common I32 expert IDs; broadcast or per-expert activations.
  • Original direct-ID MMF for up to16 tokens and compact-ID MMF above16, within existing MMF eligibility. Retains original per-bank row thresholds, including M640/T129 and T512; virtual doubled rows do not trigger the sorted cutoff.
  • SwiGLU, GeGLU, ReGLU and positive finite SwiGLU clamp, in either projection order.
  • Raw projection OUTPUT flags, later logical consumers and later private NONE/view aliases remain correct because both original outputs are written.
  • Bounded canonical bank views, graph capture/replay and updated inputs, routes, weights or output bindings. Adds standard graph-optimizer allocation dependencies, then verifies actual device ranges including NONE reads before execution.
  • Existing vector fusion keeps priority. Single-token decode can pair when upstream selects routed MMF, as demonstrated by the validation model's large expert banks.

Sorted/cuBLAS routes, ragged/alignment-ineligible layouts, mixed bank types, forced F32 precision, packed producers, incompatible GLUs, intervening real operations, unsafe aliases, unsupported view ownership and concurrent streams retain upstream execution. Private allocators that bypass graph optimization still need their own dependency integration; this PR does not weaken ownership checks or add private cache/runtime behavior.

Validation and performance

RTX5070Ti, sm120, driver615.71.09, CUDA13.4.59, Release build with j18. Candidate CUDA library and sources were frozen; all other libraries match the frozen upstream control. Every suite held the shared build then GPU locks throughout. Small correctness jobs used at most2 workers; profiles, memory checks and timings ran serially.

Whole component, matching upstream control Less elapsed time
F16/BF16, K512/M128/E8/U2, T8 30.17% /29.81%
F16/BF16, same shape, T32 33.39% /33.39%
F32, same shape, T8/T32 39.86% /42.80%
F16, K2560/M640/E16/U2, T129 9.10%
F16, K2560/M640/E16/U2, T512, hot routing 1.55%
F16, K512/M2048/E8/U2, T128 8.79%

First six cases:4 alternating pairs of1000 evaluations. K2560/T129:8 alternating pairs of10000 evaluations, replacing a shorter run with a startup outlier. Timings cover complete backend graph evaluation, including original GLU and routing preparation, rather than an isolated MMA instruction.

Two sorted control shapes were slower: M640/T513 by0.85% and M2048/T129 by0.56% over8 alternating10000-evaluation pairs. Small existing vector, packed and Q8 control timings were effectively flat in the shorter suite. This is a component optimization with a measured control-path cost, not a universal speedup.

  • 103 shape/type/activation/layout/update/alias comparisons: complete3-pass GPU outputs match upstream exactly.102 pass the unchanged CPU NMSE threshold5e-4. The remaining native FP4 control does not enter this floating path and reproduces upstream's existing CPU mismatch; it is not counted as a CPU pass.
  • Raw CUDA gate/up/final buffers for F16 K512/M128/T32 compact execution match upstream over all3 passes and pass CPU checks. Later private alias consumer payloads also match.
  • 72 whole-graph cases added to the existing backend test, covering3 floating types,3 token counts, broadcast and all4 GLUs. CPU/upstream CUDA controls and final candidate CUDA pass. Both direct and compact compute-sanitizer memchecks report0 errors.
  • Captured profiles prove actual paired dispatch, one compact map and the original standalone GLU; graph-off and fusion-disabled paths are checked.
  • Default compiled kernels retain REG/STACK/SHARED/LOCAL usage, and6 inspected actual default SASS variants have equal decoded instructions, predicates and operands; this is not binary-file identity. Kernel parameters gain24 bytes of constant argument storage; module metadata also changes. New supported compact variants use6 more registers; some direct variants use up to8 more. CUDA library grows4,086,824 bytes. Older NVIDIA and AMD hardware were not tested, so occupancy/performance there is unproven.
  • A standard offline F16 validation copy of tiny Mixtral ran through the unchanged loader/model/allocator, no MTP. Complete4-batch logits match at C1/C32; profiles show10 paired launches each, including9 direct T1 launches at C1. Longer5-pair timing repeats have mixed signs, so no whole-model speedup is claimed.200ms sustained VRAM samples match controls at844MiB/C1 and846MiB/C32; these are sampled peaks, not exact allocation attribution.

No combined-series Flash Next speedup is established. Quantized Flash Next expert banks are covered by #122, while this path applies when the actual floating graph meets the contracts above.

Reproduction and raw source/build/output/profile bindings are preserved under moe-cache-tests/results/upstream-kernel-series-20261001/FLOATING-EXPERT/FLOATING-ROUTED-20261003: SOURCE-FINAL-V4.json, BUILD-V7.json, EARLY-PROOF-V7.json, COMPILED-AUDIT-V7.json, BROAD-PROOF-V8.json, REMAINING-PROOF-V9.json, FOCUSED-TIMING-V1.json, MODEL-STEADY-V1.json and ROOT-FINAL-AUDIT-V1.json. Suite JSON records contain exact commands, environment, lock scope and teardown. Historical scratch-prototype results are retained separately and are not used as final evidence.

Requirements

  • I have read and agree with the contributing guidelines, subject to the owner-authorized fork automation exception in OWNER-FORK-POLICY.md. No upstream submission is made.
  • AI usage disclosure: YES. Codex implemented and tested this change; an owner-requested GPT-6.1 Sol High reviewer independently checked sources and proof artifacts. The repository owner remains responsible for review and maintenance.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant