From bc2b299e93679587a74b19f96c8261f62d3b89c4 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Sun, 30 Aug 2026 23:55:28 +0800 Subject: [PATCH 1/6] feat(rope): static YaRN factor-4 extension --- apps/cli/main.cpp | 1 + apps/cli/options.cpp | 2 + apps/cli/options.h | 1 + include/ninfer/ops/rope.h | 13 ++++ include/ninfer/types.h | 1 + src/core/device.h | 2 + src/ops/kernel/rope.cuh | 73 ++++++++++++++++++- src/ops/launcher/rope.cu | 60 ++++++++++++++- src/ops/launcher/rope.h | 6 ++ src/ops/wrapper/rope.cpp | 57 +++++++++++++++ src/runtime/engine/engine.cpp | 1 + src/serve/generation_service.cpp | 1 + src/serve/serve_options.cpp | 2 + src/serve/serve_options.h | 1 + .../qwen3_6/impl/runtime/text_context_impl.h | 20 ++++- 15 files changed, 234 insertions(+), 7 deletions(-) diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index 933192da9b..fabca93130 100644 --- a/apps/cli/main.cpp +++ b/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. diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index b5c5798d9a..dc1aec0ed7 100644 --- a/apps/cli/options.cpp +++ b/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::numeric_limits::max())) { diff --git a/apps/cli/options.h b/apps/cli/options.h index 3c0c2e0960..e3b1ceb9c4 100644 --- a/apps/cli/options.h +++ b/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; diff --git a/include/ninfer/ops/rope.h b/include/ninfer/ops/rope.h index d308991f1e..a4fcced15a 100644 --- a/include/ninfer/ops/rope.h +++ b/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 diff --git a/include/ninfer/types.h b/include/ninfer/types.h index 677c381e88..03786d76b8 100644 --- a/include/ninfer/types.h +++ b/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; diff --git a/src/core/device.h b/src/core/device.h index b4afa8384e..2ad513a710 100644 --- a/src/core/device.h +++ b/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(); diff --git a/src/ops/kernel/rope.cuh b/src/ops/kernel/rope.cuh index 73ed115d83..b3245bdcd2 100644 --- a/src/ops/kernel/rope.cuh +++ b/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(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(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(positions[token]) * + static_cast(kDflashRopeYarn4InvFrequency[pair]); + const double turns = angle * kInvTwoPi; + const float reduced = static_cast(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(positions[token]) * + static_cast(kTextRopeYarn4InvFrequency[pair]); + const double turns = angle * kInvTwoPi; + const float reduced = static_cast(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(blockIdx.x); if (token >= tokens) { return; } diff --git a/src/ops/launcher/rope.cu b/src/ops/launcher/rope.cu index 03ca1835a9..b9fcb87183 100644 --- a/src/ops/launcher/rope.cu +++ b/src/ops/launcher/rope.cu @@ -18,8 +18,9 @@ constexpr int kLargeBlockWaveCapacity = 1020; template 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(sizeof(__nv_bfloat16)); @@ -197,4 +198,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(positions, &q, &k, stream); + } else if (rotary_dim == 64 && theta == 1.0e7F) { + if (q.ne[1] == 24 && k.ne[1] == 4) { + launch_fixed(positions, &q, &k, stream); + } else if (q.ne[1] == 16 && k.ne[1] == 2) { + launch_fixed(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(positions, x, stream); + } else if (x.ne[1] == 8) { + launch_fixed_single(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(positions, x, stream); + } else if (x.ne[1] == 16 || x.ne[1] == 2) { + launch_fixed_single(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 diff --git a/src/ops/launcher/rope.h b/src/ops/launcher/rope.h index ba35837fc5..7eea862cf4 100644 --- a/src/ops/launcher/rope.h +++ b/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 diff --git a/src/ops/wrapper/rope.cpp b/src/ops/wrapper/rope.cpp index 17d574db88..3b52cbb8b0 100644 --- a/src/ops/wrapper/rope.cpp +++ b/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 diff --git a/src/runtime/engine/engine.cpp b/src/runtime/engine/engine.cpp index 0f7f848d66..6c72c4b6e1 100644 --- a/src/runtime/engine/engine.cpp +++ b/src/runtime/engine/engine.cpp @@ -215,6 +215,7 @@ class Engine::Impl { explicit Impl(EngineOptions engine_options) : options(normalize_engine_options(std::move(engine_options))), device(options.device) { + device.yarn_enabled = options.yarn_enabled; nvtx::ScopedRange load_range(nvtx::Name::EngineLoad, nvtx::Category::Runtime); auto constructed = targets::construct_target(options, device); active = std::move(constructed.active); diff --git a/src/serve/generation_service.cpp b/src/serve/generation_service.cpp index 08d0bf7227..a29a605d24 100644 --- a/src/serve/generation_service.cpp +++ b/src/serve/generation_service.cpp @@ -235,6 +235,7 @@ GenerationService::GenerationService(ServeOptions options, LoadProgress load_pro engine_options.prefill_chunk = options_.prefill_chunk; engine_options.kv_cache = options_.kv_cache; engine_options.enable_vision = options_.enable_vision; + engine_options.yarn_enabled = options_.yarn_enabled; engine_options.use_cuda_graph = options_.use_cuda_graph; engine_options.speculative = options_.speculative; engine_options.context_cache = options_.context_cache; diff --git a/src/serve/serve_options.cpp b/src/serve/serve_options.cpp index 66fef22197..105f690395 100644 --- a/src/serve/serve_options.cpp +++ b/src/serve/serve_options.cpp @@ -279,6 +279,8 @@ ServeOptions parse_serve_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 == "--no-prefix-reuse") { options.allow_prefix_reuse = false; } else if (arg == "--lm-head-draft") { diff --git a/src/serve/serve_options.h b/src/serve/serve_options.h index c529fbc84f..003fbb92c5 100644 --- a/src/serve/serve_options.h +++ b/src/serve/serve_options.h @@ -46,6 +46,7 @@ struct ServeOptions { ContextCacheOptions context_cache; bool enable_vision = false; bool use_cuda_graph = true; + bool yarn_enabled = false; bool allow_prefix_reuse = true; bool enable_thinking = true; // default thinking mode for the generation prompt (--no-thinking opts out) 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 f3c78dd9ae..617c105f3e 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_context_impl.h @@ -382,7 +382,11 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po 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); + if (ctx_.yarn_enabled) { + ops::rope_yarn4(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s); + } else { + ops::rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s); + } Tensor a = results.attention.view({kCfg.head_dim, kCfg.n_q, T}); if (active_sequence_batch_ != 0) { @@ -492,7 +496,11 @@ void TextContext::mtp_prefill_chunk(const Tensor& ids, const Tensor& hidden, Tensor v = v_flat.view({kCfg.head_dim, kCfg.n_kv, T}); Tensor kn = work_.alloc(DType::BF16, {kCfg.head_dim, kCfg.n_kv, T}); ops::rmsnorm(k, *mtp_.k_norm, kCfg.rms_eps, true, kn, s); + if (ctx_.yarn_enabled) { + ops::rope_yarn4(rope_positions, kCfg.rotary_dim, kCfg.rope_theta, kn, s); + } else { ops::rope(rope_positions, kCfg.rotary_dim, kCfg.rope_theta, kn, s); + } ops::kv_cache_append(kn, v, positions, mtp_kv_.layer_view(0), s); if (final_chunk) { @@ -532,7 +540,11 @@ void TextContext::mtp_prefill_chunk(const Tensor& ids, const Tensor& hidden, cudaMemcpyAsync(dst, src, sizeof(std::int32_t), cudaMemcpyDeviceToDevice, s)); } } + if (ctx_.yarn_enabled) { + ops::rope_yarn4(last_rope_position, kCfg.rotary_dim, kCfg.rope_theta, qn, s); + } else { ops::rope(last_rope_position, kCfg.rotary_dim, kCfg.rope_theta, qn, s); + } Tensor a = work_.alloc(DType::BF16, {kCfg.head_dim, kCfg.n_q, 1}); ops::causal_softmax_attention_cached(qn, last_position, @@ -845,7 +857,11 @@ void TextContext::attn_mix(const FullLayerW& w, Tensor& x, int fidx, Phase ph) { 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); + if (ctx_.yarn_enabled) { + ops::rope_yarn4(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s); + } else { + ops::rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s); + } Tensor a = results.attention.view({kCfg.head_dim, kCfg.n_q, T}); const Tensor& kv_table_rows = From e4e42caaa403ea47fa3db3f7fe3bb8d8d9b1ee86 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 07:06:21 +0800 Subject: [PATCH 2/6] feat(yarn): extend context budget and attention validation to 4x native --- include/ninfer/ops/softmax_attention.h | 2 ++ src/ops/launcher/rope.cu | 1 + .../dense/causal_cache/causal_softmax_attention.cpp | 6 +++--- src/targets/qwen3_6/impl/runtime/layouts_impl.h | 10 ++++++++-- 4 files changed, 14 insertions(+), 5 deletions(-) diff --git a/include/ninfer/ops/softmax_attention.h b/include/ninfer/ops/softmax_attention.h index 6fa169e0fa..1a7180f002 100644 --- a/include/ninfer/ops/softmax_attention.h +++ b/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; diff --git a/src/ops/launcher/rope.cu b/src/ops/launcher/rope.cu index b9fcb87183..9682fbe914 100644 --- a/src/ops/launcher/rope.cu +++ b/src/ops/launcher/rope.cu @@ -5,6 +5,7 @@ #include "ops/kernel/rope.cuh" #include +#include namespace ninfer::ops::detail { namespace { diff --git a/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp b/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp index 62abe13b63..b455d42784 100644 --- a/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp +++ b/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(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(max_width)) { throw std::invalid_argument( "causal_softmax_attention workspace: invalid profile or interval"); diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index b21165ac86..bdaea49c1a 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/src/targets/qwen3_6/impl/runtime/layouts_impl.h @@ -586,8 +586,14 @@ WorkspacePlan build_workspace_plan(const SequencePlanImpl& plan) { } void validate_target_options(DeviceContext& device, const EngineOptions& options) { - if (options.max_context == 0 || options.max_context > Variant::maximum_context) { - throw std::invalid_argument("max_context exceeds the variant native context capacity"); + // Static YaRN factor-4 extends the rope domain to 4x the native context. + const std::uint32_t context_limit = + options.yarn_enabled ? 4ULL * Variant::maximum_context : Variant::maximum_context; + if (options.max_context == 0 || options.max_context > context_limit) { + throw std::invalid_argument( + options.yarn_enabled + ? "max_context exceeds the variant YaRN-extended context capacity" + : "max_context exceeds the variant native context capacity"); } if (options.prefill_chunk == 0 || options.prefill_chunk % kPrefillChunkAlignment != 0) { throw std::invalid_argument("prefill_chunk must be a nonzero multiple of 128"); From 7dec7516062a50e35d4513fcf11e02d769750071 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 07:14:37 +0800 Subject: [PATCH 3/6] feat(yarn): allow explicit --kv-capacity below max_context --- src/targets/qwen3_6/impl/runtime/layouts.h | 3 +++ src/targets/qwen3_6/impl/runtime/layouts_impl.h | 11 ++++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/targets/qwen3_6/impl/runtime/layouts.h b/src/targets/qwen3_6/impl/runtime/layouts.h index 859f4f5d7e..5be33e7210 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts.h +++ b/src/targets/qwen3_6/impl/runtime/layouts.h @@ -70,6 +70,9 @@ struct WorkspacePlan { struct SequencePlanningInputs { WeightsProfile weights_profile; std::uint32_t capacity = 0; + // Explicit --kv-capacity page count (when set below capacity) shrinks the + // device page-pool floor: max_context then bounds only the rope domain. + std::optional kv_capacity_tokens; std::uint32_t max_concurrency = 1; std::uint32_t prefill_chunk = 0; std::uint32_t draft_window = 0; diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index bdaea49c1a..37600c559b 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/src/targets/qwen3_6/impl/runtime/layouts_impl.h @@ -742,6 +742,9 @@ make_sequence_planner_impl(DeviceContext& device, const EngineOptions& options, SequencePlanningInputs inputs{ .weights_profile = weights_profile, .capacity = options.max_context, + .kv_capacity_tokens = options.kv_capacity.mode == KvCapacityMode::Explicit + ? std::optional(options.kv_capacity.explicit_tokens) + : std::nullopt, .max_concurrency = options.max_concurrency, .prefill_chunk = std::min(options.prefill_chunk, options.max_context), .draft_window = options.speculative.draft_tokens, @@ -756,7 +759,13 @@ make_sequence_planner_impl(DeviceContext& device, const EngineOptions& options, .context_cache = options.context_cache, }; const std::uint32_t logical_pages = page_count(inputs.capacity); - const std::uint32_t minimum_pages = std::max(logical_pages, inputs.max_concurrency); + // The device page pool normally covers the full max_context; an explicit + // --kv-capacity below max_context instead floors the pool at that size so + // the rope domain (4x under YaRN) can exceed what the pool can hold. + std::uint32_t minimum_pages = std::max(logical_pages, inputs.max_concurrency); + if (inputs.kv_capacity_tokens) { + minimum_pages = std::max(page_count(*inputs.kv_capacity_tokens), inputs.max_concurrency); + } const std::uint64_t maximum_pages64 = static_cast(inputs.max_concurrency) * logical_pages; if (maximum_pages64 > std::numeric_limits::max()) { From 972ddb1b2d283220b3dd882bda357ec7aa681c7c Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 07:15:16 +0800 Subject: [PATCH 4/6] fix(yarn): relax --kv-capacity >= --max-context serve/CLI validation --- apps/cli/options.cpp | 4 ++-- src/serve/serve_options.cpp | 7 +++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index dc1aec0ed7..43bb50c669 100644 --- a/apps/cli/options.cpp +++ b/apps/cli/options.cpp @@ -209,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) { diff --git a/src/serve/serve_options.cpp b/src/serve/serve_options.cpp index 105f690395..9f98bde036 100644 --- a/src/serve/serve_options.cpp +++ b/src/serve/serve_options.cpp @@ -334,9 +334,12 @@ ServeOptions parse_serve_options(int argc, char** argv) { throw std::invalid_argument("--port must be in [1,65535]"); } if (options.max_context == 0) { throw std::invalid_argument("--max-context must be positive"); } + // An explicit --kv-capacity may floor the device page pool below max_context: + // YaRN extends the rope domain beyond the pool, and the cold pool recycles + // committed pages under pressure so the context can keep growing. 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"); } if (options.max_concurrency == 0 || options.max_concurrency > kMaximumConcurrency) { throw std::invalid_argument("--max-concurrency must be in [1,8]"); From 5d22054a1ddfaeaf1f2de8d71a587b25e06e4fe7 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 07:21:41 +0800 Subject: [PATCH 5/6] fix(yarn): relax engine kv_capacity floor to explicit pool size --- src/targets/qwen3_6/impl/runtime/layouts_impl.h | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index 37600c559b..da0034c616 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/src/targets/qwen3_6/impl/runtime/layouts_impl.h @@ -602,7 +602,13 @@ void validate_target_options(DeviceContext& device, const EngineOptions& options throw std::invalid_argument("max_concurrency must be in [1,8]"); } const std::uint32_t logical_pages = page_count(options.max_context); - const std::uint32_t minimum_pages = std::max(logical_pages, options.max_concurrency); + // An explicit --kv-capacity may floor the device page pool below max_context + // (YaRN extends the rope domain beyond the pool; the cold pool recycles pages). + std::uint32_t minimum_pages = std::max(logical_pages, options.max_concurrency); + if (options.kv_capacity.mode == KvCapacityMode::Explicit) { + minimum_pages = std::max(page_count(options.kv_capacity.explicit_tokens), + options.max_concurrency); + } const std::uint64_t maximum_pages64 = static_cast(options.max_concurrency) * logical_pages; if (maximum_pages64 > std::numeric_limits::max()) { @@ -610,8 +616,8 @@ void validate_target_options(DeviceContext& device, const EngineOptions& options } switch (options.kv_capacity.mode) { case KvCapacityMode::Explicit: { - if (options.kv_capacity.explicit_tokens < options.max_context) { - throw std::invalid_argument("kv_capacity must be at least max_context"); + if (options.kv_capacity.explicit_tokens == 0) { + throw std::invalid_argument("kv_capacity must be positive"); } const std::uint32_t requested_pages = page_count(options.kv_capacity.explicit_tokens); if (requested_pages < minimum_pages || requested_pages > maximum_pages64) { From b9ac21f3065c7c07d0c44b465fedb8836d872a82 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 07:29:39 +0800 Subject: [PATCH 6/6] fix(yarn): allow physical page pool below logical max_context --- src/targets/qwen3_6/impl/state/decoder_state.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/targets/qwen3_6/impl/state/decoder_state.cpp b/src/targets/qwen3_6/impl/state/decoder_state.cpp index 5e7372df45..1e9910496a 100644 --- a/src/targets/qwen3_6/impl/state/decoder_state.cpp +++ b/src/targets/qwen3_6/impl/state/decoder_state.cpp @@ -32,8 +32,11 @@ PagedKVCacheLayout plan_cache(LayoutBuilder& builder, std::uint32_t layers, std: } const std::uint32_t logical_pages = page_count(capacity); - if (physical_page_groups < logical_pages) { - throw std::invalid_argument("Paged KV physical pages are below logical capacity"); + // An explicit --kv-capacity may floor the device page pool below max_context: + // the rope domain (4x under YaRN) can exceed the pool, and the cold pool + // recycles committed pages under pressure so the context keeps growing. + if (physical_page_groups == 0) { + throw std::invalid_argument("Paged KV physical page capacity is zero"); } KVPageGeometry geometry;