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/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/src/targets/qwen3_6/impl/runtime/text_context_impl.h b/src/targets/qwen3_6/impl/runtime/text_context_impl.h index 8cdef782a1..a8bd731c9f 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context_impl.h +++ b/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(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(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 = 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;