Extend rmsnorm_rope to the text profile: one graph node instead of three, +0.60% decode on Qwen3.6-35B-A3B

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/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

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 <cuda_bf16.h>
 #include <cstdint>
 
@@ -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 <int QHeads, int KHeads, int HeadsPerBlock>
+__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<int>(blockIdx.x) / kGroups;
+    if (token >= tokens) { return; }
+    const int group    = static_cast<int>(blockIdx.x) % kGroups;
+    const int lane     = static_cast<int>(threadIdx.x) & 31;
+    const int warp     = static_cast<int>(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<std::int64_t>(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<RopeKernelMode::Text1D>(positions, tokens, token, coefficient_pair, &s0, &c0);
+    fixed_sincos<RopeKernelMode::Text1D>(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

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 <int QHeads, int KHeads>
+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<QHeads, KHeads, kTextHeadsPerBlock>
+        <<<static_cast<unsigned>(tokens * kGroups), kTextHeadsPerBlock * 32, 0, stream>>>(
+            static_cast<const std::int32_t*>(positions.data),
+            static_cast<const __nv_bfloat162*>(q_norm_weight.data),
+            static_cast<const __nv_bfloat162*>(k_norm_weight.data),
+            static_cast<const __nv_bfloat162*>(q_in.data),
+            static_cast<const __nv_bfloat162*>(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

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

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<std::uintptr_t>(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<std::int64_t>(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

src/targets/qwen3_6/impl/runtime/text_context_impl.h

@@ -23,6 +23,7 @@
 #include "ninfer/ops/position.h"
 #include "ninfer/ops/residual_add.h"
 #include "ninfer/ops/rmsnorm.h"
+#include "ninfer/ops/rmsnorm_rope.h"
 #include "ninfer/ops/rope.h"
 #include "ninfer/ops/sparse_moe.h"
 #include "ninfer/ops/scatter.h"
@@ -45,6 +46,32 @@
 namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule {
 namespace {
 
+// The fused Q/K norm + RoPE Op covers the two text head geometries with a 1-D position axis. The
+// mrope path and any future geometry keep the three calls it replaces. Both branches are the same
+// arithmetic - the Op is bit-exact against them - so this chooses a schedule, not a result.
+inline constexpr bool kFusedQkNormRope =
+    kCfg.head_dim == 256 && kCfg.rotary_dim == 64 &&
+    ((kCfg.n_q == 16 && kCfg.n_kv == 2) || (kCfg.n_q == 24 && kCfg.n_kv == 4));
+
+void split_qk_norm_rope(const Tensor& positions, const Tensor& q_norm, const Tensor& k_norm,
+                        const Tensor& q, const Tensor& k, Tensor& qn, Tensor& kn,
+                        cudaStream_t stream) {
+    ops::rmsnorm(q, q_norm, kCfg.rms_eps, true, qn, stream);
+    ops::rmsnorm(k, k_norm, kCfg.rms_eps, true, kn, stream);
+    ops::rope(positions, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, stream);
+}
+
+void qk_norm_rope(const Tensor& positions, const Tensor& q_norm, const Tensor& k_norm,
+                  const Tensor& q, const Tensor& k, Tensor& qn, Tensor& kn, cudaStream_t stream) {
+    if constexpr (kFusedQkNormRope) {
+        if (positions.ne[1] == 1) {
+            ops::rmsnorm_rope(positions, q_norm, k_norm, q, k, qn, kn, stream);
+            return;
+        }
+    }
+    split_qk_norm_rope(positions, q_norm, k_norm, q, k, qn, kn, stream);
+}
+
 void copy_i32(const std::int32_t* source, Tensor& destination, cudaStream_t stream) {
     if (source == nullptr || destination.dtype != DType::I32 || !destination.is_contiguous() ||
         destination.data == nullptr) {
@@ -380,10 +407,8 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po
     const auto results = workspace_recipe::mtp_attention_results<TextConfig>(work_, T);
     Tensor qn          = results.normalized_query.view({kCfg.head_dim, kCfg.n_q, T});
     Tensor kn          = results.normalized_key.view({kCfg.head_dim, kCfg.n_kv, T});
-    ops::rmsnorm(q, *mtp_.q_norm, kCfg.rms_eps, true, qn, s);
-    ops::rmsnorm(k, *mtp_.k_norm, kCfg.rms_eps, true, kn, s);
     Tensor rope_for_op = active_sequence_batch_ != 0 ? rope_positions.view({T}) : rope_positions;
-    ops::rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s);
+    qk_norm_rope(rope_for_op, *mtp_.q_norm, *mtp_.k_norm, q, k, qn, kn, s);
 
     Tensor a = results.attention.view({kCfg.head_dim, kCfg.n_q, T});
     if (active_sequence_batch_ != 0) {
@@ -839,14 +864,12 @@ void TextContext::attn_mix(const FullLayerW& w, Tensor& x, int fidx, Phase ph) {
     const auto results = workspace_recipe::text_attention_results<TextConfig>(work_, T);
     Tensor qn          = results.normalized_query.view({kCfg.head_dim, kCfg.n_q, T});
     Tensor kn          = results.normalized_key.view({kCfg.head_dim, kCfg.n_kv, T});
-    ops::rmsnorm(q, *w.q_norm, kCfg.rms_eps, true, qn, s);
-    ops::rmsnorm(k, *w.k_norm, kCfg.rms_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;
-    ops::rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s);
+    qk_norm_rope(rope_for_op, *w.q_norm, *w.k_norm, q, k, qn, kn, s);
 
     Tensor a = results.attention.view({kCfg.head_dim, kCfg.n_q, T});
     const Tensor& kv_table_rows =