PR5: static YaRN factor-4 rope extension

apps/cli/main.cpp

@@ -282,6 +282,7 @@ int main(int argc, char** argv) {
         engine_options.kv_cache       = cli.kv_cache;
         engine_options.speculative    = cli.speculative;
         engine_options.enable_vision  = cli.enable_vision;
+        engine_options.yarn_enabled  = cli.yarn_enabled;
         engine_options.use_cuda_graph = cli.use_cuda_graph;
         // One CLI invocation owns exactly one request, so retained cross-request context has no
         // consumer and must not reserve an extra Device StateImage or run terminal capture.

apps/cli/options.cpp

@@ -154,6 +154,8 @@ Options parse_options(int argc, char** argv) {
             options.enable_vision = true;
         } else if (arg == "--no-cuda-graph") {
             options.use_cuda_graph = false;
+        } else if (arg == "--yarn") {
+            options.yarn_enabled = true;
         } else if (arg == "--stop-token-id") {
             const std::uint32_t token = parse_u32(value(arg), "stop-token-id", true);
             if (token > static_cast<std::uint32_t>(std::numeric_limits<TokenId>::max())) {
@@ -207,8 +209,8 @@ Options parse_options(int argc, char** argv) {
         throw std::invalid_argument("--prefill-chunk must be a multiple of 128");
     }
     if (options.kv_capacity.mode == KvCapacityMode::Explicit &&
-        options.kv_capacity.explicit_tokens < options.max_context) {
-        throw std::invalid_argument("--kv-capacity must be at least --max-context");
+        options.kv_capacity.explicit_tokens == 0) {
+        throw std::invalid_argument("--kv-capacity must be positive");
     }
     product::validate_speculative_cli_options(options.speculative);
     if (options.speculative.backend == SpeculativeBackend::DFlash && options.enable_vision) {

apps/cli/options.h

@@ -27,6 +27,7 @@ struct Options {
     SpeculativeOptions speculative;
     bool enable_vision  = false;
     bool use_cuda_graph = true;
+    bool yarn_enabled     = false;
 
     bool raw_output      = false;
     bool print_token_ids = false;

include/ninfer/ops/rope.h

@@ -40,4 +40,17 @@ void rope(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tenso
 // from x; Q versus K role does not change the transformation.
 void rope(const Tensor& positions, int rotary_dim, float theta, Tensor& x, cudaStream_t stream);
 
+/**
+ * Static YaRN factor-4 rope: the Text 1-D / DFlash 1-D transformations with the yarn4 frequency
+ * tables (theta must be 1e7) and the yarn4 attention scaling 1.1386 folded into the sincos
+ * coefficients. Registered domains are Text D256/R64 (heads 24/4, 16/2) and DFlash D128/R128
+ * (32/8). Same storage contract as rope().
+ */
+void rope_yarn4(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tensor& k,
+                cudaStream_t stream);
+
+// Single-tensor form of rope_yarn4.
+void rope_yarn4(const Tensor& positions, int rotary_dim, float theta, Tensor& x,
+                cudaStream_t stream);
+
 } // namespace ninfer::ops

include/ninfer/ops/softmax_attention.h

@@ -13,7 +13,9 @@
 
 namespace ninfer::ops {
 
+// Native causal-attention key budget; the YaRN extension multiplies it by 4.
 inline constexpr std::uint32_t kCausalAttentionMaximumVisibleKeys = 262144;
+inline constexpr std::uint32_t kCausalAttentionMaximumVisibleKeysYarn = 4 * 262144;
 
 struct CausalAttentionExecutionEnvelope {
     std::uint32_t min_visible_keys = 0;

include/ninfer/types.h

@@ -122,6 +122,7 @@ struct EngineOptions {
     std::uint32_t media_preprocess_threads = 0;
     bool enable_vision                     = false;
     bool use_cuda_graph                    = true;
+    bool yarn_enabled                      = false;
     ContextCacheOptions context_cache;
     ContextCostOptions context_cost;
     LoadProgress load_progress;

src/core/device.h

@@ -15,6 +15,8 @@ struct DeviceContext {
     cudaStream_t stream          = nullptr;
     cudaStream_t transfer_stream = nullptr;
     cudaDeviceProp props{};
+    // Static YaRN factor-4 rope extension (Text D256/R64 and DFlash D128/R128).
+    bool yarn_enabled = false;
 
     explicit DeviceContext(int device_id = 0);
     ~DeviceContext();

src/ops/kernel/rope.cuh

@@ -14,7 +14,9 @@ namespace ninfer::ops {
 
 enum class RopeKernelMode : std::int32_t {
     Text1D,
+    Text1DYarn4,
     DflashText1D,
+    DflashText1DYarn4,
     TextMrope,
     Vision2D,
 };
@@ -56,6 +58,41 @@ static __device__ __constant__ double kDflashRopeInvFrequency[64] = {
     1.28639694493697462e-07,
 };
 
+static __device__ __constant__ double kTextRopeYarn4InvFrequency[32] = {
+    1.000000000e+00, 6.042963902e-01, 3.651741273e-01, 2.206734069e-01, 1.333521432e-01,
+    8.058421878e-02, 4.869675252e-02, 2.942727176e-02, 1.778279410e-02, 1.074607828e-02,
+    6.493816316e-03, 3.924189758e-03, 2.371373706e-03, 1.433012570e-03, 8.659643234e-04,
+    4.742398227e-04, 2.569350599e-04, 1.373497451e-04, 7.217387404e-05, 3.707224982e-05,
+    1.844922203e-05, 8.759770071e-06, 3.849816315e-06, 2.326430102e-06, 1.405853313e-06,
+    8.495520822e-07, 5.133812566e-07, 3.102344402e-07, 1.874735523e-07, 1.132895909e-07,
+    6.846049086e-08, 4.137042750e-08,
+};
+
+static __device__ __constant__ double kDflashRopeYarn4InvFrequency[64] = {
+    1.00000000000000000e+00, 7.77365030238775789e-01, 6.04296390238132863e-01,
+    4.69758881670649164e-01, 3.65174127254837722e-01, 2.83873596475875456e-01,
+    2.20673406908458991e-01, 1.71543789634287902e-01, 1.33352143216332403e-01,
+    1.03663292843769794e-01, 8.05842187761481865e-02, 6.26433536656885587e-02,
+    4.86967525165863113e-02, 3.78551524925863012e-02, 2.94272717620928173e-02,
+    2.28757320031839559e-02, 1.77827941003892293e-02, 1.38237222735789964e-02,
+    1.07460782832131743e-02, 8.35362546957826163e-03, 6.49381631576211298e-03,
+    5.04806571666747105e-03, 3.92418975848453627e-03, 3.05052789026702539e-03,
+    2.37137370566165538e-03, 1.84342299240911056e-03, 1.43301257023696268e-03,
+    1.11397385999480246e-03, 8.65964323360065387e-04, 6.73170382414498242e-04,
+    5.23299114681494734e-04, 4.06794432108304740e-04, 3.16227766016837939e-04,
+    2.45824406892019762e-04, 1.91095297497044048e-04, 1.48550801717277505e-04,
+    1.15478198468945822e-04, 8.97687132447314224e-05, 6.97830584859866353e-05,
+    5.42469093701132573e-05, 4.21696503428582224e-05, 3.27812115139345850e-05,
+    2.54829674797934641e-05, 1.98095677855033870e-05, 1.53992652605949185e-05,
+    1.19708503049572999e-05, 9.30572040929699043e-06, 7.23394162736674728e-06,
+    5.62341325190349121e-06, 4.37144481261108992e-06, 3.39820832894255927e-06,
+    2.64164832038609264e-06, 2.05352502645714607e-06, 1.59633854428794220e-06,
+    1.24093776075171953e-06, 9.64661619911199141e-07, 7.49894209332455848e-07,
+    5.82941534713607427e-07, 4.53158363760081793e-07, 3.52269465147310129e-07,
+    2.73841963426436139e-07, 2.12875166179637264e-07, 1.65481709994318135e-07,
+    1.28639694493697462e-07,
+};
+
 static __device__ __constant__ float kVisionRopeInvFrequency[18] = {
     1.000000000e+00F, 5.994842503e-01F, 3.593813664e-01F, 2.154434690e-01F, 1.291549665e-01F,
     7.742636827e-02F, 4.641588834e-02F, 2.782559402e-02F, 1.668100537e-02F, 1.000000000e-02F,
@@ -68,9 +105,13 @@ __device__ __forceinline__ void fixed_axis_frequency(int pair, int* axis, float*
     if constexpr (Mode == RopeKernelMode::Vision2D) {
         *axis      = pair / 18;
         *frequency = kVisionRopeInvFrequency[pair % 18];
-    } else if constexpr (Mode == RopeKernelMode::DflashText1D) {
+    } else if constexpr (Mode == RopeKernelMode::DflashText1D ||
+                         Mode == RopeKernelMode::DflashText1DYarn4) {
         *axis      = 0;
         *frequency = static_cast<float>(kDflashRopeInvFrequency[pair]);
+    } else if constexpr (Mode == RopeKernelMode::Text1DYarn4) {
+        *axis      = 0;
+        *frequency = kTextRopeYarn4InvFrequency[pair];
     } else {
         *axis      = Mode == RopeKernelMode::TextMrope ? pair % 3 : 0;
         *frequency = kTextRopeInvFrequency[pair];
@@ -87,6 +128,30 @@ __device__ __forceinline__ void fixed_sincos(const std::int32_t* positions, int
         const double turns  = angle * kInvTwoPi;
         const float reduced = static_cast<float>(angle - nearbyint(turns) * kTwoPi);
         sincosf(reduced, sine, cosine);
+    } else if constexpr (Mode == RopeKernelMode::DflashText1DYarn4) {
+        constexpr double kInvTwoPi = 1.59154943091895336e-01;
+        constexpr double kTwoPi    = 6.28318530717958648e+00;
+        constexpr float kYarn4AttentionScaling = 1.138629436111989f;
+        const double angle =
+            static_cast<double>(positions[token]) *
+            static_cast<double>(kDflashRopeYarn4InvFrequency[pair]);
+        const double turns  = angle * kInvTwoPi;
+        const float reduced = static_cast<float>(angle - nearbyint(turns) * kTwoPi);
+        sincosf(reduced, sine, cosine);
+        *sine *= kYarn4AttentionScaling;
+        *cosine *= kYarn4AttentionScaling;
+    } else if constexpr (Mode == RopeKernelMode::Text1DYarn4) {
+        constexpr double kInvTwoPi = 1.59154943091895336e-01;
+        constexpr double kTwoPi    = 6.28318530717958648e+00;
+        constexpr float kYarn4AttentionScaling = 1.138629436111989f;
+        const double angle =
+            static_cast<double>(positions[token]) *
+            static_cast<double>(kTextRopeYarn4InvFrequency[pair]);
+        const double turns  = angle * kInvTwoPi;
+        const float reduced = static_cast<float>(angle - nearbyint(turns) * kTwoPi);
+        sincosf(reduced, sine, cosine);
+        *sine *= kYarn4AttentionScaling;
+        *cosine *= kYarn4AttentionScaling;
     } else {
         int axis = 0;
         float frequency;
@@ -120,10 +185,12 @@ __global__ void rope_fixed_kernel(const std::int32_t* positions, __nv_bfloat16*
                                   std::int64_t k_token_stride) {
     constexpr int kHeadDim = Mode == RopeKernelMode::Vision2D       ? 72
                              : Mode == RopeKernelMode::DflashText1D ? 128
-                                                                    : 256;
+                             : Mode == RopeKernelMode::DflashText1DYarn4 ? 128
+                                                                         : 256;
     constexpr int kHalf    = Mode == RopeKernelMode::Vision2D       ? 36
                              : Mode == RopeKernelMode::DflashText1D ? 64
-                                                                    : 32;
+                             : Mode == RopeKernelMode::DflashText1DYarn4 ? 64
+                                                                         : 32;
     const int token        = static_cast<int>(blockIdx.x);
     if (token >= tokens) { return; }
 

src/ops/launcher/rope.cu

@@ -5,6 +5,7 @@
 #include "ops/kernel/rope.cuh"
 
 #include <cstdint>
+#include <stdexcept>
 
 namespace ninfer::ops::detail {
 namespace {
@@ -18,8 +19,9 @@ constexpr int kLargeBlockWaveCapacity = 1020;
 
 template <RopeKernelMode Mode>
 inline constexpr bool kTextMode =
-    Mode == RopeKernelMode::Text1D || Mode == RopeKernelMode::TextMrope ||
-    Mode == RopeKernelMode::DflashText1D;
+    Mode == RopeKernelMode::Text1D || Mode == RopeKernelMode::Text1DYarn4 ||
+    Mode == RopeKernelMode::TextMrope || Mode == RopeKernelMode::DflashText1D ||
+    Mode == RopeKernelMode::DflashText1DYarn4;
 
 std::int64_t token_stride(const Tensor* tensor) {
     return tensor == nullptr ? 0 : tensor->nb[2] / static_cast<std::int64_t>(sizeof(__nv_bfloat16));
@@ -197,4 +199,59 @@ void rope_single_launch(const Tensor& positions, int rotary_dim, float theta, Te
     CUDA_CHECK(cudaGetLastError());
 }
 
+// Static YaRN factor-4 rope: the same fixed geometry as the native modes but
+// with the yarn4 frequency tables and the yarn4 attention scaling folded into
+// the sincos coefficients. Text is D256/R64 (24Q/4K or 16Q/2K), DFlash is
+// D128/R128 (32Q/8K).
+void rope_yarn4_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tensor& k,
+                       cudaStream_t stream) {
+    if (!bf16x2_aligned(q) || !bf16x2_aligned(k)) {
+        throw std::invalid_argument("rope_yarn4: q/k must be bf16x2 aligned");
+    }
+    const int axes = positions.ne[1];
+    if (rotary_dim == 128 && theta == 1.0e7F && axes == 1 && q.ne[0] == 128 && q.ne[1] == 32 &&
+        k.ne[1] == 8) {
+        launch_fixed<RopeKernelMode::DflashText1DYarn4, 32, 8>(positions, &q, &k, stream);
+    } else if (rotary_dim == 64 && theta == 1.0e7F) {
+        if (q.ne[1] == 24 && k.ne[1] == 4) {
+            launch_fixed<RopeKernelMode::Text1DYarn4, 24, 4>(positions, &q, &k, stream);
+        } else if (q.ne[1] == 16 && k.ne[1] == 2) {
+            launch_fixed<RopeKernelMode::Text1DYarn4, 16, 2>(positions, &q, &k, stream);
+        } else {
+            throw std::invalid_argument("rope_yarn4: unsupported Text head geometry");
+        }
+    } else {
+        throw std::invalid_argument("rope_yarn4: expected Text D256/R64 or DFlash D128/R128");
+    }
+    CUDA_CHECK(cudaGetLastError());
+}
+
+void rope_yarn4_single_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& x,
+                              cudaStream_t stream) {
+    if (!bf16x2_aligned(x)) {
+        throw std::invalid_argument("rope_yarn4: tensor must be bf16x2 aligned");
+    }
+    const int axes = positions.ne[1];
+    if (rotary_dim == 128 && theta == 1.0e7F && axes == 1 && x.ne[0] == 128) {
+        if (x.ne[1] == 32) {
+            launch_fixed_single<RopeKernelMode::DflashText1DYarn4, 32>(positions, x, stream);
+        } else if (x.ne[1] == 8) {
+            launch_fixed_single<RopeKernelMode::DflashText1DYarn4, 8>(positions, x, stream);
+        } else {
+            throw std::invalid_argument("rope_yarn4: unsupported DFlash head count");
+        }
+    } else if (rotary_dim == 64 && theta == 1.0e7F) {
+        if (x.ne[1] == 24 || x.ne[1] == 4) {
+            launch_fixed_single<RopeKernelMode::Text1DYarn4, 24>(positions, x, stream);
+        } else if (x.ne[1] == 16 || x.ne[1] == 2) {
+            launch_fixed_single<RopeKernelMode::Text1DYarn4, 16>(positions, x, stream);
+        } else {
+            throw std::invalid_argument("rope_yarn4: unsupported Text head count");
+        }
+    } else {
+        throw std::invalid_argument("rope_yarn4: expected Text D256/R64 or DFlash D128/R128");
+    }
+    CUDA_CHECK(cudaGetLastError());
+}
+
 } // namespace ninfer::ops::detail

src/ops/launcher/rope.h

@@ -15,4 +15,10 @@ void rope_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& q
 void rope_single_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& x,
                         cudaStream_t stream);
 
+void rope_yarn4_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tensor& k,
+                       cudaStream_t stream);
+
+void rope_yarn4_single_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& x,
+                              cudaStream_t stream);
+
 } // namespace ninfer::ops::detail

src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp

@@ -161,7 +161,7 @@ void validate_envelope(CausalAttentionExecutionEnvelope envelope, const PagedKVL
                        std::int32_t tokens, const char* op) {
     const std::uint32_t capacity = validate_cache(cache, cache.num_kv_heads, op);
     if (envelope.min_visible_keys == 0 || envelope.min_visible_keys > envelope.max_visible_keys ||
-        envelope.max_visible_keys > kCausalAttentionMaximumVisibleKeys ||
+        envelope.max_visible_keys > kCausalAttentionMaximumVisibleKeysYarn ||
         envelope.max_visible_keys > capacity) {
         throw std::invalid_argument(std::string(op) + ": invalid execution envelope");
     }
@@ -242,7 +242,7 @@ void validate_batched_attention_tensors(const Tensor& q, const Tensor& positions
     const std::uint32_t capacity = validate_batch_cache(cache, kv_heads, op);
     if (cache.block_tables.ne[1] < batch || envelope.min_visible_keys == 0 ||
         envelope.min_visible_keys > envelope.max_visible_keys ||
-        envelope.max_visible_keys > kCausalAttentionMaximumVisibleKeys ||
+        envelope.max_visible_keys > kCausalAttentionMaximumVisibleKeysYarn ||
         envelope.max_visible_keys > capacity ||
         envelope.max_visible_keys < static_cast<std::uint32_t>(width)) {
         throw std::invalid_argument(std::string(op) + ": invalid execution envelope or table");
@@ -361,7 +361,7 @@ std::size_t causal_softmax_attention_workspace_capacity_bytes(
     if (!supported_dtype || batch_size <= 0 || batch_size > kMaximumBatchSize || min_width <= 0 ||
         max_width < min_width || (batch_size > 1 && max_width > kMaximumVerifyTokens) ||
         envelope.min_visible_keys == 0 || envelope.min_visible_keys > envelope.max_visible_keys ||
-        envelope.max_visible_keys > kCausalAttentionMaximumVisibleKeys ||
+        envelope.max_visible_keys > kCausalAttentionMaximumVisibleKeysYarn ||
         envelope.max_visible_keys < static_cast<std::uint32_t>(max_width)) {
         throw std::invalid_argument(
             "causal_softmax_attention workspace: invalid profile or interval");

src/ops/wrapper/rope.cpp

@@ -140,4 +140,61 @@ void rope(const Tensor& positions, int rotary_dim, float theta, Tensor& x, cudaS
     detail::rope_single_launch(positions, rotary_dim, theta, x, stream);
 }
 
+void require_yarn4_domain(float theta, int rotary_dim, std::int32_t head_dim) {
+    if (theta != 1.0e7F) {
+        throw std::invalid_argument("rope_yarn4: theta must be 1e7");
+    }
+    const bool text_yarn4   = rotary_dim == 64 && head_dim == 256;
+    const bool dflash_yarn4 = rotary_dim == 128 && head_dim == 128;
+    if (!text_yarn4 && !dflash_yarn4) {
+        throw std::invalid_argument(
+            "rope_yarn4: expected Text D256/R64 or DFlash D128/R128");
+    }
+}
+
+void rope_yarn4(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tensor& k,
+                cudaStream_t stream) {
+    require_common(positions, rotary_dim, theta);
+    if (q.dtype != DType::BF16 || k.dtype != DType::BF16) {
+        throw std::invalid_argument("rope_yarn4: q/k must be BF16");
+    }
+    (void)numel_allow_zero(positions, "positions");
+    const std::int64_t q_numel = numel_allow_zero(q, "q");
+    (void)numel_allow_zero(k, "k");
+    const std::int32_t tokens   = q.ne[2];
+    const int axes              = position_axes(positions, tokens);
+    const std::int32_t head_dim = axes == 2 ? kVisionDim : q.ne[0];
+    const std::int32_t q_heads  = q.ne[1];
+    const std::int32_t k_heads  = k.ne[1];
+    if (axes != 1) { throw std::invalid_argument("rope_yarn4: requires 1-D positions"); }
+    require_yarn4_domain(theta, rotary_dim, head_dim);
+    require_tensor_layout(q, "q", head_dim, q_heads, tokens);
+    require_tensor_layout(k, "k", head_dim, k_heads, tokens);
+    if (q_numel == 0) { return; }
+    require_positions_storage(positions);
+    if (q.data == nullptr || k.data == nullptr) {
+        throw std::invalid_argument("rope_yarn4: q/k data must be non-null");
+    }
+    detail::rope_yarn4_launch(positions, rotary_dim, theta, q, k, stream);
+}
+
+void rope_yarn4(const Tensor& positions, int rotary_dim, float theta, Tensor& x,
+                cudaStream_t stream) {
+    require_common(positions, rotary_dim, theta);
+    if (x.dtype != DType::BF16) { throw std::invalid_argument("rope_yarn4: tensor must be BF16"); }
+    (void)numel_allow_zero(positions, "positions");
+    const std::int64_t x_numel  = numel_allow_zero(x, "tensor");
+    const std::int32_t tokens   = x.ne[2];
+    const int axes              = position_axes(positions, tokens);
+    const std::int32_t head_dim = axes == 2 ? kVisionDim : x.ne[0];
+    const std::int32_t heads    = x.ne[1];
+    if (axes != 1) { throw std::invalid_argument("rope_yarn4: requires 1-D positions"); }
+    require_yarn4_domain(theta, rotary_dim, head_dim);
+    require_tensor_layout(x, "tensor", head_dim, heads, tokens);
+    if (x_numel == 0) { return; }
+    require_positions_storage(positions);
+    if (x.data == nullptr) { throw std::invalid_argument("rope_yarn4: tensor data must be non-null"); }
+    detail::rope_yarn4_single_launch(positions, rotary_dim, theta, x, stream);
+}
+
 } // namespace ninfer::ops