Remove the empty graph nodes of the MTP draft phase: +0.33% decode on Qwen3.6-35B-A3B

include/ninfer/ops/mtp_pack.h

@@ -29,6 +29,73 @@ namespace ninfer::ops {
 void mtp_pack_fc_input(const Tensor& embedding_norm, const Tensor& hidden_norm, Tensor& out,
                        cudaStream_t stream);
 
+/**
+ * Op: mtp_norm_pack_fc_input
+ *
+ * Math / indexing:
+ *   out[0:D, t]  = rmsnorm(embedding[:, t], embedding_weight)
+ *   out[D:2D, t] = rmsnorm(hidden[:, t], hidden_weight)
+ *   with the unit-offset weight convention, i.e. the weight used is (1 + w).
+ *
+ * Logical shapes:
+ *   BF16 embedding and hidden [D,T], weights [D], out [2D,T], all contiguous.
+ *
+ * Numeric:
+ *   Each half is normalised over its own D values: out = x * rsqrt(mean(x^2) + eps) * (1 + w),
+ *   per column and per half. eps must be positive and finite; anything else is rejected. Inputs
+ *   and weights are read as represented BF16 and every output element is rounded to nearest BF16
+ *   once -- that rounding is the only observable boundary. Conformance is against an FP64 oracle
+ *   of the formula above under the operation's numerical criterion; reduction order and every
+ *   intermediate width are implementation choices. mtp_norm_pack_fc_input_supported() reports the
+ *   shapes this route covers, and the caller keeps the three-Op path elsewhere.
+ *
+ * Effects:
+ *   Writes the full output. Inputs, weights and output must not alias.
+ *
+ * Workspace:
+ *   None. The Op has no persistent state side effect.
+ */
+[[nodiscard]] bool mtp_norm_pack_fc_input_supported(const Tensor& embedding,
+                                                    const Tensor& embedding_weight,
+                                                    const Tensor& hidden,
+                                                    const Tensor& hidden_weight, const Tensor& out);
+
+void mtp_norm_pack_fc_input(const Tensor& embedding, const Tensor& embedding_weight,
+                            const Tensor& hidden, const Tensor& hidden_weight, Tensor& out,
+                            float eps, cudaStream_t stream);
+
+/**
+ * Op: mtp_residual_norm
+ *
+ * Math / indexing:
+ *   residual[:, t] += delta[:, t]           (pairwise, FP32 sum, round to nearest BF16)
+ *   out[:, t]       = rmsnorm(residual[:, t], weight)   with the unit-offset weight convention.
+ *
+ * Logical shapes:
+ *   BF16 delta and residual [D,T], weight [D], out [D,T], all contiguous.
+ *
+ * Numeric:
+ *   The residual becomes the sum of the two represented BF16 values, rounded to nearest BF16 once.
+ *   out is the RMSNorm of that stored residual, out = r * rsqrt(mean(r^2) + eps) * (1 + w), read
+ *   back as represented BF16 and rounded to nearest BF16 once. eps must be positive and finite.
+ *   Those two roundings and the stored residual are the observable boundaries; conformance is
+ *   against an FP64 oracle of the composition, in which the sum is exact because both addends are
+ *   BF16. mtp_residual_norm_supported() reports the shapes this route covers; the caller keeps the
+ *   two-Op path elsewhere.
+ *
+ * Effects:
+ *   Updates the full residual in place and writes the full output. delta, weight and out must not
+ *   alias the residual.
+ *
+ * Workspace:
+ *   None. The Op has no persistent state side effect.
+ */
+[[nodiscard]] bool mtp_residual_norm_supported(const Tensor& delta, const Tensor& residual,
+                                               const Tensor& weight, const Tensor& out);
+
+void mtp_residual_norm(const Tensor& delta, Tensor& residual, const Tensor& weight, Tensor& out,
+                       float eps, cudaStream_t stream);
+
 /**
  * Op: mtp_split_attn_in
  *

src/ops/kernel/argmax.cuh

@@ -92,6 +92,21 @@ __launch_bounds__(kArgmaxBlock) __global__
     if (threadIdx.x == 0) { out[t] = indices[0]; }
 }
 
+// The tiled route opens its atomic contest from row 0, so every winner slot in the slice has
+// to hold 0 before it starts. A memset and this kernel write the same bytes, but they do not cost
+// the same as graph nodes: on sm_120a the memset node measured 4.6-8.4 us of decode-graph critical
+// path against 1.1 us for this kernel, and the node count is unchanged. Numbers in the commit
+// message.
+inline constexpr int kArgmaxResetBlock = 256;
+
+__launch_bounds__(kArgmaxResetBlock) __global__
+    void argmax_reset_winners_kernel(std::int32_t* out, std::int32_t count) {
+    const std::int32_t i =
+        static_cast<std::int32_t>(blockIdx.x) * static_cast<std::int32_t>(blockDim.x) +
+        static_cast<std::int32_t>(threadIdx.x);
+    if (i < count) { out[i] = 0; }
+}
+
 __launch_bounds__(kArgmaxBlock) __global__
     void argmax_tiled_atomic_kernel(const __nv_bfloat16* logits, std::int32_t* out,
                                     std::int32_t valid_rows, std::int32_t physical_rows) {

src/ops/kernel/mtp_pack.cuh

@@ -1,5 +1,7 @@
 #pragma once
 
+#include "ops/kernel/rmsnorm.cuh"
+
 #include <cuda_bf16.h>
 
 #include <cstdint>
@@ -23,6 +25,123 @@ __global__ void mtp_pack_fc_input_kernel(const __nv_bfloat16* embedding_norm,
     out[out_base + rows + row]  = hidden_norm[in_idx];
 }
 
+// Both stem norms in one launch, written straight into the packed FC input. blockIdx.y picks
+// the half: 0 is the embedding, 1 is the hidden state. A block does exactly what the standalone
+// RMSNorm CTA kernel did for its row -- same block width, same pair decomposition, same reduction,
+// same epilogue -- so the bytes are the bytes the two separate launches wrote, and the copy that
+// used to move them into place is gone with the two launches.
+template <RmsEpilogue Epilogue, int Block, int MaxPairsPerThread, bool Prefetch>
+__launch_bounds__(Block) __global__
+    void mtp_norm_pack_fc_input_kernel(const __nv_bfloat162* embedding,
+                                       const __nv_bfloat162* embedding_weight,
+                                       const __nv_bfloat162* hidden,
+                                       const __nv_bfloat162* hidden_weight, __nv_bfloat162* out,
+                                       std::int32_t d, std::int64_t rows, float eps) {
+    static_assert(Block % kWarpSize == 0);
+    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x);
+    if (row >= rows) { return; }
+    const int half          = static_cast<int>(blockIdx.y);
+    const __nv_bfloat162* x = half == 0 ? embedding : hidden;
+    const __nv_bfloat162* w = half == 0 ? embedding_weight : hidden_weight;
+
+    const int pairs             = d / 2;
+    const int pairs_per_thread  = pairs / Block;
+    const std::int64_t row_base = row * static_cast<std::int64_t>(pairs);
+    const std::int64_t out_base =
+        row * static_cast<std::int64_t>(2 * pairs) + static_cast<std::int64_t>(half) * pairs;
+    __nv_bfloat162 values[MaxPairsPerThread];
+    __nv_bfloat162 weights[MaxPairsPerThread];
+    float sum = 0.0f;
+
+#pragma unroll
+    for (int k = 0; k < MaxPairsPerThread; ++k) {
+        if (k < pairs_per_thread) {
+            const int pair = static_cast<int>(threadIdx.x) + k * Block;
+            values[k]      = x[row_base + pair];
+            if constexpr (Prefetch) { weights[k] = w[pair]; }
+            const float2 xf = __bfloat1622float2(values[k]);
+            sum += xf.x * xf.x + xf.y * xf.y;
+        }
+    }
+
+    __shared__ float warp_sums[Block / kWarpSize];
+    __shared__ float inv_shared;
+    const float block_sum = block_reduce_sum<Block>(sum, warp_sums);
+    if (threadIdx.x == 0) { inv_shared = rsqrtf(block_sum / static_cast<float>(d) + eps); }
+    __syncthreads();
+    const float inv = inv_shared;
+
+#pragma unroll
+    for (int k = 0; k < MaxPairsPerThread; ++k) {
+        if (k < pairs_per_thread) {
+            const int pair              = static_cast<int>(threadIdx.x) + k * Block;
+            const float2 xf             = __bfloat1622float2(values[k]);
+            const __nv_bfloat162 w_pair = Prefetch ? weights[k] : w[pair];
+            const float2 wf             = __bfloat1622float2(w_pair);
+            out[out_base + pair] =
+                __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(xf.x, inv, wf.x, 0.0f),
+                                      rmsnorm_epilogue<Epilogue>(xf.y, inv, wf.y, 0.0f));
+        }
+    }
+}
+
+// The residual update and the norm that reads it back are one pass over the same 2048 values.
+// A block adds its row exactly as residual_add does -- pairwise, FP32, round to nearest -- stores
+// the result, and normalises the stored values with the same reduction the standalone RMSNorm CTA
+// kernel uses, so both outputs are the bytes the two Ops wrote.
+template <RmsEpilogue Epilogue, int Block, int MaxPairsPerThread, bool Prefetch>
+__launch_bounds__(Block) __global__
+    void mtp_residual_norm_kernel(const __nv_bfloat162* delta, __nv_bfloat162* residual,
+                                  const __nv_bfloat162* weight, __nv_bfloat162* out, std::int32_t d,
+                                  std::int64_t rows, float eps) {
+    static_assert(Block % kWarpSize == 0);
+    const std::int64_t row = static_cast<std::int64_t>(blockIdx.x);
+    if (row >= rows) { return; }
+
+    const int pairs             = d / 2;
+    const int pairs_per_thread  = pairs / Block;
+    const std::int64_t row_base = row * static_cast<std::int64_t>(pairs);
+    __nv_bfloat162 values[MaxPairsPerThread];
+    __nv_bfloat162 weights[MaxPairsPerThread];
+    float sum = 0.0f;
+
+#pragma unroll
+    for (int k = 0; k < MaxPairsPerThread; ++k) {
+        if (k < pairs_per_thread) {
+            const int pair                  = static_cast<int>(threadIdx.x) + k * Block;
+            const __nv_bfloat162 delta_pair = delta[row_base + pair];
+            const __nv_bfloat162 x_pair     = residual[row_base + pair];
+            const float low                 = __low2float(x_pair) + __low2float(delta_pair);
+            const float high                = __high2float(x_pair) + __high2float(delta_pair);
+            values[k]                       = __floats2bfloat162_rn(low, high);
+            residual[row_base + pair]       = values[k];
+            if constexpr (Prefetch) { weights[k] = weight[pair]; }
+            const float2 xf = __bfloat1622float2(values[k]);
+            sum += xf.x * xf.x + xf.y * xf.y;
+        }
+    }
+
+    __shared__ float warp_sums[Block / kWarpSize];
+    __shared__ float inv_shared;
+    const float block_sum = block_reduce_sum<Block>(sum, warp_sums);
+    if (threadIdx.x == 0) { inv_shared = rsqrtf(block_sum / static_cast<float>(d) + eps); }
+    __syncthreads();
+    const float inv = inv_shared;
+
+#pragma unroll
+    for (int k = 0; k < MaxPairsPerThread; ++k) {
+        if (k < pairs_per_thread) {
+            const int pair              = static_cast<int>(threadIdx.x) + k * Block;
+            const float2 xf             = __bfloat1622float2(values[k]);
+            const __nv_bfloat162 w_pair = Prefetch ? weights[k] : weight[pair];
+            const float2 wf             = __bfloat1622float2(w_pair);
+            out[row_base + pair] =
+                __floats2bfloat162_rn(rmsnorm_epilogue<Epilogue>(xf.x, inv, wf.x, 0.0f),
+                                      rmsnorm_epilogue<Epilogue>(xf.y, inv, wf.y, 0.0f));
+        }
+    }
+}
+
 __global__ void mtp_split_attn_in_kernel(const __nv_bfloat16* attn_in, __nv_bfloat16* q,
                                          __nv_bfloat16* k, __nv_bfloat16* gate, __nv_bfloat16* v,
                                          std::int32_t tokens) {

src/ops/launcher/argmax.cu

@@ -67,9 +67,12 @@ void argmax_tiled_atomic_launch(const Tensor& logits, Tensor& out, std::int32_t
     for_each_token_slice(t_count, 1, [&](int token_offset, int token_count) {
         const Tensor logits_slice = logits.slice(1, token_offset, token_count);
         Tensor out_slice          = out.slice(0, token_offset, token_count);
-        CUDA_CHECK(cudaMemsetAsync(out_slice.data, 0,
-                                   static_cast<std::size_t>(token_count) * sizeof(std::int32_t),
-                                   stream));
+        // token_count is a whole slice, not a single column: for_each_token_slice hands
+        // out up to the grid.y limit at once, so the reset has to cover all of it.
+        argmax_reset_winners_kernel<<<
+            static_cast<unsigned int>(div_up(token_count, kArgmaxResetBlock)), kArgmaxResetBlock, 0,
+            stream>>>(static_cast<std::int32_t*>(out_slice.data), token_count);
+        CUDA_CHECK(cudaGetLastError());
         const dim3 grid(static_cast<unsigned int>(tiled_blocks),
                         static_cast<unsigned int>(token_count));
         argmax_tiled_atomic_kernel<<<grid, block, 0, stream>>>(

src/ops/launcher/mtp_pack.cu

@@ -7,6 +7,7 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <stdexcept>
 
 namespace ninfer::ops::detail {
 
@@ -27,6 +28,116 @@ void mtp_pack_fc_input_launch(const Tensor& embedding_norm, const Tensor& hidden
     });
 }
 
+namespace {
+
+// A fused row kernel is bit-exact with ops::rmsnorm only if it lands on the SAME
+// instantiation the standalone launcher would pick for that width: a different block width or a
+// different pairs-per-thread splits the sum of squares differently, and FP32 addition is not
+// associative. This mirrors the ladder in src/ops/launcher/rmsnorm.cu for the Offset epilogue,
+// including its first branch, which intercepts d == 5120 before everything else. Widths the
+// standalone launcher sends to its warp or generic kernels are declined here.
+enum class MtpRowRoute { None, Cta256x6, Cta256x10 };
+
+// Only the two registered MTP stem widths are admitted: 2048 for Qwen3.6-27B and 5120 for
+// 35B-A3B. Everything else stays on the composed three-Op path, which serves it correctly --
+// admitting a width no target asks for would be a route kept alive for a hypothetical one. Each
+// admitted width lands on the instantiation src/ops/launcher/rmsnorm.cu would have picked for it,
+// which is what makes the fused row bit-exact with ops::rmsnorm.
+MtpRowRoute mtp_row_route(std::int32_t d) {
+    if (d == 5120) { return MtpRowRoute::Cta256x10; }
+    if (d == 2048) { return MtpRowRoute::Cta256x6; }
+    return MtpRowRoute::None;
+}
+
+// The route table and the kernel's compile-time bound have to agree; if a future width is added to
+// one and not the other the kernel would silently drop the trailing pairs from the sum and from
+// the store. Checked rather than assumed.
+void require_row_route(std::int32_t d, int block, int max_pairs) {
+    const int pairs = d / 2;
+    if (pairs % block != 0 || pairs / block < 1 || pairs / block > max_pairs) {
+        throw std::invalid_argument("mtp fused row kernel: width outside the instantiated route");
+    }
+}
+
+bool mtp_row_route_admits(std::int32_t d, std::uintptr_t pointer_bits) {
+    if ((pointer_bits & (alignof(__nv_bfloat162) - 1)) != 0) { return false; }
+    return mtp_row_route(d) != MtpRowRoute::None;
+}
+
+} // namespace
+
+bool mtp_norm_pack_fc_input_admits(std::int32_t d, const Tensor& embedding,
+                                   const Tensor& embedding_weight, const Tensor& hidden,
+                                   const Tensor& hidden_weight, const Tensor& out) {
+    return mtp_row_route_admits(d, reinterpret_cast<std::uintptr_t>(embedding.data) |
+                                       reinterpret_cast<std::uintptr_t>(embedding_weight.data) |
+                                       reinterpret_cast<std::uintptr_t>(hidden.data) |
+                                       reinterpret_cast<std::uintptr_t>(hidden_weight.data) |
+                                       reinterpret_cast<std::uintptr_t>(out.data));
+}
+
+void mtp_norm_pack_fc_input_launch(const Tensor& embedding, const Tensor& embedding_weight,
+                                   const Tensor& hidden, const Tensor& hidden_weight, Tensor& out,
+                                   float eps, cudaStream_t stream) {
+    const std::int32_t d    = embedding.ne[0];
+    const std::int64_t rows = embedding.ne[1];
+    const dim3 grid(static_cast<unsigned int>(rows), 2);
+    const auto* e  = reinterpret_cast<const __nv_bfloat162*>(embedding.data);
+    const auto* ew = reinterpret_cast<const __nv_bfloat162*>(embedding_weight.data);
+    const auto* h  = reinterpret_cast<const __nv_bfloat162*>(hidden.data);
+    const auto* hw = reinterpret_cast<const __nv_bfloat162*>(hidden_weight.data);
+    auto* o        = reinterpret_cast<__nv_bfloat162*>(out.data);
+    switch (mtp_row_route(d)) {
+    case MtpRowRoute::Cta256x6:
+        require_row_route(d, 256, 6);
+        mtp_norm_pack_fc_input_kernel<RmsEpilogue::Offset, 256, 6, true>
+            <<<grid, 256, 0, stream>>>(e, ew, h, hw, o, d, rows, eps);
+        break;
+    case MtpRowRoute::Cta256x10:
+        require_row_route(d, 256, 10);
+        mtp_norm_pack_fc_input_kernel<RmsEpilogue::Offset, 256, 10, true>
+            <<<grid, 256, 0, stream>>>(e, ew, h, hw, o, d, rows, eps);
+        break;
+    case MtpRowRoute::None:
+        throw std::invalid_argument("mtp_norm_pack_fc_input: unsupported width");
+    }
+    CUDA_CHECK(cudaGetLastError());
+}
+
+bool mtp_residual_norm_admits(std::int32_t d, const Tensor& delta, const Tensor& residual,
+                              const Tensor& weight, const Tensor& out) {
+    return mtp_row_route_admits(d, reinterpret_cast<std::uintptr_t>(delta.data) |
+                                       reinterpret_cast<std::uintptr_t>(residual.data) |
+                                       reinterpret_cast<std::uintptr_t>(weight.data) |
+                                       reinterpret_cast<std::uintptr_t>(out.data));
+}
+
+void mtp_residual_norm_launch(const Tensor& delta, Tensor& residual, const Tensor& weight,
+                              Tensor& out, float eps, cudaStream_t stream) {
+    const std::int32_t d    = residual.ne[0];
+    const std::int64_t rows = residual.ne[1];
+    const auto* dv          = reinterpret_cast<const __nv_bfloat162*>(delta.data);
+    auto* rv                = reinterpret_cast<__nv_bfloat162*>(residual.data);
+    const auto* wv          = reinterpret_cast<const __nv_bfloat162*>(weight.data);
+    auto* ov                = reinterpret_cast<__nv_bfloat162*>(out.data);
+    const auto grid         = static_cast<unsigned int>(rows);
+    switch (mtp_row_route(d)) {
+    case MtpRowRoute::Cta256x6:
+        require_row_route(d, 256, 6);
+        mtp_residual_norm_kernel<RmsEpilogue::Offset, 256, 6, true>
+            <<<grid, 256, 0, stream>>>(dv, rv, wv, ov, d, rows, eps);
+        break;
+    case MtpRowRoute::Cta256x10:
+        require_row_route(d, 256, 10);
+        mtp_residual_norm_kernel<RmsEpilogue::Offset, 256, 10, true>
+            <<<grid, 256, 0, stream>>>(dv, rv, wv, ov, d, rows, eps);
+        break;
+    case MtpRowRoute::None:
+        throw std::invalid_argument("mtp_residual_norm: unsupported width");
+    }
+    CUDA_CHECK(cudaGetLastError());
+}
+
 void mtp_split_attn_in_launch(const Tensor& attn_in, Tensor& q, Tensor& k, Tensor& gate, Tensor& v,
                               cudaStream_t stream) {
     constexpr int kBlock = 256;

src/ops/launcher/mtp_pack.h

@@ -9,6 +9,20 @@ namespace ninfer::ops::detail {
 void mtp_pack_fc_input_launch(const Tensor& embedding_norm, const Tensor& hidden_norm, Tensor& out,
                               cudaStream_t stream);
 
+bool mtp_norm_pack_fc_input_admits(std::int32_t d, const Tensor& embedding,
+                                   const Tensor& embedding_weight, const Tensor& hidden,
+                                   const Tensor& hidden_weight, const Tensor& out);
+
+void mtp_norm_pack_fc_input_launch(const Tensor& embedding, const Tensor& embedding_weight,
+                                   const Tensor& hidden, const Tensor& hidden_weight, Tensor& out,
+                                   float eps, cudaStream_t stream);
+
+bool mtp_residual_norm_admits(std::int32_t d, const Tensor& delta, const Tensor& residual,
+                              const Tensor& weight, const Tensor& out);
+
+void mtp_residual_norm_launch(const Tensor& delta, Tensor& residual, const Tensor& weight,
+                              Tensor& out, float eps, cudaStream_t stream);
+
 void mtp_split_attn_in_launch(const Tensor& attn_in, Tensor& q, Tensor& k, Tensor& gate, Tensor& v,
                               cudaStream_t stream);