From e222c6230e7ee71c2c14c62aeb66aeba8fcea48c Mon Sep 17 00:00:00 2001 From: MichaelDementii Date: Wed, 16 Sep 2026 21:12:37 +0200 Subject: [PATCH] perf(ops): register the text profile of rmsnorm_rope, and route to it Both full-attention call sites normalized the query, normalized the key and rotated both - three Op calls for one semantic step. ops::rmsnorm_rope already fuses exactly that and was already registered for two profiles, the DFlash2 pair at [128,32,W,B] and the context-K single at [128,8,T]; neither covers the text heads, which are 256 wide with a 64-channel rotation and want the result out of place. So the text stack kept issuing the three. On qwen3.6-35b-a3b that is ten of the forty text layers, plus the MTP tail once per draft step. A third overload takes q_in [256,Q,T] and k_in [256,K,T] and writes q_out and k_out, for (Q,K) of (16,2) or (24,4). Only the first 64 channels rotate; the remaining 192 carry the normalized value through. The norm weight enters as a delta around one - the Offset epilogue the text stack normalizes with - unlike the two in-place forms, which multiply by the stored weight directly. The rotation consumes the BF16 represented normalized value, so the result is bit-identical to rmsnorm(q_in) -> rmsnorm(k_in) -> rope(q_out, k_out). text_qk_norm_rope sits beside text_rope in execution/attention.cpp, where the MRoPE axis check already lives, so both call sites go through one place. It takes the fused Op where the profile matches and the three calls everywhere else: the MRoPE path, any other geometry, and above a width bound. The width bound is the part that had to be measured. The fused form wins by 41.94 to 52.07 % at decode widths and loses above about six hundred tokens: T (16,2) (24,4) 256 -30.01% -21.83% 320 -26.40% -0.90% 512 -10.99% -1.91% 576 -10.35% +2.91% 1024 +5.23% +21.87% 2048 +17.20% +32.05% The two geometries cross at different widths - (24,4) has already fallen to -0.90 % at T=320 and is positive from 576, while (16,2) stays ahead through 896 and turns at 960 - so the dispatch stops at 256, where the fused route is 30.01 % ahead on (16,2) and 21.83 % on (24,4). One width further out (24,4) is at -0.90 %, inside this stand's own null, so the margin is gone before the sign is. I have no established mechanism for the crossover; the obvious one does not predict the non-monotone approach to it, and nothing here measures occupancy. The predicate covers every model with a card in the tree: three base models, all head_dim 256 and rotary 64, and exactly two head geometries. Qwen3.6-27B and Qwen3.8-27B are (24,4) with 16 full-attention layers of 64; Qwen3.6-35B-A3B is (16,2) with 10 of 40. The call sits on the full-attention path, so it fires 16 or 10 times per pass, not once per layer. The Op itself is valid at any width - the test drives it directly out to 16384 - and this is a dispatch choice. Decode issues T = W x B, at most 128, so the bound never excludes it; a prefill chunk is a multiple of 128 and the default is 1024, so prefill keeps the three calls. The benchmark gains the text profile and a --route split|fused switch, so both routes live in one binary and the comparison is one process per pass: split issues the three calls, fused issues the Op, and split again is the identical-baseline null. Six passes with the arm order rotated through six permutations, graph execution, 20 warmups and 200 samples over 32 inner repetitions. A pass is dropped when the two identical runs disagree by more than 1 %: 29 of 240 on the wide sweep and 19 of 252 on the dense sweep of the crossover region. The rule bounds what survives, so the width of the instrument is the worst kept pass: 0.97 % and 0.96 %, with medians of 0.15 % and 0.18 %. band cells median best worst faster taken by the dispatch, T <= 256 38 -43.27% -52.07% -21.83% 38/38 declined, keeps the three calls 38 -3.58% -26.85% +32.05% 25/38 T (16,2) split -> fused (24,4) split -> fused 1 3.72 -> 2.17 us 3.81 -> 2.20 us 8 4.25 -> 2.19 us 4.50 -> 2.20 us 16 4.36 -> 2.18 us 4.61 -> 2.23 us 64 4.67 -> 2.52 us 5.02 -> 2.78 us 128 4.88 -> 2.90 us 5.34 -> 3.38 us 256 5.40 -> 3.78 us 6.16 -> 4.81 us The fused form is flat at about 2.2 us from T=1 to T=16 on both geometries while the traffic it moves grows about sixteen-fold, so its time is independent of the work; (24,4) leaves the plateau first, at 2.36 us by T=24. The split form climbs over the same range, 3.72 to 4.36 us on (16,2) and 3.81 to 4.61 on (24,4). No width the dispatch takes is slower on either geometry. Output is unchanged: NINFER_OP_REPORT_STATS=1 at %.17g over the whole test suite gives 11912 cases shared between the arms, all byte-identical, and the 56 cases the candidate adds are the text-profile checks the test gains. The same comparison on the previous base dc58675f read 11900 shared, 0 differing, 56 added - master gained 12 cases from its own LinearSwiGLu retune, and the part being claimed did not move. The test judges the new form against the independent FP64 oracle and, separately, against the three calls it replaces, bit for bit. Strength control: bit-exactness is the claim, so the test has to see a perturbation far below any oracle criterion. Scaling the rotation by one part per million - about 3900x below one BF16 ulp - fails 33 of the 60 exact-equality comparisons the suite makes against the split route, over 20 of its 30 text cases; restoring it passes. It does not fail everywhere, and should not: that far below the output quantum the nudge only moves a value already sitting on a rounding boundary. ctest 120/120 on both arms, on both bases. No end-to-end number: an identical-baseline decode arm on this machine reads +0.55 % median over six passes spanning -0.44 % to +0.89 %, which is wider than a change that saves tens of microseconds per pass. That arm was collected for another submission on the same stand and is quoted as a property of the stand, not as a measurement of this change. An earlier collection of the same arm at two repetitions read +1.81 % median with a +5.41 % worst pass; ten repetitions and three warmups is what tightened it. --- bench/ops/rmsnorm_rope_bench.cu | 91 +++++++- include/ninfer/ops/rmsnorm_rope.h | 28 +++ src/models/qwen3_5/execution/attention.cpp | 38 ++++ src/models/qwen3_5/execution/attention.h | 9 + src/models/qwen3_5/execution/text.cpp | 11 +- src/ops/rmsnorm_rope/d256.cuh | 68 ++++++ src/ops/rmsnorm_rope/kernel.cuh | 45 ++++ src/ops/rmsnorm_rope/launch.cu | 34 +++ src/ops/rmsnorm_rope/launch.h | 5 + src/ops/rmsnorm_rope/rmsnorm_rope.cpp | 46 +++++ tests/ops/test_rmsnorm_rope.cpp | 230 ++++++++++++++++++++- 11 files changed, 590 insertions(+), 15 deletions(-) create mode 100644 src/ops/rmsnorm_rope/d256.cuh 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;