diff --git a/bench/ops/rmsnorm_rope_bench.cu b/bench/ops/rmsnorm_rope_bench.cu index fed660668d..f1a1a38975 100644 --- a/bench/ops/rmsnorm_rope_bench.cu +++ b/bench/ops/rmsnorm_rope_bench.cu @@ -1,6 +1,11 @@ -// Public-Op benchmark for variable-width DFlash2 pair and context-K RMSNorm+RoPE profiles. +// Public-Op benchmark for the variable-width DFlash2 pair, context-K and text RMSNorm+RoPE +// profiles. The text profile carries both routes: the three calls the model issues today, and the +// fused Op that replaces them. #include "ninfer/ops/rmsnorm_rope.h" +#include "ninfer/ops/rmsnorm.h" +#include "ninfer/ops/rope.h" + #include "ninfer_bench_common.h" #include @@ -22,7 +27,14 @@ constexpr int kHeadDim = 128; constexpr int kQueryHeads = 32; constexpr int kKeyHeads = 8; -enum class Form : std::uint8_t { Pair, Single }; +// The text profile, as the Op's contract states it. +constexpr int kTextHeadDim = 256; +constexpr int kTextRotaryDim = 64; +constexpr float kTextRopeBase = 1.0e7F; +constexpr float kTextEps = 1.0e-6F; + +enum class Form : std::uint8_t { Pair, Single, Text }; +enum class Route : std::uint8_t { Split, Fused }; enum class Execution : std::uint8_t { Eager, Graph }; struct Options { @@ -34,12 +46,14 @@ struct Options { int warmup = 20; int repeat = 200; bool profile = false; + Route route = Route::Fused; }; [[noreturn]] void usage(const char* message) { std::fprintf(stderr, "error: %s\n" - "usage: ninfer_rmsnorm_rope_bench --form pair|single " + "usage: ninfer_rmsnorm_rope_bench --form pair|single|text " + "[--route split|fused] " "[--widths W,...] [--batches B,...] [--tokens T,...] [--execution eager|graph] " "[--warmup N] [--repeat N] [--profile]\n", message); @@ -85,8 +99,18 @@ Options parse_options(int argc, char** argv) { options.form = Form::Pair; else if (value == "single") options.form = Form::Single; + else if (value == "text") + options.form = Form::Text; else - usage("--form expects pair or single"); + usage("--form expects pair, single or text"); + } else if (argument == "--route") { + const std::string_view value(next("--route requires a value")); + if (value == "split") + options.route = Route::Split; + else if (value == "fused") + options.route = Route::Fused; + else + usage("--route expects split or fused"); } else if (argument == "--widths") { options.widths = parse_list(next("--widths requires a value"), 2, 16, "--widths"); } else if (argument == "--batches") { @@ -213,6 +237,58 @@ void run_single(const Options& options, int tokens, cudaStream_t stream) { options.execution == Execution::Graph ? 32 : 1); } +// The text profile of the two registered geometries. `split` is what every full-attention layer +// issues today: normalize q, normalize k, rotate both. `fused` is the Op that replaces the three. +void run_text(const Options& options, int query_heads, int key_heads, int tokens, + cudaStream_t stream) { + const auto positions_host = host_positions(tokens); + DeviceBuffer positions(positions_host.size() * sizeof(std::int32_t)); + positions.copy_from_host(positions_host.data(), positions.bytes); + const std::size_t q_elements = static_cast(kTextHeadDim) * query_heads * tokens; + const std::size_t k_elements = static_cast(kTextHeadDim) * key_heads * tokens; + DeviceBuffer q = bench::make_bf16(q_elements); + DeviceBuffer k = bench::make_bf16(k_elements); + DeviceBuffer qn = bench::make_bf16(q_elements); + DeviceBuffer kn = bench::make_bf16(k_elements); + DeviceBuffer q_weight = bench::make_bf16(kTextHeadDim); + DeviceBuffer k_weight = bench::make_bf16(kTextHeadDim); + Tensor t_positions(positions.p, DType::I32, {tokens}); + Tensor t_q(q.p, DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor t_k(k.p, DType::BF16, {kTextHeadDim, key_heads, tokens}); + Tensor t_qn(qn.p, DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor t_kn(kn.p, DType::BF16, {kTextHeadDim, key_heads, tokens}); + Tensor t_q_weight(q_weight.p, DType::BF16, {kTextHeadDim}); + Tensor t_k_weight(k_weight.p, DType::BF16, {kTextHeadDim}); + const auto launch = [&](cudaStream_t launch_stream) { + if (options.route == Route::Fused) { + ops::rmsnorm_rope(t_positions, t_q_weight, t_k_weight, t_q, t_k, t_qn, t_kn, + launch_stream); + return; + } + ops::rmsnorm(t_q, t_q_weight, kTextEps, true, t_qn, launch_stream); + ops::rmsnorm(t_k, t_k_weight, kTextEps, true, t_kn, launch_stream); + ops::rope(t_positions, kTextRotaryDim, kTextRopeBase, t_qn, t_kn, launch_stream); + }; + if (options.profile) { + for (int index = 0; index < options.warmup; ++index) launch(stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaProfilerStart()); + launch(stream); + CUDA_CHECK(cudaStreamSynchronize(stream)); + CUDA_CHECK(cudaProfilerStop()); + return; + } + // Read in, write out, for both operands. + const double bytes = 2.0 * 2.0 * static_cast(q_elements + k_elements); + const bench::Result timing = measure(options, launch, bytes, stream); + std::printf("form=text route=%s Q=%d K=%d T=%d execution=%s median=%.3f us min=%.3f us " + "p95=%.3f us useful=%.1f GB/s graph_repetitions=%d cache=warm\n", + options.route == Route::Fused ? "fused" : "split", query_heads, key_heads, tokens, + options.execution == Execution::Graph ? "graph" : "eager", timing.median_us, + timing.min_us, timing.p95_us, timing.gbs, + options.execution == Execution::Graph ? 32 : 1); +} + } // namespace int main(int argc, char** argv) { @@ -228,8 +304,13 @@ int main(int argc, char** argv) { if (options.form == Form::Pair) { for (int width : options.widths) for (int batch : options.batches) run_pair(options, width, batch, stream); - } else { + } else if (options.form == Form::Single) { for (int tokens : options.tokens) run_single(options, tokens, stream); + } else { + for (int tokens : options.tokens) { + run_text(options, 16, 2, tokens, stream); + run_text(options, 24, 4, tokens, stream); + } } CUDA_CHECK(cudaStreamDestroy(stream)); return 0; diff --git a/include/ninfer/ops/rmsnorm_rope.h b/include/ninfer/ops/rmsnorm_rope.h index 54c2ab121e..81541d4f2e 100644 --- a/include/ninfer/ops/rmsnorm_rope.h +++ b/include/ninfer/ops/rmsnorm_rope.h @@ -39,4 +39,32 @@ void rmsnorm_rope(const Tensor& positions, const Tensor& q_norm_weight, const Te void rmsnorm_rope(const Tensor& positions, const Tensor& norm_weight, Tensor& x, cudaStream_t stream); +/** + * Text form of the same fusion: wider heads, a narrower rotation, and out of place. + * + * The profile is q_in BF16 [256,Q,T], k_in BF16 [256,K,T], q_out and k_out of the same shapes as + * their inputs, q_norm_weight and k_norm_weight BF16 [256], and positions I32 [T], with + * (Q,K) either (16,2) or (24,4) and T any positive count the launch grid can address. For + * each head and token, + * + * inv = 1 / sqrt(sum_d x[d]^2 / 256 + 1e-6) + * n[d] = x[d] * inv * (norm_weight[d] + 1) + * angle(i) = position * (1e7)^(-2*i/64), 0<=i<32 + * out[i] = n[i] * cos(angle(i)) - n[i+32] * sin(angle(i)) + * out[i+32] = n[i+32] * cos(angle(i)) + n[i] * sin(angle(i)) + * out[d] = n[d] for d >= 64. + * + * Only the first 64 channels rotate; the remaining 192 carry the normalized value through. The + * weight enters as a delta around one - the Offset epilogue the text stack normalizes with - + * unlike the two in-place forms above, which multiply by the stored weight directly. Unlike + * the in-place forms above, n IS observable at BF16 for d >= 64, and the rotation consumes the + * BF16 represented n, so the result is bit-identical to rmsnorm(q_in) -> rmsnorm(k_in) -> + * rope(q_out, k_out) with the Offset epilogue. The outputs must not overlap each other, the + * inputs, positions, or either norm weight; read-only operands may overlap each other. All + * tensors are contiguous and 4-byte aligned. The Op owns no workspace or persistent state. + */ +void rmsnorm_rope(const Tensor& positions, const Tensor& q_norm_weight, const Tensor& k_norm_weight, + const Tensor& q_in, const Tensor& k_in, Tensor& q_out, Tensor& k_out, + cudaStream_t stream); + } // namespace ninfer::ops diff --git a/src/models/qwen3_5/execution/attention.cpp b/src/models/qwen3_5/execution/attention.cpp index 5fa6a38551..eb025bcb1f 100644 --- a/src/models/qwen3_5/execution/attention.cpp +++ b/src/models/qwen3_5/execution/attention.cpp @@ -1,6 +1,8 @@ #include "models/qwen3_5/execution/attention.h" #include "ninfer/ops/attn_input_proj.h" +#include "ninfer/ops/rmsnorm.h" +#include "ninfer/ops/rmsnorm_rope.h" #include "ninfer/ops/rope.h" #include @@ -17,6 +19,25 @@ void require_rope_axes(const Tensor& positions, const RopeConfig& config) { } } +// The fused text form is registered for the two text head geometries with a one-dimensional +// position axis. The MRoPE path and any other geometry take the three calls it replaces. +// +// It is also bounded in width. One warp owns one head, so the fused kernel stops gaining once a +// width alone fills the machine, and past that the three separate kernels - each free to choose +// its own shape - are ahead: measured on an RTX 5090, the fused form wins by 22 to 52 % through +// 256 tokens and loses by up to 22 % at 1024. The bound sits a doubling below the crossover +// because the two geometries cross at different widths. The Op itself is valid at any width; this +// is a dispatch choice, and both branches are the same arithmetic bit for bit. +constexpr std::int32_t kFusedTextQkNormRopeMaximumTokens = 256; + +bool fused_text_qk_norm_rope(const Tensor& positions, const RopeConfig& rope, + const AttentionConfig& attention, std::int32_t tokens) { + return positions.ne[1] == 1 && tokens <= kFusedTextQkNormRopeMaximumTokens && + attention.head_dim == 256 && rope.rotary_dim == 64 && + ((attention.num_attention_heads == 16 && attention.num_key_value_heads == 2) || + (attention.num_attention_heads == 24 && attention.num_key_value_heads == 4)); +} + } // namespace std::size_t attention_projection_workspace_bytes(const AttentionParameters& parameters, @@ -56,4 +77,21 @@ void text_rope(const Tensor& positions, const RopeConfig& config, Tensor& query, ops::rope(positions, dimension(config.rotary_dim), config.rope_theta, query, key, stream); } +void text_qk_norm_rope(const Tensor& positions, const RopeConfig& rope, + const AttentionConfig& attention, float rms_norm_eps, + const Tensor& q_norm_weight, const Tensor& k_norm_weight, + const Tensor& query, const Tensor& key, Tensor& normalized_query, + Tensor& normalized_key, cudaStream_t stream) { + require_rope_axes(positions, rope); + if (fused_text_qk_norm_rope(positions, rope, attention, query.ne[2])) { + ops::rmsnorm_rope(positions, q_norm_weight, k_norm_weight, query, key, normalized_query, + normalized_key, stream); + return; + } + ops::rmsnorm(query, q_norm_weight, rms_norm_eps, true, normalized_query, stream); + ops::rmsnorm(key, k_norm_weight, rms_norm_eps, true, normalized_key, stream); + ops::rope(positions, dimension(rope.rotary_dim), rope.rope_theta, normalized_query, + normalized_key, stream); +} + } // namespace ninfer::models::qwen3_5::execution diff --git a/src/models/qwen3_5/execution/attention.h b/src/models/qwen3_5/execution/attention.h index 3eaa51062b..ad4edb77fb 100644 --- a/src/models/qwen3_5/execution/attention.h +++ b/src/models/qwen3_5/execution/attention.h @@ -16,4 +16,13 @@ void text_rope(const Tensor& positions, const RopeConfig& config, Tensor& query, void text_rope(const Tensor& positions, const RopeConfig& config, Tensor& query, Tensor& key, cudaStream_t stream); +// Normalize q and k and rotate them. Where the fused Op covers the geometry this is one graph node +// instead of three; everywhere else it is the three calls it replaces, which are the same +// arithmetic bit for bit. It chooses a schedule, not a result. +void text_qk_norm_rope(const Tensor& positions, const RopeConfig& rope, + const AttentionConfig& attention, float rms_norm_eps, + const Tensor& q_norm_weight, const Tensor& k_norm_weight, + const Tensor& query, const Tensor& key, Tensor& normalized_query, + Tensor& normalized_key, cudaStream_t stream); + } // namespace ninfer::models::qwen3_5::execution diff --git a/src/models/qwen3_5/execution/text.cpp b/src/models/qwen3_5/execution/text.cpp index d9037515b3..4d43816cb3 100644 --- a/src/models/qwen3_5/execution/text.cpp +++ b/src/models/qwen3_5/execution/text.cpp @@ -351,10 +351,10 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po dimension(config_.attention->num_attention_heads), T}); Tensor kn = results.normalized_key.view({dimension(config_.attention->head_dim), dimension(config_.attention->num_key_value_heads), T}); - ops::rmsnorm(q, mtp_->query_norm, config_.rms_norm_eps, true, qn, s); - ops::rmsnorm(k, mtp_->key_norm, config_.rms_norm_eps, true, kn, s); + Tensor rope_for_op = active_sequence_batch_ != 0 ? rope_positions.view({T}) : rope_positions; - text_rope(rope_for_op, *config_.rope_parameters, qn, kn, s); + text_qk_norm_rope(rope_for_op, *config_.rope_parameters, *config_.attention, + config_.rms_norm_eps, mtp_->query_norm, mtp_->key_norm, q, k, qn, kn, s); Tensor a = results.attention.view({dimension(config_.attention->head_dim), dimension(config_.attention->num_attention_heads), T}); @@ -869,14 +869,13 @@ void TextContext::attn_mix(const BlockParameters& w, Tensor& x, int fidx, Phase dimension(config_.attention->num_attention_heads), T}); Tensor kn = results.normalized_key.view({dimension(config_.attention->head_dim), dimension(config_.attention->num_key_value_heads), T}); - ops::rmsnorm(q, p.query_norm, config_.rms_norm_eps, true, qn, s); - ops::rmsnorm(k, p.key_norm, config_.rms_norm_eps, true, kn, s); const Tensor& cache_positions = active_cache_positions_ != nullptr ? *active_cache_positions_ : io_.pos; const Tensor& rope_positions = active_rope_positions_ != nullptr ? *active_rope_positions_ : io_.rope_pos; Tensor rope_for_op = active_sequence_batch_ != 0 ? rope_positions.view({T}) : rope_positions; - text_rope(rope_for_op, *config_.rope_parameters, qn, kn, s); + text_qk_norm_rope(rope_for_op, *config_.rope_parameters, *config_.attention, + config_.rms_norm_eps, p.query_norm, p.key_norm, q, k, qn, kn, s); Tensor a = results.attention.view({dimension(config_.attention->head_dim), dimension(config_.attention->num_attention_heads), T}); diff --git a/src/ops/rmsnorm_rope/d256.cuh b/src/ops/rmsnorm_rope/d256.cuh new file mode 100644 index 0000000000..1e0044952c --- /dev/null +++ b/src/ops/rmsnorm_rope/d256.cuh @@ -0,0 +1,68 @@ +#pragma once + +#include "ops/common/warp.cuh" +#include "ops/kernel/rmsnorm.cuh" + +#include + +namespace ninfer::ops::detail { + +// One warp owns one represented BF16 D256 head. Lane l carries the pairs l, l+32, l+64, l+96, the +// layout rmsnorm_warp_bf16x2_kernel uses, so the sum of squares accumulates in the same order and +// the epilogue is the same helper: the normalized value is bit-identical to the standalone norm. +struct RmsnormRopeD256Head { + __nv_bfloat162 pair[4]; +}; + +__device__ __forceinline__ RmsnormRopeD256Head rmsnorm_rope_d256_normalize( + const __nv_bfloat162* __restrict__ input, const __nv_bfloat162* __restrict__ weight, + std::int64_t base, int lane) { + constexpr int kHeadDim = 256; + constexpr float kEpsilon = 1.0e-6F; + __nv_bfloat162 values[4]; + __nv_bfloat162 weights[4]; + float sum = 0.0F; +#pragma unroll + for (int k = 0; k < 4; ++k) { + const int pair = lane + k * 32; + values[k] = input[base + pair]; + weights[k] = weight[pair]; + const float2 xf = __bfloat1622float2(values[k]); + sum += xf.x * xf.x + xf.y * xf.y; + } + sum = warp_reduce_sum(sum); + float inv = lane == 0 ? rsqrtf(sum / static_cast(kHeadDim) + kEpsilon) : 0.0F; + inv = __shfl_sync(kFullWarpMask, inv, 0); + + RmsnormRopeD256Head out; +#pragma unroll + for (int k = 0; k < 4; ++k) { + const float2 xf = __bfloat1622float2(values[k]); + const float2 wf = __bfloat1622float2(weights[k]); + out.pair[k] = + __floats2bfloat162_rn(rmsnorm_epilogue(xf.x, inv, wf.x, 0.0F), + rmsnorm_epilogue(xf.y, inv, wf.y, 0.0F)); + } + return out; +} + +// Split-half rotation over the first 64 channels, which is what R=64 means for a 256-wide head: +// channel p pairs with p + 32. The norm layout keeps those two in different lanes, so the partner +// arrives through __shfl_xor_sync(..., 16) and the coefficients are indexed by lane & 15 - exactly +// the ones lane p < 16 receives in the standalone rope kernel. +__device__ __forceinline__ __nv_bfloat162 rmsnorm_rope_d256_rotate(__nv_bfloat162 normalized, + float c0, float c1, float s0, + float s1, int lane) { + constexpr int kHalfPair = 16; + const __nv_bfloat162 theirs = __shfl_xor_sync(kFullWarpMask, normalized, kHalfPair); + const float2 first = + lane < kHalfPair ? __bfloat1622float2(normalized) : __bfloat1622float2(theirs); + const float2 second = + lane < kHalfPair ? __bfloat1622float2(theirs) : __bfloat1622float2(normalized); + if (lane < kHalfPair) { + return __floats2bfloat162_rn(first.x * c0 - second.x * s0, first.y * c1 - second.y * s1); + } + return __floats2bfloat162_rn(second.x * c0 + first.x * s0, second.y * c1 + first.y * s1); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/rmsnorm_rope/kernel.cuh b/src/ops/rmsnorm_rope/kernel.cuh index 2a4dd596bc..271d66ae4d 100644 --- a/src/ops/rmsnorm_rope/kernel.cuh +++ b/src/ops/rmsnorm_rope/kernel.cuh @@ -1,6 +1,8 @@ #pragma once #include "ops/common/dflash_rope.cuh" +#include "ops/kernel/rope.cuh" #include "ops/rmsnorm_rope/d128.cuh" +#include "ops/rmsnorm_rope/d256.cuh" #include #include @@ -36,4 +38,47 @@ __global__ __launch_bounds__(256) void rmsnorm_rope_d128_kernel( data[base + lane] = out.first; data[base + lane + 32] = out.second; } + +// Text form: D=256 heads, rotary width 64, out of place. One warp owns one head; HeadsPerBlock +// warps share a block. The Q and K heads of one token are laid out as one combined range so a +// single grid covers both tensors and no head group is left half empty. +template +__global__ __launch_bounds__(HeadsPerBlock * 32) void rmsnorm_rope_d256_text_kernel( + const std::int32_t* __restrict__ positions, const __nv_bfloat162* __restrict__ q_norm, + const __nv_bfloat162* __restrict__ k_norm, const __nv_bfloat162* __restrict__ q_in, + const __nv_bfloat162* __restrict__ k_in, __nv_bfloat162* __restrict__ q_out, + __nv_bfloat162* __restrict__ k_out, std::int32_t tokens) { + constexpr int kPairs = 128; + constexpr int kHalfPair = 16; + constexpr int kCombined = QHeads + KHeads; + constexpr int kGroups = (kCombined + HeadsPerBlock - 1) / HeadsPerBlock; + + const int token = static_cast(blockIdx.x) / kGroups; + if (token >= tokens) { return; } + const int group = static_cast(blockIdx.x) % kGroups; + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + const int combined = group * HeadsPerBlock + warp; + if (combined >= kCombined) { return; } + + const bool query = combined < QHeads; + const int head = query ? combined : combined - QHeads; + const int heads = query ? QHeads : KHeads; + const __nv_bfloat162* __restrict__ input = query ? q_in : k_in; + const __nv_bfloat162* __restrict__ weight = query ? q_norm : k_norm; + __nv_bfloat162* __restrict__ output = query ? q_out : k_out; + + const std::int64_t base = (static_cast(token) * heads + head) * kPairs; + const auto normalized = detail::rmsnorm_rope_d256_normalize(input, weight, base, lane); +#pragma unroll + for (int k = 1; k < 4; ++k) { output[base + lane + k * 32] = normalized.pair[k]; } + + const int coefficient_pair = (lane & (kHalfPair - 1)) * 2; + float s0 = 0.0F, c0 = 0.0F, s1 = 0.0F, c1 = 0.0F; + fixed_sincos(positions, tokens, token, coefficient_pair, &s0, &c0); + fixed_sincos(positions, tokens, token, coefficient_pair + 1, &s1, &c1); + output[base + lane] = + detail::rmsnorm_rope_d256_rotate(normalized.pair[0], c0, c1, s0, s1, lane); +} + } // namespace ninfer::ops diff --git a/src/ops/rmsnorm_rope/launch.cu b/src/ops/rmsnorm_rope/launch.cu index 15888bf617..b18f452ab8 100644 --- a/src/ops/rmsnorm_rope/launch.cu +++ b/src/ops/rmsnorm_rope/launch.cu @@ -20,6 +20,26 @@ void launch_fixed(const Tensor& positions, const Tensor* q_norm_weight, const Te static_cast<__nv_bfloat16*>(k.data)); } +// One warp per head, three heads per block. Measured optimum on sm_120a; the plateau is flat from +// two to nine heads per block, and both ends are worse - one warp per block does not hide the load +// latency, all eighteen heads in one block leaves four blocks for the whole card. +constexpr int kTextHeadsPerBlock = 3; + +template +void launch_text(const Tensor& positions, const Tensor& q_norm_weight, const Tensor& k_norm_weight, + const Tensor& q_in, const Tensor& k_in, Tensor& q_out, Tensor& k_out, + std::int32_t tokens, cudaStream_t stream) { + constexpr int kGroups = (QHeads + KHeads + kTextHeadsPerBlock - 1) / kTextHeadsPerBlock; + rmsnorm_rope_d256_text_kernel + <<(tokens * kGroups), kTextHeadsPerBlock * 32, 0, stream>>>( + static_cast(positions.data), + static_cast(q_norm_weight.data), + static_cast(k_norm_weight.data), + static_cast(q_in.data), + static_cast(k_in.data), static_cast<__nv_bfloat162*>(q_out.data), + static_cast<__nv_bfloat162*>(k_out.data), tokens); +} + } // namespace void rmsnorm_rope_pair_launch(const Tensor& positions, const Tensor& q_norm_weight, @@ -37,4 +57,18 @@ void rmsnorm_rope_single_launch(const Tensor& positions, const Tensor& norm_weig CUDA_CHECK(cudaGetLastError()); } +void rmsnorm_rope_text_launch(const Tensor& positions, const Tensor& q_norm_weight, + const Tensor& k_norm_weight, const Tensor& q_in, const Tensor& k_in, + Tensor& q_out, Tensor& k_out, std::int32_t tokens, + cudaStream_t stream) { + if (q_in.ne[1] == 16) { + launch_text<16, 2>(positions, q_norm_weight, k_norm_weight, q_in, k_in, q_out, k_out, + tokens, stream); + } else { + launch_text<24, 4>(positions, q_norm_weight, k_norm_weight, q_in, k_in, q_out, k_out, + tokens, stream); + } + CUDA_CHECK(cudaGetLastError()); +} + } // namespace ninfer::ops::detail diff --git a/src/ops/rmsnorm_rope/launch.h b/src/ops/rmsnorm_rope/launch.h index 5d3d5af5b2..35c19ba892 100644 --- a/src/ops/rmsnorm_rope/launch.h +++ b/src/ops/rmsnorm_rope/launch.h @@ -15,4 +15,9 @@ void rmsnorm_rope_pair_launch(const Tensor& positions, const Tensor& q_norm_weig void rmsnorm_rope_single_launch(const Tensor& positions, const Tensor& norm_weight, Tensor& x, std::int32_t tokens, cudaStream_t stream); +void rmsnorm_rope_text_launch(const Tensor& positions, const Tensor& q_norm_weight, + const Tensor& k_norm_weight, const Tensor& q_in, const Tensor& k_in, + Tensor& q_out, Tensor& k_out, std::int32_t tokens, + cudaStream_t stream); + } // namespace ninfer::ops::detail diff --git a/src/ops/rmsnorm_rope/rmsnorm_rope.cpp b/src/ops/rmsnorm_rope/rmsnorm_rope.cpp index 0566fa46a6..886156b6ce 100644 --- a/src/ops/rmsnorm_rope/rmsnorm_rope.cpp +++ b/src/ops/rmsnorm_rope/rmsnorm_rope.cpp @@ -16,6 +16,11 @@ constexpr std::int32_t kQueryHeads = 32; constexpr std::int32_t kKeyHeads = 8; constexpr std::int32_t kMaximumBatch = 8; constexpr std::int32_t kMaximumSingle = 2048; +constexpr std::int32_t kTextHeadDim = 256; +// The text form has no width of its own to cap: one warp owns one head, so the only ceiling is the +// launch grid, and even the largest supported context stays four orders of magnitude below it. +constexpr std::int64_t kMaximumTextGrid = 2147483647; +constexpr std::int32_t kMaximumTextHeadGroups = 10; bool aligned_to(const void* pointer, std::uintptr_t alignment) { return pointer != nullptr && (reinterpret_cast(pointer) & (alignment - 1)) == 0; @@ -57,6 +62,21 @@ void require_single_nonoverlap(const Tensor& positions, const Tensor& norm_weigh } } +void require_text_nonoverlap(const Tensor& positions, const Tensor& q_norm_weight, + const Tensor& k_norm_weight, const Tensor& q_in, const Tensor& k_in, + const Tensor& q_out, const Tensor& k_out) { + for (const Tensor* mutable_tensor : {&q_out, &k_out}) { + for (const Tensor* other : {&q_in, &k_in, &positions, &q_norm_weight, &k_norm_weight}) { + if (overlaps(*mutable_tensor, *other)) { + throw std::invalid_argument("rmsnorm_rope: text output overlaps an input"); + } + } + } + if (overlaps(q_out, k_out)) { + throw std::invalid_argument("rmsnorm_rope: text outputs overlap each other"); + } +} + } // namespace void rmsnorm_rope(const Tensor& positions, const Tensor& q_norm_weight, const Tensor& k_norm_weight, @@ -90,4 +110,30 @@ void rmsnorm_rope(const Tensor& positions, const Tensor& norm_weight, Tensor& x, detail::rmsnorm_rope_single_launch(positions, norm_weight, x, tokens, stream); } +void rmsnorm_rope(const Tensor& positions, const Tensor& q_norm_weight, const Tensor& k_norm_weight, + const Tensor& q_in, const Tensor& k_in, Tensor& q_out, Tensor& k_out, + cudaStream_t stream) { + const std::int32_t tokens = q_in.ne[2]; + const std::int32_t query_heads = q_in.ne[1]; + const std::int32_t key_heads = k_in.ne[1]; + if (tokens < 1 || + static_cast(tokens) * kMaximumTextHeadGroups > kMaximumTextGrid) { + throw std::invalid_argument( + "rmsnorm_rope: text T must be positive and fit the launch grid"); + } + if (!((query_heads == 16 && key_heads == 2) || (query_heads == 24 && key_heads == 4))) { + throw std::invalid_argument("rmsnorm_rope: text (Q,K) must be (16,2) or (24,4)"); + } + require_tensor(q_in, DType::BF16, {kTextHeadDim, query_heads, tokens, 1}, "text q in"); + require_tensor(k_in, DType::BF16, {kTextHeadDim, key_heads, tokens, 1}, "text k in"); + require_tensor(q_out, DType::BF16, {kTextHeadDim, query_heads, tokens, 1}, "text q out"); + require_tensor(k_out, DType::BF16, {kTextHeadDim, key_heads, tokens, 1}, "text k out"); + require_tensor(q_norm_weight, DType::BF16, {kTextHeadDim, 1, 1, 1}, "text q norm weight"); + require_tensor(k_norm_weight, DType::BF16, {kTextHeadDim, 1, 1, 1}, "text k norm weight"); + require_tensor(positions, DType::I32, {tokens, 1, 1, 1}, "text positions"); + require_text_nonoverlap(positions, q_norm_weight, k_norm_weight, q_in, k_in, q_out, k_out); + detail::rmsnorm_rope_text_launch(positions, q_norm_weight, k_norm_weight, q_in, k_in, q_out, + k_out, tokens, stream); +} + } // namespace ninfer::ops diff --git a/tests/ops/test_rmsnorm_rope.cpp b/tests/ops/test_rmsnorm_rope.cpp index 8af64c6dc3..5223d1f27e 100644 --- a/tests/ops/test_rmsnorm_rope.cpp +++ b/tests/ops/test_rmsnorm_rope.cpp @@ -1,5 +1,7 @@ #include "ninfer/ops/rmsnorm_rope.h" #include "core/decode_graph.h" +#include "ninfer/ops/rmsnorm.h" +#include "ninfer/ops/rope.h" #include "ops/op_tester.h" #include @@ -25,6 +27,18 @@ constexpr double kTheta = 1.0e7; constexpr double kRelativeL2 = 1.85e-3; constexpr double kPairRelative = 6.9e-3; +// Text profile: wider head, narrower rotation, out of place. +constexpr int kTextHeadDim = 256; +constexpr int kTextRotaryDim = 64; +// The text profile normalizes over twice as many channels, so its FP32 reduction sits further +// from the FP64 oracle than the D128 profile does. Both limits are the measured worst case over +// the cases below with margin (relative L2 1.99e-3, pair ratio 1.17e-2). They are properties of +// this route rather than of the fusion: the fused result is bit-identical to +// rmsnorm -> rmsnorm -> rope, which the same limits therefore have to admit, and the test checks +// that equality separately. +constexpr double kTextRelativeL2 = 2.5e-3; +constexpr double kTextPairRelative = 1.4e-2; + struct OracleResult { std::vector output; std::vector pair_scale; @@ -105,7 +119,8 @@ OracleResult fused_oracle(const std::vector& input, const std::vector& got, - const OracleResult& expected) { + const OracleResult& expected, double pair_relative = kPairRelative, + double relative_l2_limit = kRelativeL2) { if (got.size() != expected.output.size() || got.size() != expected.pair_scale.size()) { std::cerr << label << ": result size mismatch\n"; return 1; @@ -121,7 +136,7 @@ int verify_profile(const std::string& label, const std::vector& got, } const double error = std::abs(got[index] - expected.output[index]); const double scale = expected.pair_scale[index]; - const double limit = kPairRelative * scale; + const double limit = pair_relative * scale; const double ratio = limit == 0.0 ? (error == 0.0 ? 0.0 : std::numeric_limits::infinity()) : error / limit; @@ -137,8 +152,9 @@ int verify_profile(const std::string& label, const std::vector& got, reference_square_sum += expected.output[index] * expected.output[index]; } const double relative_l2 = std::sqrt(error_square_sum / reference_square_sum); - if (relative_l2 > kRelativeL2) { - std::cerr << label << ": relative L2=" << relative_l2 << " exceeds " << kRelativeL2 << '\n'; + if (relative_l2 > relative_l2_limit) { + std::cerr << label << ": relative L2=" << relative_l2 << " exceeds " << relative_l2_limit + << '\n'; ++violations; } if (error_stats_enabled()) { @@ -282,6 +298,199 @@ int run_single_case(int tokens, int first_position, std::uint32_t seed, bool gra return failures; } +std::size_t text_index(int heads, int token, int head, int dim) { + return (static_cast(token) * heads + head) * kTextHeadDim + dim; +} + +// Independent FP64 oracle for the text formula: RMSNorm over 256 channels, split-half rotation +// over the first 64 channels, pass-through for the remaining 192. +OracleResult text_oracle(const std::vector& input, const std::vector& weight, + const std::vector& positions, int heads) { + const int tokens = static_cast(positions.size()); + OracleResult result{ + .output = std::vector(input.size()), + .pair_scale = std::vector(input.size()), + }; + std::vector normalized(kTextHeadDim); + for (int token = 0; token < tokens; ++token) { + for (int head = 0; head < heads; ++head) { + double sum_squares = 0.0; + for (int dim = 0; dim < kTextHeadDim; ++dim) { + const double value = input[text_index(heads, token, head, dim)]; + sum_squares += value * value; + } + const double inverse = + 1.0 / std::sqrt(sum_squares / static_cast(kTextHeadDim) + kEpsilon); + for (int dim = 0; dim < kTextHeadDim; ++dim) { + // Offset epilogue: the stored weight is a delta around one. + normalized[static_cast(dim)] = + static_cast(input[text_index(heads, token, head, dim)]) * inverse * + (static_cast(weight[static_cast(dim)]) + 1.0); + } + for (int dim = kTextRotaryDim; dim < kTextHeadDim; ++dim) { + const std::size_t index = text_index(heads, token, head, dim); + const double value = normalized[static_cast(dim)]; + result.output[index] = value; + result.pair_scale[index] = std::abs(value); + } + for (int pair = 0; pair < kTextRotaryDim / 2; ++pair) { + const double exponent = -2.0 * static_cast(pair) / kTextRotaryDim; + const double phase = + static_cast(positions[static_cast(token)]) * + std::pow(kTheta, exponent); + const double cosine = std::cos(phase); + const double sine = std::sin(phase); + const double first = normalized[static_cast(pair)]; + const double second = + normalized[static_cast(pair + kTextRotaryDim / 2)]; + const double scale = std::hypot(first, second); + const std::size_t first_index = text_index(heads, token, head, pair); + const std::size_t second_index = + text_index(heads, token, head, pair + kTextRotaryDim / 2); + result.output[first_index] = first * cosine - second * sine; + result.output[second_index] = second * cosine + first * sine; + result.pair_scale[first_index] = scale; + result.pair_scale[second_index] = scale; + } + } + } + return result; +} + +// One text case, with two independent verdicts on the same inputs: +// (a) the FP64 oracle, which says the Op computes the documented formula; +// (b) bit equality against rmsnorm -> rmsnorm -> rope, which says it computes it the same way +// the three calls it replaces do. (b) is the property a caller relies on when it swaps one +// for the other, and no tolerance can stand in for it. +int run_text_case(int query_heads, int key_heads, int tokens, int first_position, + std::uint32_t seed, bool graph = false) { + const std::size_t q_count = static_cast(kTextHeadDim) * query_heads * tokens; + const std::size_t k_count = static_cast(kTextHeadDim) * key_heads * tokens; + const auto q = make_bf16_values(q_count, seed, -4.0F, 4.0F); + const auto k = make_bf16_values(k_count, seed + 1U, -4.0F, 4.0F); + const auto q_weight = make_bf16_values(kTextHeadDim, seed + 2U, 0.25F, 1.75F); + const auto k_weight = make_bf16_values(kTextHeadDim, seed + 3U, 0.25F, 1.75F); + const auto positions = make_positions(tokens, first_position); + const OracleResult q_expected = text_oracle(q, q_weight, positions, query_heads); + const OracleResult k_expected = text_oracle(k, k_weight, positions, key_heads); + const auto q_bits = bf16_bits(q); + const auto k_bits = bf16_bits(k); + const auto q_weight_bits = bf16_bits(q_weight); + const auto k_weight_bits = bf16_bits(k_weight); + + DeviceBuffer q_in_device = to_device(q_bits); + DeviceBuffer k_in_device = to_device(k_bits); + DeviceBuffer q_weight_device = to_device(q_weight_bits); + DeviceBuffer k_weight_device = to_device(k_weight_bits); + DeviceBuffer position_device = to_device(positions); + GuardedDeviceBuffer q_out_device(q_count * sizeof(std::uint16_t)); + GuardedDeviceBuffer k_out_device(k_count * sizeof(std::uint16_t)); + GuardedDeviceBuffer q_split_device(q_count * sizeof(std::uint16_t)); + GuardedDeviceBuffer k_split_device(k_count * sizeof(std::uint16_t)); + + Tensor q_in(q_in_device.p, DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor k_in(k_in_device.p, DType::BF16, {kTextHeadDim, key_heads, tokens}); + Tensor q_weight_tensor(q_weight_device.p, DType::BF16, {kTextHeadDim}); + Tensor k_weight_tensor(k_weight_device.p, DType::BF16, {kTextHeadDim}); + Tensor position_tensor(position_device.p, DType::I32, {tokens}); + Tensor q_out(q_out_device.data(), DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor k_out(k_out_device.data(), DType::BF16, {kTextHeadDim, key_heads, tokens}); + Tensor q_split(q_split_device.data(), DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor k_split(k_split_device.data(), DType::BF16, {kTextHeadDim, key_heads, tokens}); + + execute( + [&](cudaStream_t stream) { + ops::rmsnorm_rope(position_tensor, q_weight_tensor, k_weight_tensor, q_in, k_in, q_out, + k_out, stream); + }, + [](cudaStream_t) {}, graph); + + // The route this replaces, on the same inputs. + ops::rmsnorm(q_in, q_weight_tensor, static_cast(kEpsilon), true, q_split, nullptr); + ops::rmsnorm(k_in, k_weight_tensor, static_cast(kEpsilon), true, k_split, nullptr); + ops::rope(position_tensor, kTextRotaryDim, static_cast(kTheta), q_split, k_split, + nullptr); + cuda_synchronize(); + + const std::string label = "rmsnorm_rope text Q=" + std::to_string(query_heads) + + " K=" + std::to_string(key_heads) + + " graph=" + std::to_string(graph) + " T=" + std::to_string(tokens) + + " P=" + std::to_string(first_position); + int failures = verify_profile(label + " q", from_device_bf16(q_out_device.data(), q_count), + q_expected, kTextPairRelative, kTextRelativeL2); + failures += verify_profile(label + " k", from_device_bf16(k_out_device.data(), k_count), + k_expected, kTextPairRelative, kTextRelativeL2); + failures += verify_exact((label + " q equals split route").c_str(), + from_device(q_out_device.data(), q_count), + from_device(q_split_device.data(), q_count)); + failures += verify_exact((label + " k equals split route").c_str(), + from_device(k_out_device.data(), k_count), + from_device(k_split_device.data(), k_count)); + failures += q_out_device.verify_guards(label + " q guards"); + failures += k_out_device.verify_guards(label + " k guards"); + failures += verify_exact((label + " q input unchanged").c_str(), + from_device(q_in_device, q_bits.size()), q_bits); + failures += verify_exact((label + " k input unchanged").c_str(), + from_device(k_in_device, k_bits.size()), k_bits); + failures += + verify_exact((label + " positions").c_str(), + from_device(position_device, positions.size()), positions); + return failures; +} + +// A prefill chunk may be any positive multiple of 128, so the Op has to take widths far past the +// ones the oracle can afford to check. Here the reference is the split route only: the FP64 oracle +// already covers the arithmetic at the widths above, and what is at stake here is dispatch. +int run_text_wide_case(int query_heads, int key_heads, int tokens, std::uint32_t seed) { + const std::size_t q_count = static_cast(kTextHeadDim) * query_heads * tokens; + const std::size_t k_count = static_cast(kTextHeadDim) * key_heads * tokens; + const auto q_bits = bf16_bits(make_bf16_values(q_count, seed, -4.0F, 4.0F)); + const auto k_bits = bf16_bits(make_bf16_values(k_count, seed + 1U, -4.0F, 4.0F)); + const auto q_weight_bits = bf16_bits(make_bf16_values(kTextHeadDim, seed + 2U, 0.25F, 1.75F)); + const auto k_weight_bits = bf16_bits(make_bf16_values(kTextHeadDim, seed + 3U, 0.25F, 1.75F)); + const auto positions = make_positions(tokens, 0); + + DeviceBuffer q_in_device = to_device(q_bits); + DeviceBuffer k_in_device = to_device(k_bits); + DeviceBuffer q_weight_device = to_device(q_weight_bits); + DeviceBuffer k_weight_device = to_device(k_weight_bits); + DeviceBuffer position_device = to_device(positions); + GuardedDeviceBuffer q_out_device(q_count * sizeof(std::uint16_t)); + GuardedDeviceBuffer k_out_device(k_count * sizeof(std::uint16_t)); + GuardedDeviceBuffer q_split_device(q_count * sizeof(std::uint16_t)); + GuardedDeviceBuffer k_split_device(k_count * sizeof(std::uint16_t)); + + Tensor q_in(q_in_device.p, DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor k_in(k_in_device.p, DType::BF16, {kTextHeadDim, key_heads, tokens}); + Tensor q_weight_tensor(q_weight_device.p, DType::BF16, {kTextHeadDim}); + Tensor k_weight_tensor(k_weight_device.p, DType::BF16, {kTextHeadDim}); + Tensor position_tensor(position_device.p, DType::I32, {tokens}); + Tensor q_out(q_out_device.data(), DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor k_out(k_out_device.data(), DType::BF16, {kTextHeadDim, key_heads, tokens}); + Tensor q_split(q_split_device.data(), DType::BF16, {kTextHeadDim, query_heads, tokens}); + Tensor k_split(k_split_device.data(), DType::BF16, {kTextHeadDim, key_heads, tokens}); + + ops::rmsnorm_rope(position_tensor, q_weight_tensor, k_weight_tensor, q_in, k_in, q_out, k_out, + nullptr); + ops::rmsnorm(q_in, q_weight_tensor, static_cast(kEpsilon), true, q_split, nullptr); + ops::rmsnorm(k_in, k_weight_tensor, static_cast(kEpsilon), true, k_split, nullptr); + ops::rope(position_tensor, kTextRotaryDim, static_cast(kTheta), q_split, k_split, + nullptr); + cuda_synchronize(); + + const std::string label = + "rmsnorm_rope text wide Q=" + std::to_string(query_heads) + " T=" + std::to_string(tokens); + int failures = verify_exact((label + " q equals split route").c_str(), + from_device(q_out_device.data(), q_count), + from_device(q_split_device.data(), q_count)); + failures += verify_exact((label + " k equals split route").c_str(), + from_device(k_out_device.data(), k_count), + from_device(k_split_device.data(), k_count)); + failures += q_out_device.verify_guards((label + " q guards").c_str()); + failures += k_out_device.verify_guards((label + " k guards").c_str()); + return failures; +} + } // namespace int main() { @@ -305,6 +514,19 @@ int main() { failures += run_single_case(1024, 130'048, 0x2004U); failures += run_single_case(2048, 260'032, 0x2005U); + // Text profile: both registered head geometries, widths from one token to a prefill chunk, + // and both graph modes. + for (const int tokens : {1, 2, 3, 4, 7, 8, 16, 17, 64, 128, 129, 1024, 4096}) { + failures += run_text_case(16, 2, tokens, tokens == 1 ? 0 : 131'072, 0x3000U + tokens); + failures += run_text_case(24, 4, tokens, tokens == 1 ? 0 : 262'000, 0x4000U + tokens); + } + failures += run_text_case(16, 2, 4, 0, 0x3101U, true); + failures += run_text_case(24, 4, 16, 262'000, 0x4101U, true); + // Past the widths the FP64 oracle can afford, and past any ceiling of our own: a prefill + // chunk is only required to be a positive multiple of 128. + failures += run_text_wide_case(16, 2, 8320, 0x5001U); + failures += run_text_wide_case(24, 4, 16384, 0x5002U); + if (failures != 0) { std::cerr << "rmsnorm_rope failures=" << failures << '\n'; return 1;