From 697606edea3aeb9b0ccffcd66ee4db513bf5079e Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Wed, 9 Sep 2026 19:24:12 +0200 Subject: [PATCH 1/5] perf(runtime): remove the empty graph nodes of the MTP draft phase The draft phase of an MTP decode round spends nodes that carry no work of their own. Three of them, in one pass because separately none is above what an end-to-end benchmark on this hardware can resolve: - ops::argmax opened its atomic contest with a 4-byte cudaMemsetAsync per slice; an initializer kernel writes the same bytes as a kernel node; - the MTP stem normalized the embedding and the hidden state only to lay them side by side, and the MTP tail added the residual only for the next norm to read it back: two new Ops write each result into place directly; - the autoregressive draft chain copied the new hidden state over the old one after every step, where alternating between two buffers costs nothing. Both new Ops carry a *_supported() predicate and the three separate calls remain as the fallback, so an unregistered shape keeps working. Over one traced run - 8192-token prefill, 64 generated tokens, two repetitions: 173 memset nodes become kernel nodes, 397 rmsnorm launches disappear along with 134 pack and 129 residual-add launches, and 82 of 100 device-to-device copies go away. End to end on Qwen3.6-35B-A3B, twelve mirrored passes, decode +0.334% median with ten of twelve positive; the two negative passes and one outlier of the zero-control arm coincide with a build running on the neighbouring container, which shares the host CPU quota. Output is byte-identical, ctest 114/114. Co-Authored-By: Claude Opus 5 --- include/ninfer/ops/mtp_pack.h | 61 ++++ src/ops/kernel/argmax.cuh | 15 + src/ops/kernel/mtp_pack.cuh | 119 +++++++ src/ops/launcher/argmax.cu | 9 +- src/ops/launcher/mtp_pack.cu | 117 +++++++ src/ops/launcher/mtp_pack.h | 14 + src/ops/wrapper/mtp_pack.cpp | 82 +++++ src/targets/qwen3_6/impl/runtime/mtp_impl.h | 22 +- .../qwen3_6/impl/runtime/text_context_impl.h | 28 +- tests/ops/test_mtp_pack.cpp | 307 ++++++++++++++++++ 10 files changed, 756 insertions(+), 18 deletions(-) diff --git a/include/ninfer/ops/mtp_pack.h b/include/ninfer/ops/mtp_pack.h index 083ee90e21..25a77863f0 100644 --- a/include/ninfer/ops/mtp_pack.h +++ b/include/ninfer/ops/mtp_pack.h @@ -29,6 +29,67 @@ 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: + * Identical to rmsnorm() applied to each half followed by mtp_pack_fc_input(): the same block + * width, pair decomposition, reduction and epilogue produce the same bytes. Returns false from + * mtp_norm_pack_fc_input_supported() for shapes it does not cover, and the caller keeps the + * three-Op path there. + * + * 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: + * Identical to residual_add() followed by rmsnorm() on the updated residual: the same pairwise + * FP32 addition and rounding, then the same block width, pair decomposition, reduction and + * epilogue. 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 * diff --git a/src/ops/kernel/argmax.cuh b/src/ops/kernel/argmax.cuh index b773223cc9..75b107b0c5 100644 --- a/src/ops/kernel/argmax.cuh +++ b/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(blockIdx.x) * static_cast(blockDim.x) + + static_cast(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) { diff --git a/src/ops/kernel/mtp_pack.cuh b/src/ops/kernel/mtp_pack.cuh index d724ab62ab..89ae46b7d8 100644 --- a/src/ops/kernel/mtp_pack.cuh +++ b/src/ops/kernel/mtp_pack.cuh @@ -1,5 +1,7 @@ #pragma once +#include "ops/kernel/rmsnorm.cuh" + #include #include @@ -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 +__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(blockIdx.x); + if (row >= rows) { return; } + const int half = static_cast(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(pairs); + const std::int64_t out_base = + row * static_cast(2 * pairs) + static_cast(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(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(sum, warp_sums); + if (threadIdx.x == 0) { inv_shared = rsqrtf(block_sum / static_cast(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(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(xf.x, inv, wf.x, 0.0f), + rmsnorm_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 +__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(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(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(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(sum, warp_sums); + if (threadIdx.x == 0) { inv_shared = rsqrtf(block_sum / static_cast(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(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(xf.x, inv, wf.x, 0.0f), + rmsnorm_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) { diff --git a/src/ops/launcher/argmax.cu b/src/ops/launcher/argmax.cu index 416e3e2978..216a32dbd2 100644 --- a/src/ops/launcher/argmax.cu +++ b/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(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(div_up(token_count, kArgmaxResetBlock)), kArgmaxResetBlock, 0, + stream>>>(static_cast(out_slice.data), token_count); + CUDA_CHECK(cudaGetLastError()); const dim3 grid(static_cast(tiled_blocks), static_cast(token_count)); argmax_tiled_atomic_kernel<<>>( diff --git a/src/ops/launcher/mtp_pack.cu b/src/ops/launcher/mtp_pack.cu index d8d0b59764..af7ba04334 100644 --- a/src/ops/launcher/mtp_pack.cu +++ b/src/ops/launcher/mtp_pack.cu @@ -7,6 +7,7 @@ #include #include +#include namespace ninfer::ops::detail { @@ -27,6 +28,122 @@ 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, Cta512x8 }; + +MtpRowRoute mtp_row_route(std::int32_t d) { + if (d == 5120) { return MtpRowRoute::Cta256x10; } + if (d >= 512 && d <= 3072 && d % 512 == 0) { return MtpRowRoute::Cta256x6; } + if (d > 3072 && d <= 8192 && d % 1024 == 0) { return MtpRowRoute::Cta512x8; } + 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(embedding.data) | + reinterpret_cast(embedding_weight.data) | + reinterpret_cast(hidden.data) | + reinterpret_cast(hidden_weight.data) | + reinterpret_cast(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(rows), 2); + const auto* e = reinterpret_cast(embedding.data); + const auto* ew = reinterpret_cast(embedding_weight.data); + const auto* h = reinterpret_cast(hidden.data); + const auto* hw = reinterpret_cast(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 + <<>>(e, ew, h, hw, o, d, rows, eps); + break; + case MtpRowRoute::Cta256x10: + require_row_route(d, 256, 10); + mtp_norm_pack_fc_input_kernel + <<>>(e, ew, h, hw, o, d, rows, eps); + break; + case MtpRowRoute::Cta512x8: + require_row_route(d, 512, 8); + mtp_norm_pack_fc_input_kernel + <<>>(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(delta.data) | + reinterpret_cast(residual.data) | + reinterpret_cast(weight.data) | + reinterpret_cast(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(delta.data); + auto* rv = reinterpret_cast<__nv_bfloat162*>(residual.data); + const auto* wv = reinterpret_cast(weight.data); + auto* ov = reinterpret_cast<__nv_bfloat162*>(out.data); + const auto grid = static_cast(rows); + switch (mtp_row_route(d)) { + case MtpRowRoute::Cta256x6: + require_row_route(d, 256, 6); + mtp_residual_norm_kernel + <<>>(dv, rv, wv, ov, d, rows, eps); + break; + case MtpRowRoute::Cta256x10: + require_row_route(d, 256, 10); + mtp_residual_norm_kernel + <<>>(dv, rv, wv, ov, d, rows, eps); + break; + case MtpRowRoute::Cta512x8: + require_row_route(d, 512, 8); + mtp_residual_norm_kernel + <<>>(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; diff --git a/src/ops/launcher/mtp_pack.h b/src/ops/launcher/mtp_pack.h index dd81a8d720..d7c655a5fe 100644 --- a/src/ops/launcher/mtp_pack.h +++ b/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); diff --git a/src/ops/wrapper/mtp_pack.cpp b/src/ops/wrapper/mtp_pack.cpp index 3015724a2c..ff796fbc4f 100644 --- a/src/ops/wrapper/mtp_pack.cpp +++ b/src/ops/wrapper/mtp_pack.cpp @@ -45,6 +45,88 @@ void mtp_pack_fc_input(const Tensor& embedding_norm, const Tensor& hidden_norm, detail::mtp_pack_fc_input_launch(embedding_norm, hidden_norm, out, stream); } +namespace { + +void require_stem_operands(const Tensor& embedding, const Tensor& embedding_weight, + const Tensor& hidden, const Tensor& hidden_weight, const Tensor& out) { + constexpr const char* op = "mtp_norm_pack_fc_input"; + require_bf16_contiguous_nonnull(embedding, op, "embedding"); + require_bf16_contiguous_nonnull(embedding_weight, op, "embedding_weight"); + require_bf16_contiguous_nonnull(hidden, op, "hidden"); + require_bf16_contiguous_nonnull(hidden_weight, op, "hidden_weight"); + require_bf16_contiguous_nonnull(out, op, "out"); + const std::int32_t rows = embedding.ne[0]; + const std::int32_t tokens = embedding.ne[1]; + if (rows <= 0) { throw std::invalid_argument("mtp_norm_pack_fc_input: D must be positive"); } + if (tokens <= 0) { throw std::invalid_argument("mtp_norm_pack_fc_input: T must be positive"); } + require_shape(embedding, rows, tokens, op, "embedding"); + require_shape(hidden, rows, tokens, op, "hidden"); + require_shape(out, 2 * rows, tokens, op, "out"); + if (embedding_weight.numel() != rows || hidden_weight.numel() != rows) { + throw std::invalid_argument("mtp_norm_pack_fc_input: weight length must be D"); + } +} + +} // namespace + +bool mtp_norm_pack_fc_input_supported(const Tensor& embedding, const Tensor& embedding_weight, + const Tensor& hidden, const Tensor& hidden_weight, + const Tensor& out) { + require_stem_operands(embedding, embedding_weight, hidden, hidden_weight, out); + return detail::mtp_norm_pack_fc_input_admits(embedding.ne[0], embedding, embedding_weight, + hidden, hidden_weight, 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) { + require_stem_operands(embedding, embedding_weight, hidden, hidden_weight, out); + if (!detail::mtp_norm_pack_fc_input_admits(embedding.ne[0], embedding, embedding_weight, hidden, + hidden_weight, out)) { + throw std::invalid_argument("mtp_norm_pack_fc_input: unsupported operand profile"); + } + detail::mtp_norm_pack_fc_input_launch(embedding, embedding_weight, hidden, hidden_weight, out, + eps, stream); +} + +namespace { + +void require_residual_norm_operands(const Tensor& delta, const Tensor& residual, + const Tensor& weight, const Tensor& out) { + constexpr const char* op = "mtp_residual_norm"; + require_bf16_contiguous_nonnull(delta, op, "delta"); + require_bf16_contiguous_nonnull(residual, op, "residual"); + require_bf16_contiguous_nonnull(weight, op, "weight"); + require_bf16_contiguous_nonnull(out, op, "out"); + const std::int32_t rows = residual.ne[0]; + const std::int32_t tokens = residual.ne[1]; + if (rows <= 0) { throw std::invalid_argument("mtp_residual_norm: D must be positive"); } + if (tokens <= 0) { throw std::invalid_argument("mtp_residual_norm: T must be positive"); } + require_shape(delta, rows, tokens, op, "delta"); + require_shape(residual, rows, tokens, op, "residual"); + require_shape(out, rows, tokens, op, "out"); + if (weight.numel() != rows) { + throw std::invalid_argument("mtp_residual_norm: weight length must be D"); + } +} + +} // namespace + +bool mtp_residual_norm_supported(const Tensor& delta, const Tensor& residual, const Tensor& weight, + const Tensor& out) { + require_residual_norm_operands(delta, residual, weight, out); + return detail::mtp_residual_norm_admits(residual.ne[0], delta, residual, weight, out); +} + +void mtp_residual_norm(const Tensor& delta, Tensor& residual, const Tensor& weight, Tensor& out, + float eps, cudaStream_t stream) { + require_residual_norm_operands(delta, residual, weight, out); + if (!detail::mtp_residual_norm_admits(residual.ne[0], delta, residual, weight, out)) { + throw std::invalid_argument("mtp_residual_norm: unsupported operand profile"); + } + detail::mtp_residual_norm_launch(delta, residual, weight, out, eps, stream); +} + void mtp_split_attn_in(const Tensor& attn_in, Tensor& q, Tensor& k, Tensor& gate, Tensor& v, cudaStream_t stream) { constexpr const char* op = "mtp_split_attn_in"; diff --git a/src/targets/qwen3_6/impl/runtime/mtp_impl.h b/src/targets/qwen3_6/impl/runtime/mtp_impl.h index 987e23ef82..4bd527fe38 100644 --- a/src/targets/qwen3_6/impl/runtime/mtp_impl.h +++ b/src/targets/qwen3_6/impl/runtime/mtp_impl.h @@ -9,6 +9,7 @@ #include #include +#include namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule { void mtp_bridge_and_propose(PrefillContext& state, const Tensor& next_token, @@ -166,6 +167,14 @@ auto mtp_decode_batch_body(MtpBatchContext& state, std::int32_t batch_size, std: Tensor proposal_logits = frame.proposal_logits.slice(1, 0, batch_size); Tensor draft0 = next_drafts.slice(1, 0, 1).view({batch_size}); card.mtp_propose_batch(ar_hidden, proposal_logits, draft0); + // The autoregressive chain alternates between the two hidden buffers instead of + // copying the new state back over the old one after every step. This is safe only + // because nothing reads either buffer once the loop ends: the next round rewrites + // ar_hidden from scratch through speculative_select_accepted_hidden above. A future + // reader of the final state must take source_hidden, which names the buffer the last + // step wrote; for an even number of steps that is no longer ar_hidden. + Tensor source_hidden = ar_hidden; + Tensor destination_hidden = next_hidden; for (std::uint32_t step = 0; step + 1 < k; ++step) { Tensor previous = next_drafts.slice(1, static_cast(step), 1).view({batch_size}); @@ -177,15 +186,14 @@ auto mtp_decode_batch_body(MtpBatchContext& state, std::int32_t batch_size, std: .view({1, batch_size}); Tensor valid = ar_valid_columns.slice(1, static_cast(step), 1) .view({batch_size}); - Tensor previous_batch = previous.view({1, batch_size}); - Tensor hidden_batch = ar_hidden.view({TextConfig::hidden, 1, batch_size}); - Tensor next_hidden_batch = next_hidden.view({TextConfig::hidden, 1, batch_size}); + Tensor previous_batch = previous.view({1, batch_size}); + Tensor hidden_batch = source_hidden.view({TextConfig::hidden, 1, batch_size}); + Tensor next_hidden_batch = + destination_hidden.view({TextConfig::hidden, 1, batch_size}); card.mtp_forward_decode_batch(previous_batch, hidden_batch, position, rope, valid, mtp_rows, envelopes.ar[step], next_hidden_batch); - card.mtp_propose_batch(next_hidden, proposal_logits, next); - CUDA_CHECK(cudaMemcpyAsync(ar_hidden.data, next_hidden.data, ar_hidden.bytes(), - cudaMemcpyDeviceToDevice, - state.execution.device.stream)); + card.mtp_propose_batch(destination_hidden, proposal_logits, next); + std::swap(source_hidden, destination_hidden); } } 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 8cdef782a1..bdde2fca76 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_context_impl.h @@ -343,13 +343,20 @@ void TextContext::mtp_forward_stem(const Tensor& ids, const Tensor& hidden, ops::embedding(flat_ids, *embed_, emb, s); } - Tensor e = roots.normalized_embedding; - Tensor h = roots.normalized_hidden; - ops::rmsnorm(emb, *mtp_.pre_fc_norm_embedding, kCfg.rms_eps, true, e, s); - ops::rmsnorm(flat_hidden, *mtp_.pre_fc_norm_hidden, kCfg.rms_eps, true, h, s); - + Tensor e = roots.normalized_embedding; + Tensor h = roots.normalized_hidden; Tensor fc_in = roots.packed_input; - ops::mtp_pack_fc_input(e, h, fc_in, s); + // The two stem norms exist only to be laid side by side, so where the fused Op covers the + // shape it writes them into place directly and three graph nodes become one. + if (ops::mtp_norm_pack_fc_input_supported(emb, *mtp_.pre_fc_norm_embedding, flat_hidden, + *mtp_.pre_fc_norm_hidden, fc_in)) { + ops::mtp_norm_pack_fc_input(emb, *mtp_.pre_fc_norm_embedding, flat_hidden, + *mtp_.pre_fc_norm_hidden, fc_in, kCfg.rms_eps, s); + } else { + ops::rmsnorm(emb, *mtp_.pre_fc_norm_embedding, kCfg.rms_eps, true, e, s); + ops::rmsnorm(flat_hidden, *mtp_.pre_fc_norm_hidden, kCfg.rms_eps, true, h, s); + ops::mtp_pack_fc_input(e, h, fc_in, s); + } x = roots.residual; ops::linear(fc_in, *mtp_.fc, x, s); @@ -411,10 +418,15 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po const auto post = workspace_recipe::mtp_post_attention(work_, T); Tensor o = post.output; ops::linear(a.view({kCfg.q_size, T}), *mtp_.o_proj, o, s); - ops::residual_add(o, x, s); Tensor mh = post.post_mixer_hidden; - ops::rmsnorm(x, *mtp_.post_attn_norm, kCfg.rms_eps, true, mh, s); + // The residual update and the norm that reads it back are one pass over the same values. + if (ops::mtp_residual_norm_supported(o, x, *mtp_.post_attn_norm, mh)) { + ops::mtp_residual_norm(o, x, *mtp_.post_attn_norm, mh, kCfg.rms_eps, s); + } else { + ops::residual_add(o, x, s); + ops::rmsnorm(x, *mtp_.post_attn_norm, kCfg.rms_eps, true, mh, s); + } { auto post_mixer_scope = work_.scope(); diff --git a/tests/ops/test_mtp_pack.cpp b/tests/ops/test_mtp_pack.cpp index 1f8b9d08fb..aa76f1dbe6 100644 --- a/tests/ops/test_mtp_pack.cpp +++ b/tests/ops/test_mtp_pack.cpp @@ -1,4 +1,6 @@ #include "ninfer/ops/mtp_pack.h" +#include "ninfer/ops/residual_add.h" +#include "ninfer/ops/rmsnorm.h" #include "ops/op_tester.h" #include @@ -137,6 +139,288 @@ int split_case(std::int32_t tokens) { return failures; } +// The fused stem Op has to produce the bytes the three-Op path produced, not merely bytes +// close to them: it exists to remove two graph nodes, and anything that shifts a bit would show up +// as a different draft token. The reference here is the product path itself -- two rmsnorm calls +// and the pack -- run on the same inputs in the same process. +int norm_pack_case(std::int32_t hidden, std::int32_t tokens) { + const std::int32_t output_rows = 2 * hidden; + const std::size_t count = static_cast(hidden) * tokens; + std::vector embedding(count), hidden_state(count), weight_e(hidden), weight_h(hidden); + fill_uniform(embedding, 0x51ed'0001u, -4.0F, 4.0F); + fill_uniform(hidden_state, 0x51ed'0002u, -4.0F, 4.0F); + fill_uniform(weight_e, 0x51ed'0003u, -0.5F, 0.5F); + fill_uniform(weight_h, 0x51ed'0004u, -0.5F, 0.5F); + + auto pack16 = [](const std::vector& v) { + std::vector out(v.size()); + for (std::size_t i = 0; i < v.size(); ++i) { out[i] = f32_to_bf16(v[i]); } + return out; + }; + const auto embedding16 = pack16(embedding); + const auto hidden16 = pack16(hidden_state); + const auto weight_e16 = pack16(weight_e); + const auto weight_h16 = pack16(weight_h); + const std::size_t in_bytes = count * sizeof(std::uint16_t); + const std::size_t w_bytes = static_cast(hidden) * sizeof(std::uint16_t); + const std::size_t out_bytes = + static_cast(output_rows) * tokens * sizeof(std::uint16_t); + + GuardedDeviceBuffer d_embedding(in_bytes), d_hidden(in_bytes); + GuardedDeviceBuffer d_weight_e(w_bytes), d_weight_h(w_bytes); + GuardedDeviceBuffer d_norm_e(in_bytes), d_norm_h(in_bytes); + GuardedDeviceBuffer d_reference(out_bytes), d_fused(out_bytes); + d_embedding.copy_from_host(embedding16.data(), in_bytes); + d_hidden.copy_from_host(hidden16.data(), in_bytes); + d_weight_e.copy_from_host(weight_e16.data(), w_bytes); + d_weight_h.copy_from_host(weight_h16.data(), w_bytes); + d_reference.fill(0xcd); + d_fused.fill(0xcd); + + Tensor t_embedding(d_embedding.data(), DType::BF16, {hidden, tokens}); + Tensor t_hidden(d_hidden.data(), DType::BF16, {hidden, tokens}); + Tensor t_weight_e(d_weight_e.data(), DType::BF16, {hidden}); + Tensor t_weight_h(d_weight_h.data(), DType::BF16, {hidden}); + Tensor t_norm_e(d_norm_e.data(), DType::BF16, {hidden, tokens}); + Tensor t_norm_h(d_norm_h.data(), DType::BF16, {hidden, tokens}); + Tensor t_reference(d_reference.data(), DType::BF16, {output_rows, tokens}); + Tensor t_fused(d_fused.data(), DType::BF16, {output_rows, tokens}); + + constexpr float kStemEps = 1.0e-6F; + ops::rmsnorm(t_embedding, t_weight_e, kStemEps, true, t_norm_e, nullptr); + ops::rmsnorm(t_hidden, t_weight_h, kStemEps, true, t_norm_h, nullptr); + ops::mtp_pack_fc_input(t_norm_e, t_norm_h, t_reference, nullptr); + cuda_synchronize(); + + const std::string label = + "mtp_norm_pack_fc_input D=" + std::to_string(hidden) + " T=" + std::to_string(tokens); + int failures = 0; + if (!ops::mtp_norm_pack_fc_input_supported(t_embedding, t_weight_e, t_hidden, t_weight_h, + t_fused)) { + std::cout << label << ": route declined, three-Op path stands\n"; + return 0; + } + ops::mtp_norm_pack_fc_input(t_embedding, t_weight_e, t_hidden, t_weight_h, t_fused, kStemEps, + nullptr); + cuda_synchronize(); + + const std::size_t elements = static_cast(output_rows) * tokens; + const auto reference = from_device(d_reference.data(), elements); + const auto fused = from_device(d_fused.data(), elements); + failures += verify_exact(label.c_str(), fused, reference); + failures += d_embedding.verify_guards((label + " embedding").c_str()); + failures += d_hidden.verify_guards((label + " hidden").c_str()); + failures += d_fused.verify_guards((label + " fused").c_str()); + return failures; +} + +// The route predicate is part of the contract in both directions. If a stem width stopped being +// admitted, the exactness case above would print "route declined" and pass while covering nothing; +// if a width outside the mirrored ops::rmsnorm ladder started being admitted, the kernel would run +// a reduction ops::rmsnorm does not, and the output would stop being bit-exact. +int norm_pack_route(std::int32_t hidden, bool expected) { + const std::size_t bytes = static_cast(hidden) * sizeof(std::uint16_t); + GuardedDeviceBuffer d_embedding(bytes), d_hidden(bytes), d_weight_e(bytes), d_weight_h(bytes); + GuardedDeviceBuffer d_out(2 * bytes); + Tensor t_embedding(d_embedding.data(), DType::BF16, {hidden, 1}); + Tensor t_hidden(d_hidden.data(), DType::BF16, {hidden, 1}); + Tensor t_weight_e(d_weight_e.data(), DType::BF16, {hidden}); + Tensor t_weight_h(d_weight_h.data(), DType::BF16, {hidden}); + Tensor t_out(d_out.data(), DType::BF16, {2 * hidden, 1}); + const bool admitted = + ops::mtp_norm_pack_fc_input_supported(t_embedding, t_weight_e, t_hidden, t_weight_h, t_out); + if (admitted != expected) { + std::cout << "mtp_norm_pack_fc_input route D=" << hidden << " FAILED: expected " + << (expected ? "admitted" : "declined") << "\n"; + return 1; + } + return 0; +} + +// Runs the same shape twice with one input element changed and requires the two fused outputs to +// differ. Without this, "identical" above could be reporting on a kernel that never wrote anything +// the comparison reads. The change is a whole unit and not a ulp on purpose: a ulp of an input can +// round back onto the same output BF16, which would make the control itself flaky. +int norm_pack_strength(std::int32_t hidden, std::int32_t tokens) { + const std::int32_t output_rows = 2 * hidden; + const std::size_t elements = static_cast(output_rows) * tokens; + std::vector> results; + for (bool perturb : {false, true}) { + const std::size_t count = static_cast(hidden) * tokens; + std::vector embedding(count), hidden_state(count), weight_e(hidden), + weight_h(hidden); + fill_uniform(embedding, 0x51ed'0001u, -4.0F, 4.0F); + fill_uniform(hidden_state, 0x51ed'0002u, -4.0F, 4.0F); + fill_uniform(weight_e, 0x51ed'0003u, -0.5F, 0.5F); + fill_uniform(weight_h, 0x51ed'0004u, -0.5F, 0.5F); + if (perturb) { embedding[count / 2] += 1.0F; } + auto pack16 = [](const std::vector& v) { + std::vector out(v.size()); + for (std::size_t i = 0; i < v.size(); ++i) { out[i] = f32_to_bf16(v[i]); } + return out; + }; + const auto embedding16 = pack16(embedding); + const auto hidden16 = pack16(hidden_state); + const auto weight_e16 = pack16(weight_e); + const auto weight_h16 = pack16(weight_h); + const std::size_t in_bytes = count * sizeof(std::uint16_t); + const std::size_t w_bytes = static_cast(hidden) * sizeof(std::uint16_t); + GuardedDeviceBuffer d_embedding(in_bytes), d_hidden(in_bytes); + GuardedDeviceBuffer d_weight_e(w_bytes), d_weight_h(w_bytes); + GuardedDeviceBuffer d_fused(elements * sizeof(std::uint16_t)); + d_embedding.copy_from_host(embedding16.data(), in_bytes); + d_hidden.copy_from_host(hidden16.data(), in_bytes); + d_weight_e.copy_from_host(weight_e16.data(), w_bytes); + d_weight_h.copy_from_host(weight_h16.data(), w_bytes); + d_fused.fill(0xcd); + Tensor t_embedding(d_embedding.data(), DType::BF16, {hidden, tokens}); + Tensor t_hidden(d_hidden.data(), DType::BF16, {hidden, tokens}); + Tensor t_weight_e(d_weight_e.data(), DType::BF16, {hidden}); + Tensor t_weight_h(d_weight_h.data(), DType::BF16, {hidden}); + Tensor t_fused(d_fused.data(), DType::BF16, {output_rows, tokens}); + if (!ops::mtp_norm_pack_fc_input_supported(t_embedding, t_weight_e, t_hidden, t_weight_h, + t_fused)) { + return 0; + } + ops::mtp_norm_pack_fc_input(t_embedding, t_weight_e, t_hidden, t_weight_h, t_fused, 1.0e-6F, + nullptr); + cuda_synchronize(); + results.push_back(from_device(d_fused.data(), elements)); + } + if (results[0] == results[1]) { + std::cout << "mtp_norm_pack_fc_input strength control FAILED: an input change left " + "the output unchanged\n"; + return 1; + } + return 0; +} + +// The fused residual-and-norm has to leave both the updated residual and the normalised +// output byte-for-byte where the two Ops left them, so the reference here is those two Ops run on +// the same inputs in the same process. +int residual_norm_case(std::int32_t hidden, std::int32_t tokens) { + const std::size_t count = static_cast(hidden) * tokens; + std::vector delta(count), residual(count), weight(hidden); + fill_uniform(delta, 0x9e11'0001u, -4.0F, 4.0F); + fill_uniform(residual, 0x9e11'0002u, -4.0F, 4.0F); + fill_uniform(weight, 0x9e11'0003u, -0.5F, 0.5F); + auto pack16 = [](const std::vector& v) { + std::vector out(v.size()); + for (std::size_t i = 0; i < v.size(); ++i) { out[i] = f32_to_bf16(v[i]); } + return out; + }; + const auto delta16 = pack16(delta); + const auto residual16 = pack16(residual); + const auto weight16 = pack16(weight); + const std::size_t bytes = count * sizeof(std::uint16_t); + const std::size_t w_bytes = static_cast(hidden) * sizeof(std::uint16_t); + + GuardedDeviceBuffer d_delta(bytes), d_weight(w_bytes); + GuardedDeviceBuffer d_ref_residual(bytes), d_ref_out(bytes); + GuardedDeviceBuffer d_fused_residual(bytes), d_fused_out(bytes); + d_delta.copy_from_host(delta16.data(), bytes); + d_weight.copy_from_host(weight16.data(), w_bytes); + d_ref_residual.copy_from_host(residual16.data(), bytes); + d_fused_residual.copy_from_host(residual16.data(), bytes); + d_ref_out.fill(0xcd); + d_fused_out.fill(0xcd); + + Tensor t_delta(d_delta.data(), DType::BF16, {hidden, tokens}); + Tensor t_weight(d_weight.data(), DType::BF16, {hidden}); + Tensor t_ref_residual(d_ref_residual.data(), DType::BF16, {hidden, tokens}); + Tensor t_ref_out(d_ref_out.data(), DType::BF16, {hidden, tokens}); + Tensor t_fused_residual(d_fused_residual.data(), DType::BF16, {hidden, tokens}); + Tensor t_fused_out(d_fused_out.data(), DType::BF16, {hidden, tokens}); + + constexpr float kTailEps = 1.0e-6F; + ops::residual_add(t_delta, t_ref_residual, nullptr); + ops::rmsnorm(t_ref_residual, t_weight, kTailEps, true, t_ref_out, nullptr); + cuda_synchronize(); + + const std::string label = + "mtp_residual_norm D=" + std::to_string(hidden) + " T=" + std::to_string(tokens); + if (!ops::mtp_residual_norm_supported(t_delta, t_fused_residual, t_weight, t_fused_out)) { + std::cout << label << ": route declined, two-Op path stands\n"; + return 0; + } + ops::mtp_residual_norm(t_delta, t_fused_residual, t_weight, t_fused_out, kTailEps, nullptr); + cuda_synchronize(); + + int failures = 0; + failures += verify_exact((label + " residual").c_str(), + from_device(d_fused_residual.data(), count), + from_device(d_ref_residual.data(), count)); + failures += verify_exact((label + " output").c_str(), + from_device(d_fused_out.data(), count), + from_device(d_ref_out.data(), count)); + failures += d_delta.verify_guards((label + " delta").c_str()); + failures += d_fused_residual.verify_guards((label + " residual").c_str()); + failures += d_fused_out.verify_guards((label + " out").c_str()); + return failures; +} + +// Same contract in both directions as norm_pack_route: the tail width must be admitted, and a +// width outside the mirrored ops::rmsnorm ladder must not be. +int residual_norm_route(std::int32_t hidden, bool expected) { + const std::size_t bytes = static_cast(hidden) * sizeof(std::uint16_t); + GuardedDeviceBuffer d_delta(bytes), d_residual(bytes), d_weight(bytes), d_out(bytes); + Tensor t_delta(d_delta.data(), DType::BF16, {hidden, 1}); + Tensor t_residual(d_residual.data(), DType::BF16, {hidden, 1}); + Tensor t_weight(d_weight.data(), DType::BF16, {hidden}); + Tensor t_out(d_out.data(), DType::BF16, {hidden, 1}); + const bool admitted = ops::mtp_residual_norm_supported(t_delta, t_residual, t_weight, t_out); + if (admitted != expected) { + std::cout << "mtp_residual_norm route D=" << hidden << " FAILED: expected " + << (expected ? "admitted" : "declined") << "\n"; + return 1; + } + return 0; +} + +// One changed input element has to reach the output, otherwise "identical" above would be +// reporting on a kernel that never wrote what the comparison reads. A whole unit and not a ulp: a +// ulp of an input can round back onto the same output BF16. +int residual_norm_strength(std::int32_t hidden, std::int32_t tokens) { + const std::size_t count = static_cast(hidden) * tokens; + std::vector> results; + for (bool perturb : {false, true}) { + std::vector delta(count), residual(count), weight(hidden); + fill_uniform(delta, 0x9e11'0001u, -4.0F, 4.0F); + fill_uniform(residual, 0x9e11'0002u, -4.0F, 4.0F); + fill_uniform(weight, 0x9e11'0003u, -0.5F, 0.5F); + if (perturb) { delta[count / 2] += 1.0F; } + auto pack16 = [](const std::vector& v) { + std::vector out(v.size()); + for (std::size_t i = 0; i < v.size(); ++i) { out[i] = f32_to_bf16(v[i]); } + return out; + }; + const auto delta16 = pack16(delta); + const auto residual16 = pack16(residual); + const auto weight16 = pack16(weight); + const std::size_t bytes = count * sizeof(std::uint16_t); + const std::size_t w_bytes = static_cast(hidden) * sizeof(std::uint16_t); + GuardedDeviceBuffer d_delta(bytes), d_weight(w_bytes), d_residual(bytes), d_out(bytes); + d_delta.copy_from_host(delta16.data(), bytes); + d_weight.copy_from_host(weight16.data(), w_bytes); + d_residual.copy_from_host(residual16.data(), bytes); + d_out.fill(0xcd); + Tensor t_delta(d_delta.data(), DType::BF16, {hidden, tokens}); + Tensor t_weight(d_weight.data(), DType::BF16, {hidden}); + Tensor t_residual(d_residual.data(), DType::BF16, {hidden, tokens}); + Tensor t_out(d_out.data(), DType::BF16, {hidden, tokens}); + if (!ops::mtp_residual_norm_supported(t_delta, t_residual, t_weight, t_out)) { return 0; } + ops::mtp_residual_norm(t_delta, t_residual, t_weight, t_out, 1.0e-6F, nullptr); + cuda_synchronize(); + results.push_back(from_device(d_out.data(), count)); + } + if (results[0] == results[1]) { + std::cout << "mtp_residual_norm strength control FAILED: an input change left the " + "output unchanged\n"; + return 1; + } + return 0; +} + } // namespace int main() { @@ -152,6 +436,29 @@ int main() { failures += pack_case(2048, 1); failures += pack_case(2048, 6); failures += pack_case(2048, 48); + failures += norm_pack_route(2048, true); + failures += norm_pack_route(5120, true); + failures += norm_pack_route(1536, true); + failures += norm_pack_route(384, false); + failures += norm_pack_case(2048, 1); + failures += norm_pack_case(2048, 4); + failures += norm_pack_case(2048, 48); + failures += norm_pack_case(5120, 1); + failures += norm_pack_case(5120, 6); + failures += norm_pack_case(5120, 64); + failures += residual_norm_route(2048, true); + failures += residual_norm_route(5120, true); + failures += residual_norm_route(384, false); + failures += residual_norm_case(2048, 1); + failures += residual_norm_case(2048, 4); + failures += residual_norm_case(2048, 48); + failures += residual_norm_case(5120, 1); + failures += residual_norm_case(5120, 6); + failures += residual_norm_case(5120, 64); + failures += residual_norm_strength(2048, 4); + failures += residual_norm_strength(5120, 1); + failures += norm_pack_strength(2048, 4); + failures += norm_pack_strength(5120, 1); failures += split_case(1); failures += split_case(6); failures += split_case(48); From 2a9dc0faa72971578b052d5f1720cbe586d532ae Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Thu, 10 Sep 2026 00:53:51 +0200 Subject: [PATCH 2/5] test(mtp): judge the fused stem and tail against an FP64 oracle Every check on the fused Ops used the Ops they replace as the reference, so a defect the two share passed all of them. Add the naive FP64 oracle and the criterion tests/ops/test_rmsnorm.cpp judges ops::rmsnorm by, evaluated from the represented BF16 inputs; the residual half gets an exact oracle, because a BF16 sum of two BF16 values is exact in double. The parity checks stay: they are the bit-exactness these Ops claim. Co-Authored-By: Claude Opus 5 --- tests/ops/test_mtp_pack.cpp | 61 +++++++++++++++++++++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/tests/ops/test_mtp_pack.cpp b/tests/ops/test_mtp_pack.cpp index aa76f1dbe6..9b3db1d6ad 100644 --- a/tests/ops/test_mtp_pack.cpp +++ b/tests/ops/test_mtp_pack.cpp @@ -3,6 +3,7 @@ #include "ninfer/ops/rmsnorm.h" #include "ops/op_tester.h" +#include #include #include #include @@ -14,6 +15,36 @@ using namespace ninfer::test; namespace { +// The parity checks below compare the fused Op against the Ops it replaces, which cannot see a +// defect the two share. These are the independent half: the same naive FP64 oracle and the same +// criterion tests/ops/test_rmsnorm.cpp judges ops::rmsnorm by, evaluated here from the represented +// BF16 inputs and compared against the fused output directly. +constexpr ReductionCriterion rmsnorm_bf16_criterion() { + return {/*relative_l2*/ 1.85e-3, /*gross_absolute*/ 1.0e-5, + /*gross_relative_to_max_reference*/ 3.95e-3}; +} + +std::vector rmsnorm_oracle(const std::vector& input, + const std::vector& weight, std::int32_t hidden, + std::int32_t tokens, float eps) { + std::vector output(input.size()); + for (std::int32_t token = 0; token < tokens; ++token) { + const std::size_t base = static_cast(token) * hidden; + double sum_squares = 0.0; + for (std::int32_t row = 0; row < hidden; ++row) { + const double value = bf16_to_f32(input[base + row]); + sum_squares += value * value; + } + const double inverse = + 1.0 / std::sqrt(sum_squares / static_cast(hidden) + static_cast(eps)); + for (std::int32_t row = 0; row < hidden; ++row) { + output[base + row] = + bf16_to_f32(input[base + row]) * inverse * (1.0 + bf16_to_f32(weight[row])); + } + } + return output; +} + std::vector bit_pattern(std::size_t count, std::uint32_t seed) { std::vector values(count); std::uint32_t state = seed; @@ -208,6 +239,22 @@ int norm_pack_case(std::int32_t hidden, std::int32_t tokens) { const auto reference = from_device(d_reference.data(), elements); const auto fused = from_device(d_fused.data(), elements); failures += verify_exact(label.c_str(), fused, reference); + + const auto norm_e_oracle = rmsnorm_oracle(embedding16, weight_e16, hidden, tokens, kStemEps); + const auto norm_h_oracle = rmsnorm_oracle(hidden16, weight_h16, hidden, tokens, kStemEps); + std::vector oracle(elements), produced(elements); + for (std::int32_t token = 0; token < tokens; ++token) + for (std::int32_t row = 0; row < hidden; ++row) { + const std::size_t in = static_cast(token) * hidden + row; + const std::size_t out = static_cast(token) * output_rows + row; + oracle[out] = norm_e_oracle[in]; + oracle[out + hidden] = norm_h_oracle[in]; + produced[out] = bf16_to_f32(fused[out]); + produced[out + hidden] = bf16_to_f32(fused[out + hidden]); + } + failures += verify_reduction((label + " against the oracle").c_str(), produced, oracle, + rmsnorm_bf16_criterion()); + failures += d_embedding.verify_guards((label + " embedding").c_str()); failures += d_hidden.verify_guards((label + " hidden").c_str()); failures += d_fused.verify_guards((label + " fused").c_str()); @@ -347,6 +394,20 @@ int residual_norm_case(std::int32_t hidden, std::int32_t tokens) { cuda_synchronize(); int failures = 0; + // The updated residual is a BF16 sum of two BF16 values, which is exact in double, so this + // half of the fused Op has an independent oracle that is itself exact. + std::vector residual_oracle(count); + for (std::size_t i = 0; i < count; ++i) + residual_oracle[i] = f32_to_bf16(bf16_to_f32(delta16[i]) + bf16_to_f32(residual16[i])); + failures += + verify_exact((label + " residual against the oracle").c_str(), + from_device(d_fused_residual.data(), count), residual_oracle); + { + const auto out_oracle = rmsnorm_oracle(residual_oracle, weight16, hidden, tokens, kTailEps); + const auto produced = from_device_bf16(d_fused_out.data(), count); + failures += verify_reduction((label + " output against the oracle").c_str(), produced, + out_oracle, rmsnorm_bf16_criterion()); + } failures += verify_exact((label + " residual").c_str(), from_device(d_fused_residual.data(), count), from_device(d_ref_residual.data(), count)); From 03e67bc369ca7ba68ca5e4958cdc7892c29ac242 Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Thu, 10 Sep 2026 01:15:09 +0200 Subject: [PATCH 3/5] fix(mtp): admit only the registered stem widths on the fused route The route mirrored the whole generic RMSNorm ladder, so 512, 1536, 4096 and 8192 were admitted although no target asks for them -- a path kept alive for a model that does not exist. Admit 2048 and 5120, drop the instantiation nothing reaches any more, and pin both directions in the test. Co-Authored-By: Claude Opus 5 --- src/ops/launcher/mtp_pack.cu | 20 +++++++------------- tests/ops/test_mtp_pack.cpp | 4 +++- 2 files changed, 10 insertions(+), 14 deletions(-) diff --git a/src/ops/launcher/mtp_pack.cu b/src/ops/launcher/mtp_pack.cu index af7ba04334..2219c381b5 100644 --- a/src/ops/launcher/mtp_pack.cu +++ b/src/ops/launcher/mtp_pack.cu @@ -36,12 +36,16 @@ namespace { // 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, Cta512x8 }; +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 >= 512 && d <= 3072 && d % 512 == 0) { return MtpRowRoute::Cta256x6; } - if (d > 3072 && d <= 8192 && d % 1024 == 0) { return MtpRowRoute::Cta512x8; } + if (d == 2048) { return MtpRowRoute::Cta256x6; } return MtpRowRoute::None; } @@ -94,11 +98,6 @@ void mtp_norm_pack_fc_input_launch(const Tensor& embedding, const Tensor& embedd mtp_norm_pack_fc_input_kernel <<>>(e, ew, h, hw, o, d, rows, eps); break; - case MtpRowRoute::Cta512x8: - require_row_route(d, 512, 8); - mtp_norm_pack_fc_input_kernel - <<>>(e, ew, h, hw, o, d, rows, eps); - break; case MtpRowRoute::None: throw std::invalid_argument("mtp_norm_pack_fc_input: unsupported width"); } @@ -133,11 +132,6 @@ void mtp_residual_norm_launch(const Tensor& delta, Tensor& residual, const Tenso mtp_residual_norm_kernel <<>>(dv, rv, wv, ov, d, rows, eps); break; - case MtpRowRoute::Cta512x8: - require_row_route(d, 512, 8); - mtp_residual_norm_kernel - <<>>(dv, rv, wv, ov, d, rows, eps); - break; case MtpRowRoute::None: throw std::invalid_argument("mtp_residual_norm: unsupported width"); } diff --git a/tests/ops/test_mtp_pack.cpp b/tests/ops/test_mtp_pack.cpp index 9b3db1d6ad..fe8099f22b 100644 --- a/tests/ops/test_mtp_pack.cpp +++ b/tests/ops/test_mtp_pack.cpp @@ -499,7 +499,8 @@ int main() { failures += pack_case(2048, 48); failures += norm_pack_route(2048, true); failures += norm_pack_route(5120, true); - failures += norm_pack_route(1536, true); + failures += norm_pack_route(1536, false); + failures += norm_pack_route(4096, false); failures += norm_pack_route(384, false); failures += norm_pack_case(2048, 1); failures += norm_pack_case(2048, 4); @@ -509,6 +510,7 @@ int main() { failures += norm_pack_case(5120, 64); failures += residual_norm_route(2048, true); failures += residual_norm_route(5120, true); + failures += residual_norm_route(1536, false); failures += residual_norm_route(384, false); failures += residual_norm_case(2048, 1); failures += residual_norm_case(2048, 4); From 59f80df250b7cef52e37ca390e8d973c3c46b40e Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:02:02 +0200 Subject: [PATCH 4/5] fix(mtp): refuse an eps that is not positive and finite Both fused entry points stand in for a composition containing ops::rmsnorm, which refuses such an eps; they launched instead and came back with NaN. Apply the same check and pin it in the test, on both Ops and on both registered widths. Co-Authored-By: Claude Opus 5 --- src/ops/wrapper/mtp_pack.cpp | 12 +++++++++++ tests/ops/test_mtp_pack.cpp | 39 ++++++++++++++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/src/ops/wrapper/mtp_pack.cpp b/src/ops/wrapper/mtp_pack.cpp index ff796fbc4f..33244bde87 100644 --- a/src/ops/wrapper/mtp_pack.cpp +++ b/src/ops/wrapper/mtp_pack.cpp @@ -1,6 +1,7 @@ #include "ninfer/ops/mtp_pack.h" #include "ops/launcher/mtp_pack.h" +#include #include #include @@ -69,6 +70,15 @@ void require_stem_operands(const Tensor& embedding, const Tensor& embedding_weig } // namespace +// The fused Ops stand in for a composition that contains ops::rmsnorm, so they owe the same +// contract on eps that ops::rmsnorm enforces (src/ops/wrapper/rmsnorm.cpp). Without this a +// non-positive or non-finite eps would reach the kernel and come back as NaN instead of a throw. +void require_normalization_eps(float eps, const char* op) { + if (!(eps > 0.0F) || !std::isfinite(eps)) { + throw std::invalid_argument(std::string(op) + ": eps must be positive and finite"); + } +} + bool mtp_norm_pack_fc_input_supported(const Tensor& embedding, const Tensor& embedding_weight, const Tensor& hidden, const Tensor& hidden_weight, const Tensor& out) { @@ -81,6 +91,7 @@ void mtp_norm_pack_fc_input(const Tensor& embedding, const Tensor& embedding_wei const Tensor& hidden, const Tensor& hidden_weight, Tensor& out, float eps, cudaStream_t stream) { require_stem_operands(embedding, embedding_weight, hidden, hidden_weight, out); + require_normalization_eps(eps, "mtp_norm_pack_fc_input"); if (!detail::mtp_norm_pack_fc_input_admits(embedding.ne[0], embedding, embedding_weight, hidden, hidden_weight, out)) { throw std::invalid_argument("mtp_norm_pack_fc_input: unsupported operand profile"); @@ -120,6 +131,7 @@ bool mtp_residual_norm_supported(const Tensor& delta, const Tensor& residual, co void mtp_residual_norm(const Tensor& delta, Tensor& residual, const Tensor& weight, Tensor& out, float eps, cudaStream_t stream) { + require_normalization_eps(eps, "mtp_residual_norm"); require_residual_norm_operands(delta, residual, weight, out); if (!detail::mtp_residual_norm_admits(residual.ne[0], delta, residual, weight, out)) { throw std::invalid_argument("mtp_residual_norm: unsupported operand profile"); diff --git a/tests/ops/test_mtp_pack.cpp b/tests/ops/test_mtp_pack.cpp index fe8099f22b..5e2c21fe5b 100644 --- a/tests/ops/test_mtp_pack.cpp +++ b/tests/ops/test_mtp_pack.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -261,6 +262,42 @@ int norm_pack_case(std::int32_t hidden, std::int32_t tokens) { return failures; } +// An eps that is not positive and finite has to be refused, not normalised with: the composition +// these Ops replace refuses it in ops::rmsnorm, and a kernel handed it returns NaN instead. +int eps_contract(std::int32_t hidden) { + const std::size_t bytes = static_cast(hidden) * sizeof(std::uint16_t); + GuardedDeviceBuffer d_embedding(bytes), d_hidden(bytes), d_weight_e(bytes), d_weight_h(bytes); + GuardedDeviceBuffer d_residual(bytes), d_out(bytes), d_pack(2 * bytes); + Tensor t_embedding(d_embedding.data(), DType::BF16, {hidden, 1}); + Tensor t_hidden(d_hidden.data(), DType::BF16, {hidden, 1}); + Tensor t_weight_e(d_weight_e.data(), DType::BF16, {hidden}); + Tensor t_weight_h(d_weight_h.data(), DType::BF16, {hidden}); + Tensor t_residual(d_residual.data(), DType::BF16, {hidden, 1}); + Tensor t_out(d_out.data(), DType::BF16, {hidden, 1}); + Tensor t_pack(d_pack.data(), DType::BF16, {2 * hidden, 1}); + + int failures = 0; + const auto refuses = [&](const char* what, auto&& call) { + try { + call(); + } catch (const std::invalid_argument&) { return; } + std::cout << "mtp eps contract D=" << hidden << " " << what + << " FAILED: accepted an invalid eps\n"; + ++failures; + }; + for (const float eps : {0.0F, -1.0e-6F, std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity()}) { + refuses("norm_pack", [&] { + ops::mtp_norm_pack_fc_input(t_embedding, t_weight_e, t_hidden, t_weight_h, t_pack, eps, + nullptr); + }); + refuses("residual_norm", [&] { + ops::mtp_residual_norm(t_embedding, t_residual, t_weight_e, t_out, eps, nullptr); + }); + } + return failures; +} + // The route predicate is part of the contract in both directions. If a stem width stopped being // admitted, the exactness case above would print "route declined" and pass while covering nothing; // if a width outside the mirrored ops::rmsnorm ladder started being admitted, the kernel would run @@ -522,6 +559,8 @@ int main() { failures += residual_norm_strength(5120, 1); failures += norm_pack_strength(2048, 4); failures += norm_pack_strength(5120, 1); + failures += eps_contract(2048); + failures += eps_contract(5120); failures += split_case(1); failures += split_case(6); failures += split_case(48); From d121079154ce4dd3b434695451d7676c99c7ca4f Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:34:41 +0200 Subject: [PATCH 5/5] fix(mtp): state the fused contract semantically, and hold the weight shape The Numeric fields made the current RMSNorm block width, pair decomposition and byte parity contractual, which docs/maintainer/op-development.md forbids: a valid retuning of ops::rmsnorm would have broken the contract without changing what the Op computes. State the formula, the eps domain, the observable BF16 roundings and the FP64-oracle criterion instead; parity stays where it belongs, in the tests. The weight check counted elements, so [1,D] and [D/2,2] reached the fused route although ops::rmsnorm refuses them. Check the whole shape in both Ops. Co-Authored-By: Claude Opus 5 --- include/ninfer/ops/mtp_pack.h | 22 +++++++++++------- src/ops/wrapper/mtp_pack.cpp | 15 +++++++----- tests/ops/test_mtp_pack.cpp | 43 +++++++++++++++++++++++++++++++++++ 3 files changed, 66 insertions(+), 14 deletions(-) diff --git a/include/ninfer/ops/mtp_pack.h b/include/ninfer/ops/mtp_pack.h index 25a77863f0..915a6aa5f4 100644 --- a/include/ninfer/ops/mtp_pack.h +++ b/include/ninfer/ops/mtp_pack.h @@ -41,10 +41,13 @@ void mtp_pack_fc_input(const Tensor& embedding_norm, const Tensor& hidden_norm, * BF16 embedding and hidden [D,T], weights [D], out [2D,T], all contiguous. * * Numeric: - * Identical to rmsnorm() applied to each half followed by mtp_pack_fc_input(): the same block - * width, pair decomposition, reduction and epilogue produce the same bytes. Returns false from - * mtp_norm_pack_fc_input_supported() for shapes it does not cover, and the caller keeps the - * three-Op path there. + * 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. @@ -72,10 +75,13 @@ void mtp_norm_pack_fc_input(const Tensor& embedding, const Tensor& embedding_wei * BF16 delta and residual [D,T], weight [D], out [D,T], all contiguous. * * Numeric: - * Identical to residual_add() followed by rmsnorm() on the updated residual: the same pairwise - * FP32 addition and rounding, then the same block width, pair decomposition, reduction and - * epilogue. mtp_residual_norm_supported() reports the shapes this route covers; the caller keeps - * the two-Op path elsewhere. + * 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 diff --git a/src/ops/wrapper/mtp_pack.cpp b/src/ops/wrapper/mtp_pack.cpp index 33244bde87..fafa470129 100644 --- a/src/ops/wrapper/mtp_pack.cpp +++ b/src/ops/wrapper/mtp_pack.cpp @@ -27,6 +27,12 @@ void require_shape(const Tensor& t, std::int32_t n0, std::int32_t n1, const char } } +// The contract says [D], and ops::rmsnorm holds callers to it. A length check alone would let +// [1,D] or [D/2,2] through here and take the fused route on an operand the composed path refuses. +void require_weight_vector(const Tensor& t, std::int32_t d, const char* op, const char* name) { + require_shape(t, d, 1, op, name); +} + } // namespace void mtp_pack_fc_input(const Tensor& embedding_norm, const Tensor& hidden_norm, Tensor& out, @@ -63,9 +69,8 @@ void require_stem_operands(const Tensor& embedding, const Tensor& embedding_weig require_shape(embedding, rows, tokens, op, "embedding"); require_shape(hidden, rows, tokens, op, "hidden"); require_shape(out, 2 * rows, tokens, op, "out"); - if (embedding_weight.numel() != rows || hidden_weight.numel() != rows) { - throw std::invalid_argument("mtp_norm_pack_fc_input: weight length must be D"); - } + require_weight_vector(embedding_weight, rows, op, "embedding_weight"); + require_weight_vector(hidden_weight, rows, op, "hidden_weight"); } } // namespace @@ -116,9 +121,7 @@ void require_residual_norm_operands(const Tensor& delta, const Tensor& residual, require_shape(delta, rows, tokens, op, "delta"); require_shape(residual, rows, tokens, op, "residual"); require_shape(out, rows, tokens, op, "out"); - if (weight.numel() != rows) { - throw std::invalid_argument("mtp_residual_norm: weight length must be D"); - } + require_weight_vector(weight, rows, op, "weight"); } } // namespace diff --git a/tests/ops/test_mtp_pack.cpp b/tests/ops/test_mtp_pack.cpp index 5e2c21fe5b..d3abb3bcdb 100644 --- a/tests/ops/test_mtp_pack.cpp +++ b/tests/ops/test_mtp_pack.cpp @@ -298,6 +298,47 @@ int eps_contract(std::int32_t hidden) { return failures; } +// The weight contract is [D]. A view with the right element count but the wrong shape has to be +// refused here exactly as ops::rmsnorm refuses it, or the fused route would run on an operand the +// composed path would not accept. +int weight_shape_contract(std::int32_t hidden) { + const std::size_t bytes = static_cast(hidden) * sizeof(std::uint16_t); + GuardedDeviceBuffer d_embedding(bytes), d_hidden(bytes), d_weight(bytes), d_residual(bytes); + GuardedDeviceBuffer d_out(bytes), d_pack(2 * bytes); + Tensor t_embedding(d_embedding.data(), DType::BF16, {hidden, 1}); + Tensor t_hidden(d_hidden.data(), DType::BF16, {hidden, 1}); + Tensor t_good(d_weight.data(), DType::BF16, {hidden}); + Tensor t_residual(d_residual.data(), DType::BF16, {hidden, 1}); + Tensor t_out(d_out.data(), DType::BF16, {hidden, 1}); + Tensor t_pack(d_pack.data(), DType::BF16, {2 * hidden, 1}); + + int failures = 0; + const auto refuses = [&](const char* what, auto&& call) { + try { + call(); + } catch (const std::invalid_argument&) { return; } + std::cout << "mtp weight shape D=" << hidden << " " << what + << " FAILED: accepted a weight that is not [D]\n"; + ++failures; + }; + const Tensor bad_rowed(d_weight.data(), DType::BF16, {1, hidden}); + const Tensor bad_split(d_weight.data(), DType::BF16, {hidden / 2, 2}); + for (const Tensor* bad : {&bad_rowed, &bad_split}) { + refuses("norm_pack embedding_weight", [&] { + ops::mtp_norm_pack_fc_input(t_embedding, *bad, t_hidden, t_good, t_pack, 1.0e-6F, + nullptr); + }); + refuses("norm_pack hidden_weight", [&] { + ops::mtp_norm_pack_fc_input(t_embedding, t_good, t_hidden, *bad, t_pack, 1.0e-6F, + nullptr); + }); + refuses("residual_norm weight", [&] { + ops::mtp_residual_norm(t_embedding, t_residual, *bad, t_out, 1.0e-6F, nullptr); + }); + } + return failures; +} + // The route predicate is part of the contract in both directions. If a stem width stopped being // admitted, the exactness case above would print "route declined" and pass while covering nothing; // if a width outside the mirrored ops::rmsnorm ladder started being admitted, the kernel would run @@ -559,6 +600,8 @@ int main() { failures += residual_norm_strength(5120, 1); failures += norm_pack_strength(2048, 4); failures += norm_pack_strength(5120, 1); + failures += weight_shape_contract(2048); + failures += weight_shape_contract(5120); failures += eps_contract(2048); failures += eps_contract(5120); failures += split_case(1);