Register the text profile of ops::rmsnorm_rope and route the two full-attention call sites to it

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 <cuda_profiler_api.h>
@@ -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<std::size_t>(kTextHeadDim) * query_heads * tokens;
+    const std::size_t k_elements = static_cast<std::size_t>(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<double>(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;

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

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 <stdexcept>
@@ -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

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

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});

src/ops/rmsnorm_rope/d256.cuh

@@ -0,0 +1,68 @@
+#pragma once
+
+#include "ops/common/warp.cuh"
+#include "ops/kernel/rmsnorm.cuh"
+
+#include <cuda_bf16.h>
+
+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<float>(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<RmsEpilogue::Offset>(xf.x, inv, wf.x, 0.0F),
+                                  rmsnorm_epilogue<RmsEpilogue::Offset>(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