From f860368a8a8cca0da8b39536c15df4c1e895e4fc Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Mon, 7 Sep 2026 00:01:36 +0000 Subject: [PATCH] perf(ops): warm L2 for the next projection from the MoE down tail, issuing the hint before the block barrier In the T=1 decode path sparse_moe_d4_nine_warp_kernel is the last producer before the next layer's attention or GDN projection reads its weights from HBM. Those weights are a fixed, known address range while D4 is still running, and L2 holds nothing of them: no earlier kernel in the round has touched them. The kernel now receives that range and issues prefetch.global.L2 over it, one 128-byte line per thread, clamped to 8 MiB against a 96 MiB L2. The hint is issued before the block barrier rather than after the FP32 rank-order epilogue, so it hides behind the slowest warp instead of extending the grid tail. Ownership. The range is passed explicitly and nothing is stored. A new public struct ops::SparseMoeHints names the span; a second overload of ops::sparse_moe() takes it, with the existing six-argument signature kept and forwarding a default-constructed SparseMoeHints{} to it - the idiom already at include/ninfer/ops/linear.h:107 and :121. From there the value rides in the plan, as prefill and small-T already do: two fields on SparseMoeDecodePlan, a third defaulted parameter on resolve_sparse_moe_decode_plan() so the capacity call site is untouched, and sparse_moe_decode_launch() taking a const SparseMoeDecodePlan&. The 8 MiB clamp became a pure function inside resolve_sparse_moe_decode_plan(). There is no thread-local channel, no setter and no module-global device state; the Op keeps no state between calls and is re-entrant, and CUDA Graph capture is correct by construction because the pointer is an ordinary kernel argument computed on the same path as the launch. No test and no bench file is touched. This is a cache hint: it reads no value, writes nothing, and cannot change an addition order. cuobjdump --dump-resource-usage over sparse_moe_decode_kernels.cu.o, matched by demangled name with template arguments, shows the three Rows=1 bodies gaining two parameters with registers 40 -> 40, shared 36 -> 36, stack 0 and spills 0, and every other body unchanged; issuing the hint before rather than after the barrier changes 0 of 33 bodies. That census was taken on e3aeaf8c and has not been repeated on the current base; the D4 kernel body it describes is byte-identical to the one here. Claim level: schedule, cross-kernel producer-consumer. Measured on one RTX 5090 (sm_120a, CUDA 13.1.115, driver 580.95.05, Release, -DCMAKE_CUDA_ARCHITECTURES=120a), artifact Qwen3.6-35B-A3B, batch 1, every measurement under a GPU-exclusive lock with the card-witness window opened inside the lock. All arms built from one tree and run in one pass over one set of points; seven passes with pass 0 discarded and the median over the remaining six; arm order reversed on alternate passes; one process per cell; two zero controls, a bit-identical copy of the master binary and a rebuild of master. ninfer_bench -pg 128,128;1000,128;4095,128;12000,128 -r 10 --warmup 2 --max-ctx 16384, decode output t/s against master at mtp0, i.e. with neither speculative flag: +1.944 / +1.873 / +1.808 / +1.720 %, pass spread +1.183 to +2.694 %, zero controls -0.010 to +0.101 %. The claim is the interval, +1.7 to +1.9 %; the medians are printed because that is what the raw says, but the zero-control band reaches +/-0.57 % at the extremes on the short context. An earlier edition of this change, on ad0f3d38 and split across two commits, read +1.74 / +1.70 / +1.53 / +1.54 %; both sets are left visible. The effect is flat in context length because what is warmed is a weight block of fixed size, not request data. At --spec mtp --draft-tokens 3 the change measures -0.049 to +0.034 % against a zero control of -0.139 to +0.039 %, and prefill medians are zero to within 0.06 % (a single run reaches +0.663 %). That zero is structural, not weak: the prefetch parameters exist only on Rows=1 instantiations, T >= 2 goes to sparse_moe_small_t_launch which has no prefetch at all, and the MTP post-mixer passes SparseMoeHints{} explicitly. Carried over from ad0f3d38 and not re-run: draft 1 and draft 5 read +0.01 and -0.09 to -0.10 %, the CUDA Graph decode step at draft 3 read -0.03 to +0.04 %, and ninfer_sparse_moe_bench --sweep 2:12 read +0.00 % in all 22 cells - the last a control rather than evidence, since the sweep starts at T=2 and never reaches the T=1 kernel. There is no production-shape operator benchmark for this mechanism. Output is bit-identical to master by MD5 and length over greedy generations with neither speculative flag and at --spec mtp --draft-tokens 3, three prompts, two repeats, arms alternated: 12 of 12, with a master-against-master control also at 12 of 12 and no reply below the 40-byte floor. The gate is weaker than that score suggests and the weakness is stated rather than left to be found: the repeats are the same deterministic run, and two of the three prompts return identical output under both modes, so exactly one of the twelve cells discriminates. The gate was also repaired mid-work - its first run returned 36 one-byte replies because a thinking model sent every permitted token to reasoning and reasoning goes to stderr; the 40-byte floor caught it and --no-thinking fixed it. ctest on the base this branch sits on, both arms from the same build tree: 107 of 114 registered targets executed, 106 passed, ninfer_attn_input_proj_test failed, 7 skipped as artifact-gated - identical on the arm and on the bare base, so the single failure is an upstream defect (issue #196) and not something this branch introduces. No round of "114 of 114" is claimed. Real weights were not attached, so the five tests that exercise the changed decode route did not run. clang-format 23.1.0 delta to the base is 0. Base and provenance. This change targets upstream a16b6442; it sits on a16b6442's parent 487f8977, which is the commit every measurement above was taken on, and the single commit between the two touches documentation only and none of the eleven files changed here. Earlier heads are named at the point of use for the figures taken on them. The bench flag --mtp-draft-tokens no longer exists: the spelling is --spec mtp --draft-tokens N, and --draft-tokens 0 is rejected for MTP (src/product/speculative_options.h:41-43), so the zero-draft arm is spelled by passing neither flag - the identical engine state, since SpeculativeOptions defaults to backend None with draft_tokens 0 and the validator accepts None only when draft_tokens is 0. --- include/ninfer/ops/sparse_moe.h | 22 ++++++++++ src/ops/sparse_moe/decode/sparse_moe_decode.h | 7 +++- .../decode/sparse_moe_decode_kernels.cu | 41 ++++++++++++++----- .../decode/sparse_moe_decode_plan.cpp | 15 ++++++- src/ops/wrapper/sparse_moe.cpp | 11 ++++- .../qwen3_6/impl/runtime/text_context.h | 5 ++- .../qwen3_6/impl/runtime/text_context_impl.h | 23 +++++++++-- src/targets/qwen3_6_27b/impl/variant.cpp | 3 +- src/targets/qwen3_6_27b/impl/variant.h | 15 ++++++- src/targets/qwen3_6_35b_a3b/impl/variant.cpp | 14 ++++--- src/targets/qwen3_6_35b_a3b/impl/variant.h | 15 ++++++- 11 files changed, 140 insertions(+), 31 deletions(-) diff --git a/include/ninfer/ops/sparse_moe.h b/include/ninfer/ops/sparse_moe.h index 73cc05b89c..9bce43aae9 100644 --- a/include/ninfer/ops/sparse_moe.h +++ b/include/ninfer/ops/sparse_moe.h @@ -22,6 +22,20 @@ enum class SparseMoeEpilogue : std::uint8_t { AddResidual, }; +/** + * Optional per-call execution hints. Every field is a pure cache hint with no numeric effect: + * the same call with a default-constructed SparseMoeHints produces bit-identical output. + * + * next_weight_prefetch names a weight span the next decode-step consumer will stream; the decode + * D4 epilogue issues fire-and-forget L2 prefetches over its first bytes. The span is caller-owned + * and read once, inside the call: the Op keeps no state between calls, and no hidden channel + * carries it. + */ +struct SparseMoeHints { + const void* next_weight_prefetch = nullptr; + std::size_t next_weight_prefetch_bytes = 0; +}; + /** * Returns the transient capacity required by SparseMoe for every T in the inclusive * [min_tokens,max_tokens] interval. The routed QTypes are the fixed implementation profile. @@ -62,4 +76,12 @@ enum class SparseMoeEpilogue : std::uint8_t { void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue, Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream); +/** + * The same Op with caller-supplied execution hints. Semantics, workspace requirement and output + * are exactly those of the overload above; hints only steer cache warming. + */ +void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue, + Tensor& destination, const SparseMoeHints& hints, WorkspaceArena& workspace, + cudaStream_t stream); + } // namespace ninfer::ops diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode.h b/src/ops/sparse_moe/decode/sparse_moe_decode.h index 30da87bf71..b8545bd586 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode.h +++ b/src/ops/sparse_moe/decode/sparse_moe_decode.h @@ -16,6 +16,9 @@ enum class SparseMoeSmallTD4Schedule : std::uint8_t; struct SparseMoeDecodePlan { std::size_t workspace_bytes = 0; + // L2 cache hint resolved from the caller's SparseMoeHints. No numeric effect. + const void* next_weight_prefetch = nullptr; + std::size_t next_weight_prefetch_bytes = 0; }; struct SparseMoeDecodeWorkspace { @@ -39,7 +42,8 @@ SparseMoeDecodeWorkspace allocate_sparse_moe_decode_workspace(Arena& arena) { [[nodiscard]] std::size_t sparse_moe_decode_workspace_bytes(); [[nodiscard]] SparseMoeDecodePlan resolve_sparse_moe_decode_plan(QType routed_gate_up, - QType routed_down); + QType routed_down, + const SparseMoeHints& hints = {}); void sparse_moe_decode_launch_d3_small_t(const Tensor& x, const SparseMoeWeights& weights, const int* token_ids, float* token_activations, @@ -53,6 +57,7 @@ void sparse_moe_decode_launch_d4_small_t(const SparseMoeWeights& weights, Tensor cudaStream_t stream, const int* adaptive_route_jobs = nullptr); void sparse_moe_decode_launch(const Tensor& x, const SparseMoeWeights& weights, Tensor& destination, + const SparseMoeDecodePlan& plan, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream); } // namespace ninfer::ops::detail diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 12a68e0c18..c91b4611c2 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -378,7 +378,8 @@ __global__ void sparse_moe_d4_nine_warp_kernel( const float* __restrict__ shared_scale, const float* __restrict__ act, const std::uint8_t* __restrict__ routed_codes, const std::uint8_t* __restrict__ routed_high, const std::uint8_t* __restrict__ routed_scales, const std::uint8_t* __restrict__ shared_codes, - const std::uint8_t* __restrict__ shared_scales, __nv_bfloat16* __restrict__ destination) { + const std::uint8_t* __restrict__ shared_scales, __nv_bfloat16* __restrict__ destination, + const char* __restrict__ prefetch_data, unsigned long long prefetch_bytes) { __shared__ float paths[kTopK + 1][Rows]; pdl::wait_for_dependencies(); const int warp = static_cast(threadIdx.x) >> 5; @@ -405,6 +406,16 @@ __global__ void sparse_moe_d4_nine_warp_kernel( for (int row = 0; row < Rows; ++row) { paths[kTopK][row] = *shared_scale * dot[row]; } } } + if (prefetch_data != nullptr) { + // Fire-and-forget L2 warmup of the next consumer's weight codes, + // issued BEFORE the block barrier so it hides behind the slowest warp instead + // of extending the grid tail. Pure cache hint; no values, no addition order. + const unsigned long long offset = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) * 128ull; + if (offset < prefetch_bytes) { + asm volatile("prefetch.global.L2 [%0];" ::"l"(prefetch_data + offset)); + } + } __syncthreads(); if (warp == 0 && lane < Rows) { float value = __bfloat162float(destination[row_base + lane]); @@ -528,7 +539,8 @@ void launch_d2_d3(const Tensor& x, const SparseMoeWeights& weights, template void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destination, - const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { + const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, + const void* prefetch_data, std::size_t prefetch_bytes) { const auto* ids = static_cast(workspace.ids.data); const auto* alpha = static_cast(workspace.alpha.data); const auto* shared_scale = static_cast(workspace.shared_scale.data); @@ -539,23 +551,28 @@ void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destinat const auto* shared_codes = static_cast(weights.shared_down.qdata); const auto* shared_scales = static_cast(weights.shared_down.scales); auto* output = static_cast<__nv_bfloat16*>(destination.data); - CUDA_CHECK(pdl::launch_dependent({dim3(kHidden), dim3(9 * 32), 0, stream}, - sparse_moe_d4_nine_warp_kernel, ids, alpha, - shared_scale, act, routed_codes, routed_high, routed_scales, - shared_codes, shared_scales, output)); + CUDA_CHECK(pdl::launch_dependent( + {dim3(kHidden), dim3(9 * 32), 0, stream}, sparse_moe_d4_nine_warp_kernel, ids, + alpha, shared_scale, act, routed_codes, routed_high, routed_scales, shared_codes, + shared_scales, output, static_cast(prefetch_data), + static_cast(prefetch_bytes))); } void launch_d4_dependent(const SparseMoeWeights& weights, Tensor& destination, - const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { + const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, + const void* prefetch_data, std::size_t prefetch_bytes) { switch (weights.routed_down.qtype) { case QType::Q5G64_F16S: - launch_d4_dependent_codec(weights, destination, workspace, stream); + launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, + prefetch_bytes); return; case QType::Q6G64_F16S: - launch_d4_dependent_codec(weights, destination, workspace, stream); + launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, + prefetch_bytes); return; case QType::W8G32_F16S: - launch_d4_dependent_codec(weights, destination, workspace, stream); + launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, + prefetch_bytes); return; default: throw std::invalid_argument("sparse_moe: unsupported D4 codec"); @@ -718,10 +735,12 @@ void sparse_moe_decode_launch_d4_small_t(const SparseMoeWeights& weights, Tensor } void sparse_moe_decode_launch(const Tensor& x, const SparseMoeWeights& weights, Tensor& destination, + const SparseMoeDecodePlan& plan, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { launch_d1(x, weights.router_shared_gate, workspace, stream); launch_d2_d3(x, weights, workspace, stream); - launch_d4_dependent(weights, destination, workspace, stream); + launch_d4_dependent(weights, destination, workspace, stream, plan.next_weight_prefetch, + plan.next_weight_prefetch_bytes); } } // namespace ninfer::ops::detail diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp b/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp index 736aba32e3..52a375a160 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_plan.cpp @@ -5,6 +5,11 @@ #include namespace ninfer::ops::detail { +namespace { +// The D4 epilogue warms at most this many bytes of the next consumer's weights: far below the +// 96 MiB L2, so the warmed block cannot evict what it was meant to help. +constexpr std::size_t kNextWeightPrefetchLimit = std::size_t{8} << 20; +} // namespace std::size_t sparse_moe_decode_workspace_bytes() { WorkspaceLayoutBuilder layout; @@ -12,7 +17,8 @@ std::size_t sparse_moe_decode_workspace_bytes() { return layout.peak_bytes(1); } -SparseMoeDecodePlan resolve_sparse_moe_decode_plan(QType routed_gate_up, QType routed_down) { +SparseMoeDecodePlan resolve_sparse_moe_decode_plan(QType routed_gate_up, QType routed_down, + const SparseMoeHints& hints) { const bool main_profile = routed_gate_up == QType::Q4G64_F16S && (routed_down == QType::Q5G64_F16S || routed_down == QType::Q6G64_F16S); @@ -24,6 +30,13 @@ SparseMoeDecodePlan resolve_sparse_moe_decode_plan(QType routed_gate_up, QType r SparseMoeDecodePlan plan; plan.workspace_bytes = sparse_moe_decode_workspace_bytes(); + if (hints.next_weight_prefetch != nullptr) { + plan.next_weight_prefetch = hints.next_weight_prefetch; + plan.next_weight_prefetch_bytes = + hints.next_weight_prefetch_bytes < kNextWeightPrefetchLimit + ? hints.next_weight_prefetch_bytes + : kNextWeightPrefetchLimit; + } return plan; } diff --git a/src/ops/wrapper/sparse_moe.cpp b/src/ops/wrapper/sparse_moe.cpp index 48e9c725fc..d1a8ddef8c 100644 --- a/src/ops/wrapper/sparse_moe.cpp +++ b/src/ops/wrapper/sparse_moe.cpp @@ -191,6 +191,12 @@ std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType rout void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue, Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream) { + sparse_moe(x, weights, epilogue, destination, SparseMoeHints{}, workspace, stream); +} + +void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue, + Tensor& destination, const SparseMoeHints& hints, WorkspaceArena& workspace, + cudaStream_t stream) { if (epilogue != SparseMoeEpilogue::AddResidual) { throw std::invalid_argument("sparse_moe: unsupported epilogue"); } @@ -252,13 +258,14 @@ void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilo } const detail::SparseMoeDecodePlan plan = detail::resolve_sparse_moe_decode_plan( - weights.routed_gate_up.qtype, weights.routed_down.qtype); + weights.routed_gate_up.qtype, weights.routed_down.qtype, hints); const detail::SparseMoeDecodeWorkspace views = detail::allocate_sparse_moe_decode_workspace(workspace); for (std::int32_t token = 0; token < tokens; ++token) { const Tensor x_column = x.slice(1, token, 1); Tensor destination_column = destination.slice(1, token, 1); - detail::sparse_moe_decode_launch(x_column, weights, destination_column, views, stream); + detail::sparse_moe_decode_launch(x_column, weights, destination_column, plan, views, + stream); } } diff --git a/src/targets/qwen3_6/impl/runtime/text_context.h b/src/targets/qwen3_6/impl/runtime/text_context.h index 3b4deda322..42e5f17b9e 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context.h +++ b/src/targets/qwen3_6/impl/runtime/text_context.h @@ -10,6 +10,7 @@ #include "core/weight.h" #include "ninfer/ops/sampling.h" #include "ninfer/ops/softmax_attention.h" +#include "ninfer/ops/sparse_moe.h" #include #include #include @@ -248,7 +249,9 @@ class TextContext { [[nodiscard]] const MtpW& mtp_weights() const; void attn_mix(const FullLayerW& weights, Tensor& x, int index, Phase phase); void gdn_mix(const GdnLayerW& weights, Tensor& x, int index, Phase phase); - void mlp_tail(const Tensor* post_norm, const MlpW& weights, Tensor& x, Phase phase); + void mlp_tail(const Tensor* post_norm, const MlpW& weights, Tensor& x, Phase phase, + const ops::SparseMoeHints& hints); + [[nodiscard]] ops::SparseMoeHints next_projection_hints(int layer) const; void run_layers(Tensor& x, Phase phase); template void run_layers(Tensor& x, Phase phase, Tap& tap); diff --git a/src/targets/qwen3_6/impl/runtime/text_context_impl.h b/src/targets/qwen3_6/impl/runtime/text_context_impl.h index 872c13531a..8cdef782a1 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_context_impl.h @@ -24,6 +24,7 @@ #include "ninfer/ops/residual_add.h" #include "ninfer/ops/rmsnorm.h" #include "ninfer/ops/rope.h" +#include "ninfer/ops/sparse_moe.h" #include "ninfer/ops/scatter.h" #include "ninfer/ops/scalar.h" #include "ninfer/ops/sigmoid_mul.h" @@ -984,13 +985,27 @@ void TextContext::gdn_mix(const GdnLayerW& w, Tensor& x, int gidx, Phase ph) { Variant::gdn_output_projection(on.view({kCfg.value_dim, T}), *w.out_proj, x, ph, work_, s); } -void TextContext::mlp_tail(const Tensor* post_norm, const MlpW& m, Tensor& x, Phase ph) { +ops::SparseMoeHints TextContext::next_projection_hints(int layer) const { + // Name the next layer's projection codes so this layer's MoE D4 epilogue can warm L2 for + // them. The last layer names nothing. A pure hint: the value never reaches arithmetic. + const int next = layer + 1; + if (next >= kCfg.n_layers) { return {}; } + if (ModelConfig::is_full(next)) { + return Variant::projection_prefetch_hints( + *full_.at(static_cast(ModelConfig::full_idx(next))).projection); + } + return Variant::projection_prefetch_hints( + *gdn_.at(static_cast(ModelConfig::gdn_idx(next))).projection); +} + +void TextContext::mlp_tail(const Tensor* post_norm, const MlpW& m, Tensor& x, Phase ph, + const ops::SparseMoeHints& hints) { cudaStream_t s = ctx_.stream; const int T = x.ne[1]; Tensor h = workspace_recipe::post_mixer_hidden(work_, T); ops::rmsnorm(x, *post_norm, kCfg.rms_eps, true, h, s); - Variant::post_mixer(h, *m.payload, x, ph, work_, s); + Variant::post_mixer(h, *m.payload, x, ph, hints, work_, s); } template @@ -1015,7 +1030,7 @@ void TextContext::run_layers(Tensor& x, Phase ph, Tap& tap) { prefill ? nvtx::Name::PrefillPostMixer : nvtx::Name::VerifyPostMixer, nvtx::Category::PostMixer, static_cast(layer)); auto mlp_scope = work_.scope(); - mlp_tail(full.post_attn_norm, full.mlp, x, ph); + mlp_tail(full.post_attn_norm, full.mlp, x, ph, next_projection_hints(layer)); if constexpr (Tap::enabled) { tap.capture_layer(layer, x, ctx_.stream); } } } else { @@ -1036,7 +1051,7 @@ void TextContext::run_layers(Tensor& x, Phase ph, Tap& tap) { prefill ? nvtx::Name::PrefillPostMixer : nvtx::Name::VerifyPostMixer, nvtx::Category::PostMixer, static_cast(layer)); auto mlp_scope = work_.scope(); - mlp_tail(gdn.post_attn_norm, gdn.mlp, x, ph); + mlp_tail(gdn.post_attn_norm, gdn.mlp, x, ph, next_projection_hints(layer)); if constexpr (Tap::enabled) { tap.capture_layer(layer, x, ctx_.stream); } } } diff --git a/src/targets/qwen3_6_27b/impl/variant.cpp b/src/targets/qwen3_6_27b/impl/variant.cpp index db107ce55b..66d5c957eb 100644 --- a/src/targets/qwen3_6_27b/impl/variant.cpp +++ b/src/targets/qwen3_6_27b/impl/variant.cpp @@ -301,7 +301,8 @@ void Variant::gdn_norm_control_projection(const Tensor& residual, const Tensor& } void Variant::post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, - qwen3_6::TextPhase, WorkspaceArena& workspace, cudaStream_t stream) { + qwen3_6::TextPhase, const ::ninfer::ops::SparseMoeHints&, + WorkspaceArena& workspace, cudaStream_t stream) { auto scope = workspace.scope(); Tensor activation = workspace.alloc(DType::BF16, {TextConfig::intermediate, hidden.ne[1]}); ops::linear_swiglu(hidden, weights.gate_up, activation, text_policy(weights.gate_up), workspace, diff --git a/src/targets/qwen3_6_27b/impl/variant.h b/src/targets/qwen3_6_27b/impl/variant.h index b9c60984a7..d871054e06 100644 --- a/src/targets/qwen3_6_27b/impl/variant.h +++ b/src/targets/qwen3_6_27b/impl/variant.h @@ -2,6 +2,7 @@ #include "core/device.h" #include "targets/qwen3_6_27b/impl/config.h" +#include "ninfer/ops/sparse_moe.h" #include "targets/qwen3_6_27b/impl/load/bindings.h" #include @@ -29,6 +30,16 @@ struct Variant { using VisionWeights = qwen3_6::VisionWeights; using GraphExecutionProfile = detail::GraphExecutionProfile; + // The dense post-mixer streams no MoE weights, so this target names no prefetch span. + static ::ninfer::ops::SparseMoeHints + projection_prefetch_hints(const FullAttentionProjectionWeights&) { + return {}; + } + + static ::ninfer::ops::SparseMoeHints projection_prefetch_hints(const GdnProjectionWeights&) { + return {}; + } + static constexpr float attention_scale = kAttentionScale; static constexpr float gdn_scale = kGdnScale; static constexpr std::uint32_t prefill_chunk_alignment = kPrefillChunkAlignment; @@ -80,8 +91,8 @@ struct Variant { WorkspaceArena& workspace, DeviceExecutionView execution); static void post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, - qwen3_6::TextPhase phase, WorkspaceArena& workspace, - cudaStream_t stream); + qwen3_6::TextPhase phase, const ::ninfer::ops::SparseMoeHints& hints, + WorkspaceArena& workspace, cudaStream_t stream); static void mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights, Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream); [[nodiscard]] static std::size_t diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp index bf9f11fe6f..f5455be8c6 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp @@ -64,13 +64,14 @@ bool dflash_target_uses_chunked_small_t(std::uint32_t draft_window, std::uint32_ } void run_sparse_moe(const Tensor& hidden, const ops::SparseMoeWeights& weights, Tensor& residual, - WorkspaceArena& workspace, cudaStream_t stream) { + const ops::SparseMoeHints& hints, WorkspaceArena& workspace, + cudaStream_t stream) { auto scope = workspace.scope(); const DeviceSpan storage = workspace.alloc_bytes(ops::sparse_moe_workspace_capacity_bytes( weights.routed_gate_up.qtype, weights.routed_down.qtype, hidden.ne[1], hidden.ne[1])); WorkspaceArena leaf_workspace(storage); - ops::sparse_moe(hidden, weights, ops::SparseMoeEpilogue::AddResidual, residual, leaf_workspace, - stream); + ops::sparse_moe(hidden, weights, ops::SparseMoeEpilogue::AddResidual, residual, hints, + leaf_workspace, stream); } void validate_token_interval(std::int32_t first, std::int32_t last) { @@ -214,13 +215,14 @@ void Variant::gdn_norm_control_projection(const Tensor& residual, const Tensor& } void Variant::post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, - qwen3_6::TextPhase, WorkspaceArena& workspace, cudaStream_t stream) { - run_sparse_moe(hidden, weights.op, residual, workspace, stream); + qwen3_6::TextPhase, const ::ninfer::ops::SparseMoeHints& hints, + WorkspaceArena& workspace, cudaStream_t stream) { + run_sparse_moe(hidden, weights.op, residual, hints, workspace, stream); } void Variant::mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights, Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream) { - run_sparse_moe(hidden, weights.op, residual, workspace, stream); + run_sparse_moe(hidden, weights.op, residual, ops::SparseMoeHints{}, workspace, stream); } std::size_t Variant::mtp_attention_projection_workspace_capacity_bytes(std::int32_t first, diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.h b/src/targets/qwen3_6_35b_a3b/impl/variant.h index 532fc82633..602c701bc0 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.h +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.h @@ -2,6 +2,7 @@ #include "core/device.h" #include "targets/qwen3_6_35b_a3b/impl/config.h" +#include "ninfer/ops/sparse_moe.h" #include "targets/qwen3_6_35b_a3b/impl/load/bindings.h" #include @@ -27,6 +28,16 @@ struct Variant { using VisionWeights = qwen3_6::VisionWeights; using GraphExecutionProfile = detail::GraphExecutionProfile; + static ::ninfer::ops::SparseMoeHints + projection_prefetch_hints(const FullAttentionProjectionWeights& weights) { + return {weights.query_key_gate_value.qdata, std::size_t{9216} * 2048}; + } + + static ::ninfer::ops::SparseMoeHints + projection_prefetch_hints(const GdnProjectionWeights& weights) { + return {weights.query_key_value_z.qdata, std::size_t{12288} * 2048}; + } + static constexpr float attention_scale = kAttentionScale; static constexpr float gdn_scale = kGdnScale; static constexpr std::uint32_t prefill_chunk_alignment = kPrefillChunkAlignment; @@ -86,8 +97,8 @@ struct Variant { WorkspaceArena& workspace, DeviceExecutionView execution); static void post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, - qwen3_6::TextPhase phase, WorkspaceArena& workspace, - cudaStream_t stream); + qwen3_6::TextPhase phase, const ::ninfer::ops::SparseMoeHints& hints, + WorkspaceArena& workspace, cudaStream_t stream); static void mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights, Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream);