Neroued · Neroued · Sep 8, 2026 · Sep 7, 2026
diff --git 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
@@ -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
@@ -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<int>(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<unsigned long long>(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 <class Codec>
 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<const int*>(workspace.ids.data);
     const auto* alpha         = static_cast<const float*>(workspace.alpha.data);
     const auto* shared_scale  = static_cast<const float*>(workspace.shared_scale.data);
@@ -539,23 +551,28 @@ void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destinat
     const auto* shared_codes  = static_cast<const std::uint8_t*>(weights.shared_down.qdata);
     const auto* shared_scales = static_cast<const std::uint8_t*>(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<Codec, 1>, 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<Codec, 1>, ids,
+        alpha, shared_scale, act, routed_codes, routed_high, routed_scales, shared_codes,
+        shared_scales, output, static_cast<const char*>(prefetch_data),
+        static_cast<unsigned long long>(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<Q5Codec>(weights, destination, workspace, stream);
+        launch_d4_dependent_codec<Q5Codec>(weights, destination, workspace, stream, prefetch_data,
+                                           prefetch_bytes);
         return;
     case QType::Q6G64_F16S:
-        launch_d4_dependent_codec<Q6Codec>(weights, destination, workspace, stream);
+        launch_d4_dependent_codec<Q6Codec>(weights, destination, workspace, stream, prefetch_data,
+                                           prefetch_bytes);
         return;
     case QType::W8G32_F16S:
-        launch_d4_dependent_codec<W8Codec>(weights, destination, workspace, stream);
+        launch_d4_dependent_codec<W8Codec>(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
@@ -5,14 +5,20 @@
 #include <stdexcept>
 
 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;
     (void)allocate_sparse_moe_decode_workspace(layout);
     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
@@ -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
@@ -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 <ninfer/targets/qwen3_6/decoder_state.h>
 #include <ninfer/targets/qwen3_6/prepared_prompt.h>
 #include <ninfer/targets/qwen3_6/round_state.h>
@@ -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 <class Tap>
     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
@@ -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<std::size_t>(ModelConfig::full_idx(next))).projection);
+    }
+    return Variant::projection_prefetch_hints(
+        *gdn_.at(static_cast<std::size_t>(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<TextConfig>(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 <class Tap>
@@ -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<std::uint64_t>(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<std::uint64_t>(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
@@ -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
@@ -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 <ninfer/targets/qwen3_6/runtime.h>
 
@@ -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