diff --git a/include/ninfer/ops/mtp_pack.h b/include/ninfer/ops/mtp_pack.h index 083ee90e21..915a6aa5f4 100644 --- a/include/ninfer/ops/mtp_pack.h +++ b/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 * 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..2219c381b5 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,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(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::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::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..fafa470129 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 @@ -26,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, @@ -45,6 +52,96 @@ 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"); + require_weight_vector(embedding_weight, rows, op, "embedding_weight"); + require_weight_vector(hidden_weight, rows, op, "hidden_weight"); +} + +} // 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) { + 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); + 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"); + } + 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"); + require_weight_vector(weight, rows, op, "weight"); +} + +} // 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_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"); + } + 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..d3abb3bcdb 100644 --- a/tests/ops/test_mtp_pack.cpp +++ b/tests/ops/test_mtp_pack.cpp @@ -1,7 +1,11 @@ #include "ninfer/ops/mtp_pack.h" +#include "ninfer/ops/residual_add.h" +#include "ninfer/ops/rmsnorm.h" #include "ops/op_tester.h" +#include #include +#include #include #include #include @@ -12,6 +16,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; @@ -137,6 +171,395 @@ 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); + + 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()); + 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 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 +// 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; + // 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)); + 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 +575,35 @@ 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, 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); + 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(1536, false); + 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 += weight_shape_contract(2048); + failures += weight_shape_contract(5120); + failures += eps_contract(2048); + failures += eps_contract(5120); failures += split_case(1); failures += split_case(6); failures += split_case(48);