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);