From 8bf12349aafaa197cfd665f6c89b20f4da24140c Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Sun, 30 Aug 2026 23:39:13 +0800 Subject: [PATCH 1/4] feat(spec): length-based backend decision policy + DFlash2 enum/CLI --- include/ninfer/types.h | 1 + src/product/speculative_options.h | 10 ++- .../qwen3_6/impl/runtime/spec_decision.h | 81 +++++++++++++++++++ 3 files changed, 91 insertions(+), 1 deletion(-) create mode 100644 src/targets/qwen3_6/impl/runtime/spec_decision.h diff --git a/include/ninfer/types.h b/include/ninfer/types.h index 677c381e88..44726959b5 100644 --- a/include/ninfer/types.h +++ b/include/ninfer/types.h @@ -69,6 +69,7 @@ enum class SpeculativeBackend : std::uint8_t { None, Mtp, DFlash, + DFlash2, }; struct SpeculativeOptions { diff --git a/src/product/speculative_options.h b/src/product/speculative_options.h index a50ea5fc82..aa307d21b6 100644 --- a/src/product/speculative_options.h +++ b/src/product/speculative_options.h @@ -11,6 +11,7 @@ namespace ninfer::product { [[nodiscard]] inline SpeculativeBackend parse_speculative_backend(std::string_view value) { if (value == "mtp") { return SpeculativeBackend::Mtp; } if (value == "dflash") { return SpeculativeBackend::DFlash; } + if (value == "dflash2") { return SpeculativeBackend::DFlash2; } throw std::invalid_argument("invalid speculative backend: " + std::string(value)); } @@ -22,6 +23,8 @@ namespace ninfer::product { return "mtp"; case SpeculativeBackend::DFlash: return "dflash"; + case SpeculativeBackend::DFlash2: + return "dflash2"; } return "unknown"; } @@ -31,7 +34,7 @@ inline void validate_speculative_cli_options(const SpeculativeOptions& options) case SpeculativeBackend::None: if (options.draft_tokens != 0 || options.proposal_head != ProposalHead::Full) { throw std::invalid_argument( - "--draft-tokens and --lm-head-draft require --spec mtp|dflash"); + "--draft-tokens and --lm-head-draft require --spec mtp|dflash|dflash2"); } return; case SpeculativeBackend::Mtp: @@ -44,6 +47,11 @@ inline void validate_speculative_cli_options(const SpeculativeOptions& options) throw std::invalid_argument("--spec dflash requires --draft-tokens in [1,15]"); } return; + case SpeculativeBackend::DFlash2: + if (options.draft_tokens == 0 || options.draft_tokens > 15) { + throw std::invalid_argument("--spec dflash2 requires --draft-tokens in [1,15]"); + } + return; } throw std::invalid_argument("invalid speculative backend"); } diff --git a/src/targets/qwen3_6/impl/runtime/spec_decision.h b/src/targets/qwen3_6/impl/runtime/spec_decision.h new file mode 100644 index 0000000000..94b6f7a5cd --- /dev/null +++ b/src/targets/qwen3_6/impl/runtime/spec_decision.h @@ -0,0 +1,81 @@ +#pragma once + +// Per-request speculative backend decision policy (host-only, no CUDA deps). +// +// Measured on the 32 GiB RTX 5090 (Qwen3.8-27B NVFP4 artifact, 256-token +// greedy decodes, Chinese wiki prompts; /tmp/spec_threshold_results.txt): +// prompt tok | MTP d3 | DFlash2 d7 +// 2,047 | 58.4 | 75.3 (+29%) +// 4,066 | 58.3 | 71.2 (+22%) +// 7,924 | 58.8 | 71.9 (+22%) +// 11,585 | 58.1 | 60.4 (+4%) +// 15,226 | 57.6 | 59.1 (+2.6%) +// MTP is flat to 21k+; the curves cross near 17-20k. DSpark is dominated at +// every point (65.8 at 2k, 35.3 at 15k) and excluded from selection. +// +// Policy (user decision 2026-08-30): +// * admission: projected context <= kSpecDemoteTokens -> DFlash2, else MTP +// * mid-flight: while DFlash2 is active, demote the engine to MTP once ALL +// active requests' frontiers pass the threshold, or on VRAM pressure. +// Demotion is low-frequency (at most once per growth crossing) and uses +// the existing spec-degrade resume machinery. +// * promotion back to DFlash2 happens only when the engine is fully idle. + +#include +#include +#include + +namespace ninfer::targets::qwen3_6 { + +inline constexpr std::uint32_t kSpecDemoteTokens = 20480; + +enum class SpecDecision : std::uint8_t { + Mtp, + DFlash2, +}; + +struct SpecAdmissionInputs { + std::uint32_t prompt_tokens = 0; + // Generation budget if the caller supplied one; 0 = unknown. When known, + // a projection that is guaranteed to cross the threshold starts on MTP + // directly instead of demoting mid-flight. + std::uint32_t max_new_tokens = 0; + bool dflash2_available = false; + std::size_t free_device_bytes = 0; + // Projected device bytes the DFlash2 context caches need for this request + // (BF16 full-context draft caches). 0 = unknown / not tracked yet. + std::size_t dflash2_context_bytes = 0; +}; + +[[nodiscard]] inline SpecDecision choose_spec_backend(const SpecAdmissionInputs& in) noexcept { + if (!in.dflash2_available) { return SpecDecision::Mtp; } + const std::uint64_t projection = static_cast(in.prompt_tokens) + + (in.max_new_tokens != 0 ? in.max_new_tokens : 0); + if (projection > kSpecDemoteTokens) { return SpecDecision::Mtp; } + if (in.dflash2_context_bytes != 0 && in.free_device_bytes != 0 && + in.dflash2_context_bytes > in.free_device_bytes) { + return SpecDecision::Mtp; + } + return SpecDecision::DFlash2; +} + +// Engine-wide mid-flight demotion check (called at decode round boundaries; +// cheap by design). Demote when every active request has crossed the +// threshold, or when the DFlash2 footprint no longer fits free VRAM. +[[nodiscard]] inline bool should_demote_dflash2( + const std::vector& active_frontiers, std::size_t free_device_bytes, + std::size_t dflash2_footprint_bytes) noexcept { + if (active_frontiers.empty()) { return false; } + bool all_beyond = true; + for (const std::uint32_t frontier : active_frontiers) { + if (frontier <= kSpecDemoteTokens) { + all_beyond = false; + break; + } + } + if (all_beyond) { return true; } + return dflash2_footprint_bytes != 0 && free_device_bytes != 0 && + dflash2_footprint_bytes > free_device_bytes; +} + +} // namespace ninfer::targets::qwen3_6 From 3ac3b04e827d0925e77ce24f92258395c73e8e65 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 08:04:14 +0800 Subject: [PATCH 2/4] feat(spec): port DFlash2 draft backend + length-based switching --- $f | 0 .../ninfer/ops/bidirectional_gqa_attention.h | 68 ++ include/ninfer/ops/dflash2_grouped_conv.h | 38 + include/ninfer/ops/dflash2_selector.h | 47 ++ include/ninfer/ops/kv_cache_append.h | 14 +- include/ninfer/ops/swa.h | 64 ++ include/ninfer/types.h | 1 + src/CMakeLists.txt | 8 + .../kernel/bidirectional_gqa_attention.cuh | 719 ++++++++++++++++++ src/ops/kernel/dflash2_grouped_conv.cuh | 44 ++ src/ops/kernel/dflash2_selector.cuh | 191 +++++ src/ops/kv_cache/append/kernel.cuh | 4 +- src/ops/kv_cache/append/kv_cache_append.cpp | 15 +- src/ops/kv_cache/append/launch.cu | 28 +- src/ops/kv_cache/append/launch.h | 3 +- .../launcher/bidirectional_gqa_attention.cu | 185 +++++ .../launcher/bidirectional_gqa_attention.h | 33 + src/ops/launcher/dflash2_grouped_conv.cu | 29 + src/ops/launcher/dflash2_grouped_conv.h | 14 + src/ops/launcher/dflash2_selector.cu | 48 ++ src/ops/launcher/dflash2_selector.h | 20 + src/ops/launcher/swa.cu | 166 ++++ src/ops/launcher/swa.h | 33 + .../wrapper/bidirectional_gqa_attention.cpp | 168 ++++ src/ops/wrapper/dflash2_grouped_conv.cpp | 45 ++ src/ops/wrapper/dflash2_selector.cpp | 57 ++ src/ops/wrapper/swa.cpp | 144 ++++ src/product/speculative_options.h | 16 +- .../ninfer/targets/qwen3_6/model_view.h | 32 +- .../ninfer/targets/qwen3_6/round_state.h | 2 + .../ninfer/targets/qwen3_6/startup_features.h | 6 + .../qwen3_6/impl/runtime/dflash2_impl.h | 475 ++++++++++++ .../qwen3_6/impl/runtime/dflash_context.h | 14 + .../impl/runtime/dflash_context_impl.h | 32 + .../qwen3_6/impl/runtime/dflash_impl.h | 2 +- src/targets/qwen3_6/impl/runtime/instance.h | 8 + .../qwen3_6/impl/runtime/instantiate.h | 1 + src/targets/qwen3_6/impl/runtime/layouts.h | 15 + .../qwen3_6/impl/runtime/layouts_impl.h | 129 +++- src/targets/qwen3_6/impl/runtime/program.h | 9 + .../qwen3_6/impl/runtime/program_impl.h | 312 +++++++- .../qwen3_6/impl/runtime/request_plan_impl.h | 9 + src/targets/qwen3_6/impl/runtime/schedule.h | 44 ++ .../qwen3_6/impl/runtime/text_prefill_impl.h | 17 + .../qwen3_6/impl/runtime/workspace_recipe.h | 16 + .../qwen3_6/impl/state/round_state.cpp | 2 + .../ninfer/targets/qwen3_6_27b/package.h | 1 + src/targets/qwen3_6_27b/impl/config.h | 242 +++--- .../qwen3_6_27b/impl/load/bindings.cpp | 108 +++ src/targets/qwen3_6_27b/impl/load/bindings.h | 33 +- src/targets/qwen3_6_27b/impl/package.cpp | 35 +- src/targets/qwen3_6_27b/impl/variant.cpp | 52 ++ src/targets/qwen3_6_27b/impl/variant.h | 9 + src/targets/qwen3_6_35b_a3b/impl/config.h | 34 + .../qwen3_6_35b_a3b/impl/load/bindings.h | 2 +- src/targets/qwen3_6_35b_a3b/impl/variant.cpp | 4 + src/targets/qwen3_6_35b_a3b/impl/variant.h | 5 + 57 files changed, 3722 insertions(+), 130 deletions(-) create mode 100644 $f create mode 100644 include/ninfer/ops/bidirectional_gqa_attention.h create mode 100644 include/ninfer/ops/dflash2_grouped_conv.h create mode 100644 include/ninfer/ops/dflash2_selector.h create mode 100644 include/ninfer/ops/swa.h create mode 100644 src/ops/kernel/bidirectional_gqa_attention.cuh create mode 100644 src/ops/kernel/dflash2_grouped_conv.cuh create mode 100644 src/ops/kernel/dflash2_selector.cuh create mode 100644 src/ops/launcher/bidirectional_gqa_attention.cu create mode 100644 src/ops/launcher/bidirectional_gqa_attention.h create mode 100644 src/ops/launcher/dflash2_grouped_conv.cu create mode 100644 src/ops/launcher/dflash2_grouped_conv.h create mode 100644 src/ops/launcher/dflash2_selector.cu create mode 100644 src/ops/launcher/dflash2_selector.h create mode 100644 src/ops/launcher/swa.cu create mode 100644 src/ops/launcher/swa.h create mode 100644 src/ops/wrapper/bidirectional_gqa_attention.cpp create mode 100644 src/ops/wrapper/dflash2_grouped_conv.cpp create mode 100644 src/ops/wrapper/dflash2_selector.cpp create mode 100644 src/ops/wrapper/swa.cpp create mode 100644 src/targets/qwen3_6/impl/runtime/dflash2_impl.h diff --git a/$f b/$f new file mode 100644 index 0000000000..e69de29bb2 diff --git a/include/ninfer/ops/bidirectional_gqa_attention.h b/include/ninfer/ops/bidirectional_gqa_attention.h new file mode 100644 index 0000000000..ed1570f65c --- /dev/null +++ b/include/ninfer/ops/bidirectional_gqa_attention.h @@ -0,0 +1,68 @@ +#pragma once + +#include "core/arena.h" +#include "core/paged_kv_cache.h" +#include "core/tensor.h" + +#include + +#include +#include + +namespace ninfer::ops { + +/** + * Host execution-resource promise for bidirectional_gqa_attention. + * + * Device context_lengths define the exact per-row mathematical contexts. This envelope bounds + * every row so a fixed launch can be captured and replayed without a host read. + */ +struct GqaContextExecutionEnvelope { + std::uint32_t min_context = 0; + std::uint32_t max_context = 0; +}; + +/** + * Op: bidirectional grouped-query attention over persistent context and one query block + * + * For D=128, Hq=32, Hkv=8, group=4, query row i, query head h, and kvh=floor(h/4): + * + * keys = context K rows [0,L) followed logically by every live query K row [0,V) + * score = scale * dot(q[:,h,i], key[:,kvh,j]) + * prob = softmax over the complete logical key set + * ideal[:,h,i] = sum_j prob[j] * value[:,kvh,j] + * + * q/out are contiguous BF16 [128,32,W,B]. query_k/query_v are contiguous BF16 [128,8,W,B]. + * context_lengths, valid_columns, and table_rows are contiguous device I32 [B]. Row b has + * V=valid_columns[b] live query columns and reads logical context [0,context_lengths[b]) through + * table row table_rows[b]. Columns i>=V are an inert physical tail and produce zero output. + * context is a read-only paged BF16 cache with head-major page planes [128,64,Nphysical,8]. scale + * is 1/sqrt(128). + * + * There is no causal triangle: every live query row attends every other live query K/V row in the + * same batch row. Context and query K/V remain separate physical segments and every input/cache + * byte is unchanged. The oracle evaluates `ideal` naively in FP64 from represented inputs. The + * BF16 out is promoted and compared directly with that result; output storage rounding belongs to + * the Op's numerical criterion, not the oracle. out is the only observable mutation and is + * completely overwritten. The current optimized implementation domain is W=1..16 on sm_120a. + * + * The caller guarantees min_context <= L <= max_context and that every logical page intersecting + * [0,L) is materialized. The execution envelope may affect finite launch selection and workspace + * capacity, never the admitted key set or numerical result. + */ +void bidirectional_gqa_attention(const Tensor& q, const Tensor& query_k, const Tensor& query_v, + const Tensor& context_lengths, const Tensor& valid_columns, + const Tensor& table_rows, float scale, + const PagedKVBatchLayerView& context, + GqaContextExecutionEnvelope envelope, WorkspaceArena& workspace, + Tensor& out, cudaStream_t stream); + +/** + * Returns the transient arena capacity required for every T in the inclusive optimized interval. + * The execution envelope is the fixed profile; invalid profiles or intervals throw. + */ +[[nodiscard]] std::size_t bidirectional_gqa_attention_workspace_capacity_bytes( + GqaContextExecutionEnvelope envelope, std::int32_t min_tokens, std::int32_t max_tokens, + std::int32_t batch_size); + +} // namespace ninfer::ops diff --git a/include/ninfer/ops/dflash2_grouped_conv.h b/include/ninfer/ops/dflash2_grouped_conv.h new file mode 100644 index 0000000000..6ddfb8b61f --- /dev/null +++ b/include/ninfer/ops/dflash2_grouped_conv.h @@ -0,0 +1,38 @@ +#pragma once + +#include "core/tensor.h" + +#include + +#include + +namespace ninfer::ops { + +/** + * Op: DFlash2 grouped convolution pass + * + * `hidden` is contiguous BF16 [H,T]: the token-major block produced by the + * draft projections. `delta` is contiguous BF16 [N,T] with N=2*taps*G and is + * the exact token-major output of the kernel projection (N fastest). `side` + * selects the pass's tap bank. `base` is a BF16 [taps,H] Weight (one side of + * the stored [2,taps,H] base kernel). `out` is contiguous BF16 [H,T] and is + * completely overwritten. + * + * For token t with in-block position p = t & (block_size-1) and channel h in + * group g = h / group_size, + * + * out[h,t] = hidden[h,t] * (base[0,h] + delta[G*(side*taps)+g,t]) + * + sum_{tap>=1, p>=tap} hidden[h,t-tap] + * * (base[tap,h] + delta[G*(side*taps+tap)+g,t]). + * + * Cross-block reads introduced by t-tap are multiplied by the zero in-block + * position mask, so rows never leak into a neighbour block. The registered + * domain is H=5120, T=1..64, taps=2, G=320, group_size=16, block_size=8, + * side in {0,1}. Inputs and base are unchanged and the Op owns no workspace or + * persistent state. Intermediate arithmetic is FP32 and out is rounded to BF16. + */ +void dflash2_grouped_conv(const Tensor& hidden, const Tensor& delta, const Weight& base, + std::int32_t block_size, std::int32_t group_size, std::int32_t taps, + std::int32_t side, Tensor& out, cudaStream_t stream); + +} // namespace ninfer::ops diff --git a/include/ninfer/ops/dflash2_selector.h b/include/ninfer/ops/dflash2_selector.h new file mode 100644 index 0000000000..beffed7cd1 --- /dev/null +++ b/include/ninfer/ops/dflash2_selector.h @@ -0,0 +1,47 @@ +#pragma once + +#include "core/tensor.h" + +#include + +#include + +namespace ninfer::ops { + +/** + * Op: DFlash2 candidate selector with a greedy path walk + * + * `unary_logits` is contiguous BF16 [V,S*B]: full-vocab base logits of S + * proposal positions for B batch rows, token-major (each column holds one + * token's V logits). `projected_hidden` is contiguous BF16 [R,S*B]: the + * candidate selector's hidden projection of the same proposal columns. + * `predecessor_codebook` and `successor_codebook` are BF16 [V,R] Weights. + * `anchors` is contiguous I32 [B] and holds the previous target token per row. + * + * `candidates` (I32), `unary` (F32), and `scores` (F32) are caller-owned + * scratch with shapes [B,S,K], [B,S,K], and [B,S,K,K] respectively; all are + * completely overwritten. `drafts` is contiguous I32 [S*B] and receives the + * selected token for every proposal position. + * + * Selection: per row/step the top-K base logits (higher value first, lower + * token id breaking ties) form the candidate set. Each candidate c at step s + * receives the pair score + * + * score(s,p,c) = unary[s,c] + * + sum_r projected[r,s] * predecessor_codebook[pred,r] + * * successor_codebook[candidate[s,c],r], + * + * where pred is the anchor token at s=0 and candidate[s-1,p] afterwards; only + * p=0 is admitted at s=0. The walk starts at previous candidate index 0 and + * greedily picks the highest-score current candidate at every step (lowest + * candidate index breaking ties). + * + * The registered domain is V=248320, R=256, S=7, B=1..8, K=16. Inputs and + * weights are unchanged and the Op owns no workspace or persistent state. + */ +void dflash2_selector(const Tensor& unary_logits, const Tensor& projected_hidden, + const Weight& predecessor_codebook, const Weight& successor_codebook, + const Tensor& anchors, Tensor& candidates, Tensor& unary, Tensor& scores, + Tensor& drafts, std::int32_t steps, std::int32_t top_k, cudaStream_t stream); + +} // namespace ninfer::ops diff --git a/include/ninfer/ops/kv_cache_append.h b/include/ninfer/ops/kv_cache_append.h index ddddee74bd..61ee3fd1a0 100644 --- a/include/ninfer/ops/kv_cache_append.h +++ b/include/ninfer/ops/kv_cache_append.h @@ -79,15 +79,17 @@ void kv_cache_append_prefix(const Tensor& k, const Tensor& v, const Tensor& posi * Append device-selected exact BF16 prefixes to lane-owned cyclic storage. * * k/v, positions, counts, and their exact-copy and mutation contracts match the paged overload; - * lanes[b] selects the destination lane. The fixed geometry is D=128, Hkv=8, capacity=4096, and - * absolute position p maps to slot p mod 4096. The caller guarantees that each row's existing live - * interval ends immediately before positions[0,b], advancing it by counts[b] makes every - * overwritten old slot dead, and one row commits at most the ring capacity. Consequently, no two - * live writes race for one physical slot. The Op does not own or publish the lane frontier. + * lanes[b] selects the destination lane. The registered geometry is D=128, Hkv=8, capacity equal + * to the window in {2048, 4096}, and absolute position p maps to slot p mod window. The caller + * guarantees that each row's existing live interval ends immediately before positions[0,b], + * advancing it by counts[b] makes every overwritten old slot dead, and one row commits at most the + * ring capacity. Consequently, no two live writes race for one physical slot. The Op does not own + * or publish the lane frontier. */ void kv_cache_append_prefix(const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& counts, const Tensor& lanes, KVCacheAppendPrefixExecutionEnvelope envelope, - CyclicKVCacheLayerView cache, cudaStream_t stream); + CyclicKVCacheLayerView cache, std::uint32_t window, + cudaStream_t stream); } // namespace ninfer::ops diff --git a/include/ninfer/ops/swa.h b/include/ninfer/ops/swa.h new file mode 100644 index 0000000000..ab1ee5b881 --- /dev/null +++ b/include/ninfer/ops/swa.h @@ -0,0 +1,64 @@ +#pragma once + +#include "core/arena.h" +#include "core/cyclic_kv_cache.h" +#include "core/tensor.h" + +#include + +#include +#include + +namespace ninfer::ops { + +/** + * Host execution-resource promise for swa. + * + * positions[0,b] is row b's exact device-resident committed-context frontier. This envelope bounds + * every row so a fixed launch can be captured and replayed without a host read. + */ +struct SwaContextExecutionEnvelope { + std::uint32_t min_context = 0; + std::uint32_t max_context = 0; +}; + +/** + * Op: symmetric non-causal sliding-window grouped-query attention + * + * The fixed optimized geometry is D=128, Hq=32, Hkv=8, group=4, and a registered + * window W in {2048, 4096}. q/out are contiguous BF16 [128,32,W,B], query_k/query_v are + * contiguous BF16 [128,8,W,B], positions is contiguous device I32 [W,B], valid_columns and + * lanes are contiguous device I32 [B]. Row b has V=valid_columns[b] live query columns with + * positions[i,b]=L[b]+i for i=V are an + * inert physical tail and produce zero output. + * + * The read-only cyclic context contains committed absolute positions + * [max(0,L-window),L), with absolute position p stored at physical slot p mod padded_capacity + * (the registered window). Query K/V is a separate temporary segment at positions [L,L+V). For + * every live query position p_i, admitted populated keys satisfy abs(p_j-p_i) +#include + +#include + +namespace ninfer::ops { + +inline constexpr int kBidirectionalGqaHeadDim = 128; +inline constexpr int kBidirectionalGqaQHeads = 32; +inline constexpr int kBidirectionalGqaKVHeads = 8; +inline constexpr int kBidirectionalGqaGroup = 4; +inline constexpr int kBidirectionalGqaMaxSplit = 85; + +__device__ __forceinline__ int bidirectional_gqa_swz(int row, int col) { + return (((col >> 3) ^ (row & 7)) << 3) | (col & 7); +} + +__device__ __forceinline__ unsigned bidirectional_gqa_swz_addr(unsigned lane_base, unsigned ck, + unsigned as, unsigned r) { + return lane_base + ((ck | as) ^ r); +} + +__device__ __forceinline__ std::int64_t bidirectional_gqa_q_index(int q_head, int d, int token) { + return static_cast(d) + + static_cast(kBidirectionalGqaHeadDim) * + (static_cast(q_head) + + static_cast(kBidirectionalGqaQHeads) * token); +} + +__device__ __forceinline__ std::int64_t bidirectional_gqa_query_kv_index(int kv_head, int d, + int token) { + return static_cast(d) + + static_cast(kBidirectionalGqaHeadDim) * + (static_cast(kv_head) + + static_cast(kBidirectionalGqaKVHeads) * token); +} + +__device__ __forceinline__ std::int64_t +bidirectional_gqa_cyclic_context_index(int kv_head, int d, int position, int padded_context) { + return static_cast(d) + static_cast(kBidirectionalGqaHeadDim) * + (static_cast(position) + + static_cast(padded_context) * kv_head); +} + +template +__device__ __forceinline__ std::int64_t bidirectional_gqa_partial_index(int q_head, int d, + int token, int split) { + return static_cast(d) + + static_cast(kBidirectionalGqaHeadDim) * + (static_cast(q_head) + + static_cast(kBidirectionalGqaQHeads) * + (static_cast(token) + static_cast(Tokens) * split)); +} + +template +__device__ __forceinline__ std::int64_t bidirectional_gqa_stat_index(int q_head, int token, + int split) { + return static_cast(q_head) + + static_cast(kBidirectionalGqaQHeads) * + (static_cast(token) + static_cast(Tokens) * split); +} + +__device__ __forceinline__ void noncausal_gqa_row_to_qt(int row, int kv_head, int& q_head, + int& token) { + token = row / kBidirectionalGqaGroup; + const int q_local = row - token * kBidirectionalGqaGroup; + q_head = kv_head * kBidirectionalGqaGroup + q_local; +} + +template +__device__ __forceinline__ void +bidirectional_gqa_stage_tile(__nv_bfloat16* dst, const __nv_bfloat16* context, + const __nv_bfloat16* query, int key0, int valid_keys, bool query_tile, + int kv_head, int context_stride, int physical_page, int tid) { + constexpr int VecsPerRow = kBidirectionalGqaHeadDim / 8; + constexpr int Page = 64; + const std::int64_t paged_base = + static_cast(kBidirectionalGqaHeadDim) * + ((key0 & (Page - 1)) + Page * (physical_page + context_stride * kv_head)); + for (int chunk = tid; chunk < KeyBlock * VecsPerRow; chunk += Threads) { + const int row = chunk / VecsPerRow; + const int d = (chunk - row * VecsPerRow) * 8; + const bool live = row < valid_keys; + const int safe_row = live ? row : 0; + std::int64_t src_index; + if constexpr (CyclicSwa) { + const int context_position = (live ? key0 + row : 0) & (Window - 1); + src_index = query_tile ? bidirectional_gqa_query_kv_index(kv_head, d, safe_row) + : bidirectional_gqa_cyclic_context_index( + kv_head, d, context_position, context_stride); + } else { + src_index = query_tile + ? bidirectional_gqa_query_kv_index(kv_head, d, safe_row) + : paged_base + d + + static_cast(kBidirectionalGqaHeadDim) * safe_row; + } + const __nv_bfloat16* src = query_tile ? query + src_index : context + src_index; + __nv_bfloat16* smem = &dst[row * kBidirectionalGqaHeadDim + bidirectional_gqa_swz(row, d)]; + cp_async_zfill<16, Cache::cg>(smem, src, live ? 16 : 0); + } +} + +template +__device__ __forceinline__ void noncausal_gqa_split_partial_body( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ query_k, + const __nv_bfloat16* __restrict__ query_v, const std::int32_t* __restrict__ context_state, + const std::int32_t* __restrict__ valid_columns, const std::int32_t* __restrict__ selectors, + const __nv_bfloat16* __restrict__ context_k, const __nv_bfloat16* __restrict__ context_v, + const std::int32_t* __restrict__ block_tables, int context_stride, int logical_pages, + int max_context, int split_capacity, float scale, __nv_bfloat16* __restrict__ partial_acc, + float* __restrict__ partial_m, float* __restrict__ partial_l, __nv_bfloat16* __restrict__ out) { + static_assert(Tokens >= 1 && Tokens <= 16); + static_assert(WarpsPerCta == (Tokens + 3) / 4); + static_assert(KeyBlock == 32 || KeyBlock == 64); + + constexpr int D = kBidirectionalGqaHeadDim; + constexpr int Wc = WarpsPerCta; + constexpr int Threads = Wc * 32; + constexpr int Br = Wc * 16; + constexpr int RowCount = Tokens * kBidirectionalGqaGroup; + constexpr int QKNt = KeyBlock / 8; + constexpr int QKKs = D / 16; + constexpr int PVNt = D / 8; + constexpr int PVKs = KeyBlock / 16; + constexpr int RowBytes = D * static_cast(sizeof(__nv_bfloat16)); + constexpr float Log2E = 1.4426950408889634074f; + constexpr unsigned FullMask = 0xffffffffu; + + static_assert(RowCount <= Br); + static_assert(Br <= 2 * KeyBlock); + const int kv_head = static_cast(blockIdx.x); + const int split = static_cast(blockIdx.y); + const int batch = static_cast(blockIdx.z); + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + const int lane = tid & 31; + + constexpr std::int64_t QueryElements = + static_cast(D) * kBidirectionalGqaQHeads * Tokens; + constexpr std::int64_t QueryKvElements = + static_cast(D) * kBidirectionalGqaKVHeads * Tokens; + constexpr std::int64_t PartialElements = QueryElements; + constexpr std::int64_t StatElements = + static_cast(kBidirectionalGqaQHeads) * Tokens; + q += QueryElements * batch; + query_k += QueryKvElements * batch; + query_v += QueryKvElements * batch; + out += QueryElements * batch; + partial_acc += PartialElements * split_capacity * batch; + partial_m += StatElements * split_capacity * batch; + partial_l += StatElements * split_capacity * batch; + const int valid = valid_columns[batch]; + if constexpr (CyclicSwa) { + context_state += static_cast(Tokens) * batch; + const std::int64_t lane_elements = + static_cast(D) * context_stride * kBidirectionalGqaKVHeads; + context_k += lane_elements * selectors[batch]; + context_v += lane_elements * selectors[batch]; + } else { + context_state += batch; + block_tables += static_cast(logical_pages) * selectors[batch]; + } + const int length = context_state[0]; + if (kv_head >= kBidirectionalGqaKVHeads || split >= split_capacity || length < 0 || + length > max_context || valid < 1 || valid > Tokens) { + return; + } + + const int context_count = CyclicSwa ? min(length, Window - 1) : length; + const int context_start = length - context_count; + const int context_tiles = (context_count + KeyBlock - 1) / KeyBlock; + const int active_splits = context_tiles > 0 ? min(context_tiles, split_capacity) : 1; + if (split >= active_splits) { return; } + + const int tile_begin = + static_cast((static_cast(context_tiles) * split) / active_splits); + const int tile_end = + static_cast((static_cast(context_tiles) * (split + 1)) / active_splits); + const bool owns_query = split == active_splits - 1; + const int context_tile_count = tile_end - tile_begin; + const int iterations = context_tile_count + (owns_query ? 1 : 0); + + int table_group = -1; + int table_lane_page = 0; + const auto paged_page = [&](int key0) { + const int logical_page = key0 >> 6; + if constexpr (Tokens == 4) { + const int group = logical_page & ~3; + if (group != table_group) { + const int table_index = group + lane; + if (lane < 4 && table_index < logical_pages) { + table_lane_page = __ldg(block_tables + table_index); + } + table_group = group; + } + return __shfl_sync(FullMask, table_lane_page, logical_page - group); + } else { + const int physical_page = lane == 0 ? __ldg(block_tables + logical_page) : 0; + return __shfl_sync(FullMask, physical_page, 0); + } + }; + if constexpr (!CyclicSwa && Tokens == 4) { + if (context_tile_count > 0) { + const int first_logical_page = + (context_start + static_cast(tile_begin) * KeyBlock) >> 6; + const int first_group = first_logical_page & ~3; + const int table_index = first_group + lane; + if (lane < 4 && table_index < logical_pages) { + table_lane_page = __ldg(block_tables + table_index); + } + table_group = first_group; + } + } + + extern __shared__ __align__(16) __nv_bfloat16 shared[]; + __nv_bfloat16* k_s = shared; + __nv_bfloat16* v_s = shared + KeyBlock * D; + + // The two K/V buffers together hold at least Br rows. Use them once as Q staging, then retain + // all Q MMA fragments in registers for the complete split. + for (int chunk = tid; chunk < Br * (D / 8); chunk += Threads) { + const int row = chunk / (D / 8); + const int d = (chunk - row * (D / 8)) * 8; + int q_head = 0, token = 0; + noncausal_gqa_row_to_qt(row, kv_head, q_head, token); + const bool live = row < RowCount && token < valid; + const __nv_bfloat16* src = + q + bidirectional_gqa_q_index(live ? q_head : 0, d, live ? token : 0); + __nv_bfloat16* dst = &shared[row * D + bidirectional_gqa_swz(row, d)]; + cp_async_zfill<16, Cache::cg>(dst, src, live ? 16 : 0); + } + cp_commit(); + cp_wait<0>(); + __syncthreads(); + + const int gid = lane >> 2; + const int lid = lane & 3; + + const int a_mat = lane >> 3; + const int a_rin = lane & 7; + const int a_rowoff = a_rin + ((a_mat & 1) << 3); + const int a_coloff = (a_mat >> 1) << 3; + const int b_rin = lane & 7; + const int b_koff = ((lane >> 3) & 1) << 3; + + const int warp_row0 = warp * 16; + const int row0 = warp_row0 + gid; + const int row1 = row0 + 8; + const int q_position0 = + CyclicSwa ? context_state[row0 < RowCount ? row0 / kBidirectionalGqaGroup : 0] : 0; + const int q_position1 = + CyclicSwa ? context_state[row1 < RowCount ? row1 / kBidirectionalGqaGroup : 0] : 0; + unsigned af_q[QKKs][4]; +#pragma unroll + for (int ks = 0; ks < QKKs; ++ks) { + const int row = warp_row0 + a_rowoff; + const int col = ks * 16 + a_coloff; + ldmatrix_x4(af_q[ks][0], af_q[ks][1], af_q[ks][2], af_q[ks][3], + smem_addr(&shared[row * D + bidirectional_gqa_swz(row, col)])); + } + __syncthreads(); + + float acc[PVNt][4]; +#pragma unroll + for (int n = 0; n < PVNt; ++n) { +#pragma unroll + for (int item = 0; item < 4; ++item) { acc[n][item] = 0.0f; } + } + float m0 = -CUDART_INF_F; + float m1 = -CUDART_INF_F; + float l0 = 0.0f; + float l1 = 0.0f; + + const unsigned v_sbase = smem_addr(v_s); + const unsigned v_lane_base = v_sbase + static_cast(((lane >> 3) & 1) * 8 * RowBytes) + + static_cast(b_rin * RowBytes); + const unsigned v_as = static_cast((lane >> 4) << 4); + const unsigned v_r = static_cast(b_rin << 4); + + auto tile_metadata = [&](int iteration, bool& is_query, int& key0, int& valid_keys) { + is_query = iteration >= context_tile_count; + if (is_query) { + key0 = 0; + valid_keys = valid; + } else { + key0 = context_start + (tile_begin + iteration) * KeyBlock; + valid_keys = min(KeyBlock, length - key0); + } + }; + const auto tile_page = [&](bool is_query, int key0) { + if constexpr (CyclicSwa) { + return 0; + } else { + return is_query ? 0 : paged_page(key0); + } + }; + + bool current_is_query = false; + int current_key0 = 0; + int current_valid = 0; + tile_metadata(0, current_is_query, current_key0, current_valid); + int current_page = tile_page(current_is_query, current_key0); + bidirectional_gqa_stage_tile( + k_s, context_k, query_k, current_key0, current_valid, current_is_query, kv_head, + context_stride, current_page, tid); + cp_commit(); + + for (int iteration = 0; iteration < iterations; ++iteration) { + cp_wait<0>(); + __syncthreads(); + + bidirectional_gqa_stage_tile( + v_s, context_v, query_v, current_key0, current_valid, current_is_query, kv_head, + context_stride, current_page, tid); + cp_commit(); + + float score[QKNt][4]; +#pragma unroll + for (int nt = 0; nt < QKNt; ++nt) { + score[nt][0] = score[nt][1] = score[nt][2] = score[nt][3] = 0.0f; +#pragma unroll + for (int ks = 0; ks < QKKs; ++ks) { + unsigned bf[2]; + const int brow = nt * 8 + b_rin; + const int bcol = ks * 16 + b_koff; + ldmatrix_x2(bf[0], bf[1], + smem_addr(&k_s[brow * D + bidirectional_gqa_swz(brow, bcol)])); + mma_bf16(score[nt][0], score[nt][1], score[nt][2], score[nt][3], af_q[ks][0], + af_q[ks][1], af_q[ks][2], af_q[ks][3], bf[0], bf[1]); + } + } + + cp_wait<0>(); + __syncthreads(); + + bool next_is_query = false; + int next_key0 = 0; + int next_valid = 0; + int next_page = 0; + if (iteration + 1 < iterations) { + tile_metadata(iteration + 1, next_is_query, next_key0, next_valid); + if constexpr (CyclicSwa) { + next_page = tile_page(next_is_query, next_key0); + } else { + next_page = + !next_is_query && !current_is_query && (next_key0 >> 6) == (current_key0 >> 6) + ? current_page + : tile_page(next_is_query, next_key0); + } + bidirectional_gqa_stage_tile( + k_s, context_k, query_k, next_key0, next_valid, next_is_query, kv_head, + context_stride, next_page, tid); + cp_commit(); + } + + float block_m0 = -CUDART_INF_F; + float block_m1 = -CUDART_INF_F; +#pragma unroll + for (int nt = 0; nt < QKNt; ++nt) { + const int col0 = nt * 8 + 2 * lid; + const int col1 = col0 + 1; + const bool row0_live = row0 < RowCount && row0 / kBidirectionalGqaGroup < valid; + const bool row1_live = row1 < RowCount && row1 / kBidirectionalGqaGroup < valid; + const bool allow00 = + row0_live && col0 < current_valid && + (!CyclicSwa || current_is_query || current_key0 + col0 >= q_position0 - (Window - 1)); + const bool allow01 = + row0_live && col1 < current_valid && + (!CyclicSwa || current_is_query || current_key0 + col1 >= q_position0 - (Window - 1)); + const bool allow10 = + row1_live && col0 < current_valid && + (!CyclicSwa || current_is_query || current_key0 + col0 >= q_position1 - (Window - 1)); + const bool allow11 = + row1_live && col1 < current_valid && + (!CyclicSwa || current_is_query || current_key0 + col1 >= q_position1 - (Window - 1)); + score[nt][0] = allow00 ? score[nt][0] * scale : -CUDART_INF_F; + score[nt][1] = allow01 ? score[nt][1] * scale : -CUDART_INF_F; + score[nt][2] = allow10 ? score[nt][2] * scale : -CUDART_INF_F; + score[nt][3] = allow11 ? score[nt][3] * scale : -CUDART_INF_F; + block_m0 = fmaxf(block_m0, fmaxf(score[nt][0], score[nt][1])); + block_m1 = fmaxf(block_m1, fmaxf(score[nt][2], score[nt][3])); + } + block_m0 = warp_max<4>(block_m0, FullMask); + block_m1 = warp_max<4>(block_m1, FullMask); + + const float next_m0 = fmaxf(m0, block_m0); + const float next_m1 = fmaxf(m1, block_m1); + const float alpha0 = m0 == -CUDART_INF_F ? 0.0f : exp2_approx((m0 - next_m0) * Log2E); + const float alpha1 = m1 == -CUDART_INF_F ? 0.0f : exp2_approx((m1 - next_m1) * Log2E); + + unsigned p_frag[PVKs][4]; + float block_l0 = 0.0f; + float block_l1 = 0.0f; +#pragma unroll + for (int nt = 0; nt < QKNt; ++nt) { + const float p00 = + score[nt][0] > -CUDART_INF_F ? exp2_approx((score[nt][0] - next_m0) * Log2E) : 0.0f; + const float p01 = + score[nt][1] > -CUDART_INF_F ? exp2_approx((score[nt][1] - next_m0) * Log2E) : 0.0f; + const float p10 = + score[nt][2] > -CUDART_INF_F ? exp2_approx((score[nt][2] - next_m1) * Log2E) : 0.0f; + const float p11 = + score[nt][3] > -CUDART_INF_F ? exp2_approx((score[nt][3] - next_m1) * Log2E) : 0.0f; + block_l0 += p00 + p01; + block_l1 += p10 + p11; + const int pk = nt >> 1; + if ((nt & 1) == 0) { + p_frag[pk][0] = pack_bf16x2(p00, p01); + p_frag[pk][1] = pack_bf16x2(p10, p11); + } else { + p_frag[pk][2] = pack_bf16x2(p00, p01); + p_frag[pk][3] = pack_bf16x2(p10, p11); + } + } + block_l0 = warp_sum<4>(block_l0, FullMask); + block_l1 = warp_sum<4>(block_l1, FullMask); + + l0 = l0 * alpha0 + block_l0; + l1 = l1 * alpha1 + block_l1; + m0 = next_m0; + m1 = next_m1; +#pragma unroll + for (int n = 0; n < PVNt; ++n) { + acc[n][0] *= alpha0; + acc[n][1] *= alpha0; + acc[n][2] *= alpha1; + acc[n][3] *= alpha1; + } + + constexpr int PVTilePairs = (PVNt + 1) / 2; + constexpr int PVLoads = PVKs * PVTilePairs; + unsigned vf[2][4]; + ldmatrix_x4_t(vf[0][0], vf[0][1], vf[0][2], vf[0][3], + bidirectional_gqa_swz_addr(v_lane_base, 0u, v_as, v_r)); +#pragma unroll + for (int load = 0; load < PVLoads; ++load) { + const int pk = load / PVTilePairs; + const int n2 = (load % PVTilePairs) * 2; + const int cur = load & 1; + const int next = cur ^ 1; + if (load + 1 < PVLoads) { + const int next_pk = (load + 1) / PVTilePairs; + const int next_n2 = ((load + 1) % PVTilePairs) * 2; + ldmatrix_x4_t(vf[next][0], vf[next][1], vf[next][2], vf[next][3], + bidirectional_gqa_swz_addr( + v_lane_base + static_cast(next_pk * 16 * RowBytes), + static_cast(next_n2 << 4), v_as, v_r)); + } + mma_bf16(acc[n2][0], acc[n2][1], acc[n2][2], acc[n2][3], p_frag[pk][0], p_frag[pk][1], + p_frag[pk][2], p_frag[pk][3], vf[cur][0], vf[cur][1]); + if (n2 + 1 < PVNt) { + mma_bf16(acc[n2 + 1][0], acc[n2 + 1][1], acc[n2 + 1][2], acc[n2 + 1][3], + p_frag[pk][0], p_frag[pk][1], p_frag[pk][2], p_frag[pk][3], vf[cur][2], + vf[cur][3]); + } + } + + current_is_query = next_is_query; + current_key0 = next_key0; + current_valid = next_valid; + current_page = next_page; + } + + if constexpr (!DirectOutput) { + if (lid == 0) { + const int row0 = warp_row0 + gid; + const int row1 = row0 + 8; + if (row0 < RowCount) { + int q_head = 0, token = 0; + noncausal_gqa_row_to_qt(row0, kv_head, q_head, token); + partial_m[bidirectional_gqa_stat_index(q_head, token, split)] = m0; + partial_l[bidirectional_gqa_stat_index(q_head, token, split)] = l0; + } + if (row1 < RowCount) { + int q_head = 0, token = 0; + noncausal_gqa_row_to_qt(row1, kv_head, q_head, token); + partial_m[bidirectional_gqa_stat_index(q_head, token, split)] = m1; + partial_l[bidirectional_gqa_stat_index(q_head, token, split)] = l1; + } + } + } + +#pragma unroll + for (int n = 0; n < PVNt; ++n) { + const int d0 = n * 8 + 2 * lid; + const int row0 = warp_row0 + gid; + const int row1 = row0 + 8; + if (row0 < RowCount) { + int q_head = 0, token = 0; + noncausal_gqa_row_to_qt(row0, kv_head, q_head, token); + if constexpr (DirectOutput) { + const float inv_l = l0 > 0.0f ? 1.0f / l0 : 0.0f; + const auto dst = bidirectional_gqa_q_index(q_head, d0, token); + store_vec(&out[dst], pack_bf16x2(acc[n][0] * inv_l, acc[n][1] * inv_l)); + } else { + const auto dst = bidirectional_gqa_partial_index(q_head, d0, token, split); + store_vec(&partial_acc[dst], pack_bf16x2(acc[n][0], acc[n][1])); + } + } + if (row1 < RowCount) { + int q_head = 0, token = 0; + noncausal_gqa_row_to_qt(row1, kv_head, q_head, token); + if constexpr (DirectOutput) { + const float inv_l = l1 > 0.0f ? 1.0f / l1 : 0.0f; + const auto dst = bidirectional_gqa_q_index(q_head, d0, token); + store_vec(&out[dst], pack_bf16x2(acc[n][2] * inv_l, acc[n][3] * inv_l)); + } else { + const auto dst = bidirectional_gqa_partial_index(q_head, d0, token, split); + store_vec(&partial_acc[dst], pack_bf16x2(acc[n][2], acc[n][3])); + } + } + } +} + +template +__launch_bounds__(WarpsPerCta * 32, 2) __global__ void bidirectional_gqa_split_partial_kernel( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ query_k, + const __nv_bfloat16* __restrict__ query_v, const std::int32_t* __restrict__ context_length, + const std::int32_t* __restrict__ valid_columns, const std::int32_t* __restrict__ table_rows, + const __nv_bfloat16* __restrict__ context_k, const __nv_bfloat16* __restrict__ context_v, + const std::int32_t* __restrict__ block_tables, int physical_pages, int logical_pages, + int max_context, int split_capacity, float scale, __nv_bfloat16* __restrict__ partial_acc, + float* __restrict__ partial_m, float* __restrict__ partial_l, __nv_bfloat16* __restrict__ out) { + noncausal_gqa_split_partial_body( + q, query_k, query_v, context_length, valid_columns, table_rows, context_k, context_v, + block_tables, physical_pages, logical_pages, max_context, split_capacity, scale, + partial_acc, partial_m, partial_l, out); +} + +template +__launch_bounds__(WarpsPerCta * 32, 2) __global__ void swa_split_partial_kernel( + const __nv_bfloat16* __restrict__ q, const __nv_bfloat16* __restrict__ query_k, + const __nv_bfloat16* __restrict__ query_v, const std::int32_t* __restrict__ positions, + const std::int32_t* __restrict__ valid_columns, const std::int32_t* __restrict__ lanes, + const __nv_bfloat16* __restrict__ context_k, const __nv_bfloat16* __restrict__ context_v, + int padded_context, int max_context, int split_capacity, float scale, + __nv_bfloat16* __restrict__ partial_acc, float* __restrict__ partial_m, + float* __restrict__ partial_l, __nv_bfloat16* __restrict__ out) { + noncausal_gqa_split_partial_body( + q, query_k, query_v, positions, valid_columns, lanes, context_k, context_v, nullptr, + padded_context, 0, max_context, split_capacity, scale, partial_acc, partial_m, partial_l, + out); +} + +template +__device__ __forceinline__ void +noncausal_gqa_reduce_body(const __nv_bfloat16* __restrict__ partial_acc, + const float* __restrict__ partial_m, const float* __restrict__ partial_l, + const std::int32_t* __restrict__ context_state, + const std::int32_t* __restrict__ valid_columns, int max_context, + int split_capacity, __nv_bfloat16* __restrict__ out) { + const int q_head = static_cast(blockIdx.x); + const int token = static_cast(blockIdx.y); + const int batch = static_cast(blockIdx.z); + const int tid = static_cast(threadIdx.x); + constexpr std::int64_t QueryElements = + static_cast(kBidirectionalGqaHeadDim) * kBidirectionalGqaQHeads * Tokens; + constexpr std::int64_t StatElements = + static_cast(kBidirectionalGqaQHeads) * Tokens; + partial_acc += QueryElements * split_capacity * batch; + partial_m += StatElements * split_capacity * batch; + partial_l += StatElements * split_capacity * batch; + out += QueryElements * batch; + if constexpr (CyclicSwa) { + context_state += static_cast(Tokens) * batch; + } else { + context_state += batch; + } + const int length = context_state[0]; + if (q_head >= kBidirectionalGqaQHeads || token >= Tokens) { return; } + if (length < 0 || length > max_context || token >= valid_columns[batch]) { + if (tid < kBidirectionalGqaHeadDim) { + out[bidirectional_gqa_q_index(q_head, tid, token)] = __float2bfloat16(0.0f); + } + return; + } + + const int context_count = CyclicSwa ? min(length, Window - 1) : length; + const int context_tiles = (context_count + KeyBlock - 1) / KeyBlock; + const int active_splits = context_tiles > 0 ? min(context_tiles, split_capacity) : 1; + __shared__ float reduce[128]; + + float local_m = -CUDART_INF_F; + for (int split = tid; split < active_splits; split += blockDim.x) { + local_m = + fmaxf(local_m, partial_m[bidirectional_gqa_stat_index(q_head, token, split)]); + } + reduce[tid] = local_m; + __syncthreads(); + for (int stride = 64; stride > 0; stride >>= 1) { + if (tid < stride) { reduce[tid] = fmaxf(reduce[tid], reduce[tid + stride]); } + __syncthreads(); + } + const float global_m = reduce[0]; + __syncthreads(); + + float local_l = 0.0f; + for (int split = tid; split < active_splits; split += blockDim.x) { + const auto idx = bidirectional_gqa_stat_index(q_head, token, split); + local_l += partial_l[idx] * expf(partial_m[idx] - global_m); + } + reduce[tid] = local_l; + __syncthreads(); + for (int stride = 64; stride > 0; stride >>= 1) { + if (tid < stride) { reduce[tid] += reduce[tid + stride]; } + __syncthreads(); + } + const float global_l = reduce[0]; + + if (tid < kBidirectionalGqaHeadDim) { + float numerator = 0.0f; + for (int split = 0; split < active_splits; ++split) { + const auto stat = bidirectional_gqa_stat_index(q_head, token, split); + const float weight = expf(partial_m[stat] - global_m); + numerator += __bfloat162float(partial_acc[bidirectional_gqa_partial_index( + q_head, tid, token, split)]) * + weight; + } + const float value = global_l > 0.0f ? numerator / global_l : 0.0f; + out[bidirectional_gqa_q_index(q_head, tid, token)] = __float2bfloat16(value); + } +} + +template +__launch_bounds__(128, 2) __global__ + void bidirectional_gqa_reduce_kernel(const __nv_bfloat16* __restrict__ partial_acc, + const float* __restrict__ partial_m, + const float* __restrict__ partial_l, + const std::int32_t* __restrict__ context_length, + const std::int32_t* __restrict__ valid_columns, + int max_context, int split_capacity, + __nv_bfloat16* __restrict__ out) { + noncausal_gqa_reduce_body(partial_acc, partial_m, partial_l, + context_length, valid_columns, max_context, + split_capacity, out); +} + +template +__launch_bounds__(WarpsPerBlock * 32, 2) __global__ + void swa_reduce_kernel(const __nv_bfloat16* __restrict__ partial_acc, + const float* __restrict__ partial_m, const float* __restrict__ partial_l, + const std::int32_t* __restrict__ positions, + const std::int32_t* __restrict__ valid_columns, int max_context, + int split_capacity, __nv_bfloat16* __restrict__ out) { + static_assert(WarpsPerBlock >= 1 && WarpsPerBlock <= 8); + constexpr int MaxSplits = 128; + constexpr unsigned Mask = 0xffffffffu; + __shared__ float weights[WarpsPerBlock][MaxSplits]; + + const int warp = static_cast(threadIdx.x) >> 5; + const int lane = static_cast(threadIdx.x) & 31; + const int batch = static_cast(blockIdx.z); + const int output_row = static_cast(blockIdx.x) * WarpsPerBlock + warp; + const int token = output_row / kBidirectionalGqaQHeads; + const int q_head = output_row - token * kBidirectionalGqaQHeads; + if (warp >= WarpsPerBlock || token >= Tokens) return; + + constexpr std::int64_t QueryElements = + static_cast(kBidirectionalGqaHeadDim) * kBidirectionalGqaQHeads * Tokens; + constexpr std::int64_t StatElements = + static_cast(kBidirectionalGqaQHeads) * Tokens; + partial_acc += QueryElements * split_capacity * batch; + partial_m += StatElements * split_capacity * batch; + partial_l += StatElements * split_capacity * batch; + positions += static_cast(Tokens) * batch; + out += QueryElements * batch; + + const int length = positions[0]; + if (length < 0 || length > max_context || token >= valid_columns[batch]) { +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int d = lane + item * 32; + out[bidirectional_gqa_q_index(q_head, d, token)] = __float2bfloat16(0.0f); + } + return; + } + const int context_count = min(length, Window - 1); + const int context_tiles = (context_count + KeyBlock - 1) / KeyBlock; + const int active_splits = context_tiles > 0 ? min(context_tiles, split_capacity) : 1; + + float local_m = -CUDART_INF_F; + for (int split = lane; split < active_splits; split += 32) { + local_m = + fmaxf(local_m, partial_m[bidirectional_gqa_stat_index(q_head, token, split)]); + } + const float global_m = warp_max<32>(local_m, Mask); + + float local_l = 0.0f; + for (int split = lane; split < active_splits; split += 32) { + const auto stat = bidirectional_gqa_stat_index(q_head, token, split); + const float weight = expf(partial_m[stat] - global_m); + weights[warp][split] = weight; + local_l += partial_l[stat] * weight; + } + const float global_l = warp_sum<32>(local_l, Mask); + __syncwarp(Mask); + +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int d = lane + item * 32; + float numerator = 0.0f; + for (int split = 0; split < active_splits; ++split) { + numerator += + __bfloat162float( + partial_acc[bidirectional_gqa_partial_index(q_head, d, token, split)]) * + weights[warp][split]; + } + const float value = global_l > 0.0f ? numerator / global_l : 0.0f; + out[bidirectional_gqa_q_index(q_head, d, token)] = __float2bfloat16(value); + } +} + +} // namespace ninfer::ops diff --git a/src/ops/kernel/dflash2_grouped_conv.cuh b/src/ops/kernel/dflash2_grouped_conv.cuh new file mode 100644 index 0000000000..8b3b5b3530 --- /dev/null +++ b/src/ops/kernel/dflash2_grouped_conv.cuh @@ -0,0 +1,44 @@ +#pragma once + +// Implements: include/ninfer/ops/dflash2_grouped_conv.h +// Match: contiguous BF16 hidden/delta/base with the documented token-major layout. + +#include + +#include + +namespace ninfer::ops { + +template +__launch_bounds__(Block) __global__ void dflash2_grouped_conv_kernel( + const __nv_bfloat16* __restrict__ hidden, const __nv_bfloat16* __restrict__ delta, + const __nv_bfloat16* __restrict__ base, __nv_bfloat16* __restrict__ out, int hidden_size, + int tokens, int taps, int groups, int group_size, int block_size, int side) { + const int h = static_cast(blockIdx.x) * Block + static_cast(threadIdx.x); + if (h >= hidden_size) { return; } + const int g = h / group_size; + for (int t = 0; t < tokens; ++t) { + const std::int64_t hidden_offset = static_cast(h) + + static_cast(hidden_size) * t; + const std::int64_t delta_n = + static_cast(groups) * (side * taps) + g; + const std::int64_t delta_offset = + delta_n + static_cast(2) * taps * groups * t; + float value = __bfloat162float(hidden[hidden_offset]) * + (__bfloat162float(base[h]) + __bfloat162float(delta[delta_offset])); + const int position = t & (block_size - 1); + for (int tap = 1; tap < taps; ++tap) { + if (position < tap) { continue; } + const std::int64_t prev = static_cast(h) + + static_cast(hidden_size) * (t - tap); + const std::int64_t coefficient = + delta_offset + static_cast(tap) * groups; + value += __bfloat162float(hidden[prev]) * + (__bfloat162float(base[static_cast(tap) * hidden_size + h]) + + __bfloat162float(delta[coefficient])); + } + out[hidden_offset] = __float2bfloat16(value); + } +} + +} // namespace ninfer::ops diff --git a/src/ops/kernel/dflash2_selector.cuh b/src/ops/kernel/dflash2_selector.cuh new file mode 100644 index 0000000000..296535e5b0 --- /dev/null +++ b/src/ops/kernel/dflash2_selector.cuh @@ -0,0 +1,191 @@ +#pragma once + +// Implements: include/ninfer/ops/dflash2_selector.h +// Match: token-major BF16 logits/projection and the documented scratch layouts. + +#include "ops/launcher/dflash2_selector.h" + +#include +#include + +#include + +namespace ninfer::ops { + +__device__ __forceinline__ int dflash2_selector_candidate_offset(int b, int batch, int s, + int steps, int c) { + return b + batch * (s + steps * c); +} + +__device__ __forceinline__ int dflash2_selector_score_offset(int b, int batch, int s, int steps, + int p, int c, int top_k) { + return b + batch * (s + steps * (p + top_k * c)); +} + +template +__launch_bounds__(Block) __global__ void dflash2_selector_topk_kernel( + const __nv_bfloat16* __restrict__ logits, std::int32_t* __restrict__ candidates, + float* __restrict__ unary, int vocab, int batch, int steps, int columns) { + const int column = static_cast(blockIdx.x); + const int tid = static_cast(threadIdx.x); + if (column >= columns) { return; } + + float local_value[K]; + std::int32_t local_index[K]; +#pragma unroll + for (int k = 0; k < K; ++k) { + local_value[k] = -CUDART_INF_F; + local_index[k] = vocab; + } + const __nv_bfloat16* col_logits = logits + static_cast(column) * vocab; + for (int v = tid; v < vocab; v += Block) { + const float value = __bfloat162float(col_logits[v]); + // Insertion into the sorted list while keeping (value, -id) ordering. + for (int k = 0; k < K; ++k) { + const bool better = + value > local_value[k] || (value == local_value[k] && v < local_index[k]); + if (better) { + for (int tail = K - 1; tail > k; --tail) { + local_value[tail] = local_value[tail - 1]; + local_index[tail] = local_index[tail - 1]; + } + local_value[k] = value; + local_index[k] = v; + break; + } + } + } + + __shared__ float shared_value[Block][K]; + __shared__ std::int32_t shared_index[Block][K]; +#pragma unroll + for (int k = 0; k < K; ++k) { + shared_value[tid][k] = local_value[k]; + shared_index[tid][k] = local_index[k]; + } + __syncthreads(); + + if (tid == 0) { + float merged_value[K]; + std::int32_t merged_index[K]; +#pragma unroll + for (int k = 0; k < K; ++k) { + merged_value[k] = -CUDART_INF_F; + merged_index[k] = vocab; + } + for (int thread = 0; thread < Block; ++thread) { + for (int k = 0; k < K; ++k) { + const float value = shared_value[thread][k]; + const std::int32_t index = shared_index[thread][k]; + for (int slot = 0; slot < K; ++slot) { + const bool better = + value > merged_value[slot] || + (value == merged_value[slot] && index < merged_index[slot]); + if (better) { + for (int tail = K - 1; tail > slot; --tail) { + merged_value[tail] = merged_value[tail - 1]; + merged_index[tail] = merged_index[tail - 1]; + } + merged_value[slot] = value; + merged_index[slot] = index; + break; + } + } + } + } + const int b = column / steps; + const int s = column - b * steps; + for (int k = 0; k < K; ++k) { + const int offset = + dflash2_selector_candidate_offset(b, batch, s, steps, k); + candidates[offset] = merged_index[k]; + unary[offset] = merged_value[k]; + } + } +} + +template +__launch_bounds__(Block) __global__ void dflash2_selector_scores_kernel( + const std::int32_t* __restrict__ candidates, const float* __restrict__ unary, + const __nv_bfloat16* __restrict__ projected, + const __nv_bfloat16* __restrict__ predecessor_codebook, + const __nv_bfloat16* __restrict__ successor_codebook, + const std::int32_t* __restrict__ anchors, float* __restrict__ scores, int vocab, int batch, + int steps, int top_k, int columns) { + constexpr int K = detail::kDflash2SelectorTopK; + constexpr int R = detail::kDflash2SelectorRank; + const int flat = static_cast(blockIdx.x); + const int tid = static_cast(threadIdx.x); + if (flat >= columns || tid >= K * K) { return; } + const int b = flat / steps; + const int s = flat - b * steps; + const int p = tid / K; + const int c = tid - p * K; + + __shared__ std::int32_t shared_candidates[K]; + __shared__ float shared_unary[K]; + if (tid < K) { + const int offset = dflash2_selector_candidate_offset(b, batch, s, steps, tid); + shared_candidates[tid] = candidates[offset]; + shared_unary[tid] = unary[offset]; + } + __syncthreads(); + + if (s == 0 && p > 0) { + scores[dflash2_selector_score_offset(b, batch, s, steps, p, c, top_k)] = -CUDART_INF_F; + return; + } + const std::int32_t predecessor = + s == 0 ? anchors[b] + : candidates[dflash2_selector_candidate_offset(b, batch, s - 1, steps, p)]; + const std::int32_t successor = shared_candidates[c]; + const __nv_bfloat16* hidden = projected + static_cast(R) * flat; + const __nv_bfloat16* predecessor_row = + predecessor_codebook + static_cast(predecessor) * R; + const __nv_bfloat16* successor_row = + successor_codebook + static_cast(successor) * R; + float pair = 0.0f; +#pragma unroll 8 + for (int r = 0; r < R; ++r) { + pair += __bfloat162float(hidden[r]) * __bfloat162float(predecessor_row[r]) * + __bfloat162float(successor_row[r]); + } + scores[dflash2_selector_score_offset(b, batch, s, steps, p, c, top_k)] = + pair + shared_unary[c]; +} + +__global__ void dflash2_selector_walk_kernel(const std::int32_t* __restrict__ candidates, + const float* __restrict__ scores, + std::int32_t* __restrict__ drafts, int batch, + int steps, int top_k) { + constexpr int K = detail::kDflash2SelectorTopK; + const int b = static_cast(blockIdx.x); + const int lane = static_cast(threadIdx.x); + constexpr unsigned Mask = 0xffffffffu; + int previous = 0; + for (int s = 0; s < steps; ++s) { + float value = -CUDART_INF_F; + if (lane < K) { + value = scores[dflash2_selector_score_offset(b, batch, s, steps, previous, lane, + top_k)]; + } +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + value = fmaxf(value, __shfl_xor_sync(Mask, value, offset)); + } + const float best = __shfl_sync(Mask, value, 0); + const bool equal = lane < K && value == best; + const std::int32_t candidate_rank = equal ? lane : K; + std::int32_t chosen = candidate_rank; +#pragma unroll + for (int offset = 16; offset > 0; offset >>= 1) { + chosen = min(chosen, __shfl_xor_sync(Mask, chosen, offset)); + } + const std::int32_t token = + candidates[dflash2_selector_candidate_offset(b, batch, s, steps, chosen)]; + drafts[b * steps + s] = token; + previous = chosen; + } +} + +} // namespace ninfer::ops diff --git a/src/ops/kv_cache/append/kernel.cuh b/src/ops/kv_cache/append/kernel.cuh index 7ea9aebfb0..07a0f217c9 100644 --- a/src/ops/kv_cache/append/kernel.cuh +++ b/src/ops/kv_cache/append/kernel.cuh @@ -337,7 +337,6 @@ inline constexpr int kKVCacheAppendPrefixHeadDim = 128; inline constexpr int kKVCacheAppendPrefixHeads = 8; inline constexpr int kKVCacheAppendPrefixWindow = 4096; inline constexpr int kKVCacheAppendPrefixPage = 64; - __device__ __forceinline__ void kv_cache_append_prefix_copy_cyclic_unit( const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, __nv_bfloat16* __restrict__ cache_k, __nv_bfloat16* __restrict__ cache_v, int token, @@ -389,6 +388,7 @@ __device__ __forceinline__ void kv_cache_append_prefix_copy_paged_unit( *reinterpret_cast(&cache_v[dst + 8]) = v1; } +template __global__ void kv_cache_append_prefix_cyclic_kernel( const __nv_bfloat16* __restrict__ k, const __nv_bfloat16* __restrict__ v, const std::int32_t* __restrict__ positions, const std::int32_t* __restrict__ counts, @@ -419,7 +419,7 @@ __global__ void kv_cache_append_prefix_cyclic_kernel( const int token = static_cast(blockIdx.x) * TokensPerBlock + local_token; if (token >= count) return; const int position = positions[token]; - const int slot = position & (kKVCacheAppendPrefixWindow - 1); + const int slot = position & (Window - 1); kv_cache_append_prefix_copy_cyclic_unit(k, v, cache_k, cache_v, token, unit_in_token, slot, padded_capacity); } diff --git a/src/ops/kv_cache/append/kv_cache_append.cpp b/src/ops/kv_cache/append/kv_cache_append.cpp index fdfcbb36db..c27b473ee9 100644 --- a/src/ops/kv_cache/append/kv_cache_append.cpp +++ b/src/ops/kv_cache/append/kv_cache_append.cpp @@ -151,8 +151,11 @@ void validate_paged_cache(const PagedKVBatchLayerView& cache, } void validate_cyclic_cache(const CyclicKVCacheLayerView& cache, - KVCacheAppendPrefixExecutionEnvelope envelope) { - if (cache.num_kv_heads != kKVHeads || cache.head_dim != kHeadDim || cache.capacity != kWindow || + KVCacheAppendPrefixExecutionEnvelope envelope, + std::uint32_t window) { + if (cache.num_kv_heads != kKVHeads || cache.head_dim != kHeadDim || + cache.capacity != window || + (window != 2048 && window != 4096) || cache.padded_capacity < cache.capacity || cache.padded_capacity > static_cast(std::numeric_limits::max()) || @@ -212,10 +215,12 @@ void kv_cache_append_prefix(const Tensor& k, const Tensor& v, const Tensor& posi void kv_cache_append_prefix(const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& counts, const Tensor& lanes, KVCacheAppendPrefixExecutionEnvelope envelope, - CyclicKVCacheLayerView cache, cudaStream_t stream) { + CyclicKVCacheLayerView cache, std::uint32_t window, + cudaStream_t stream) { const auto plan = validate_inputs(k, v, positions, counts, lanes, envelope); - validate_cyclic_cache(cache, envelope); - detail::kv_cache_append_prefix_launch(k, v, positions, counts, lanes, cache, plan, stream); + validate_cyclic_cache(cache, envelope, window); + detail::kv_cache_append_prefix_launch(k, v, positions, counts, lanes, cache, plan, window, + stream); } } // namespace ninfer::ops diff --git a/src/ops/kv_cache/append/launch.cu b/src/ops/kv_cache/append/launch.cu index 811941d135..087575c6d3 100644 --- a/src/ops/kv_cache/append/launch.cu +++ b/src/ops/kv_cache/append/launch.cu @@ -128,7 +128,8 @@ void launch_paged(const Tensor& k, const Tensor& v, const Tensor& positions, con void launch_cyclic(const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& counts, const Tensor& lanes, CyclicKVCacheLayerView cache, - const KVCacheAppendPrefixPlan& plan, cudaStream_t stream) { + const KVCacheAppendPrefixPlan& plan, std::uint32_t window, + cudaStream_t stream) { validate_plan(k, plan); if (plan.max_count == 0) return; auto* cache_k = static_cast<__nv_bfloat16*>(cache.k.data); @@ -141,10 +142,22 @@ void launch_cyclic(const Tensor& k, const Tensor& v, const Tensor& positions, co const int padded = static_cast(cache.padded_capacity); const dim3 grid(1 + (plan.max_count - 1) / 4, k.ne[3], 1); - kv_cache_append_prefix_cyclic_kernel<<>>( - input_k, input_v, pos, count, lane, cache_k, cache_v, plan.min_count, plan.max_count, - plan.tokens, padded); - CUDA_CHECK(cudaGetLastError()); + const auto launch_window = [&]() { + kv_cache_append_prefix_cyclic_kernel<<>>( + input_k, input_v, pos, count, lane, cache_k, cache_v, plan.min_count, plan.max_count, + plan.tokens, padded); + CUDA_CHECK(cudaGetLastError()); + }; + switch (window) { + case 2048: + launch_window.template operator()<2048>(); + break; + case 4096: + launch_window.template operator()<4096>(); + break; + default: + throw std::invalid_argument("kv_cache_append_prefix: unsupported cyclic window"); + } } } // namespace @@ -211,8 +224,9 @@ void kv_cache_append_prefix_launch(const Tensor& k, const Tensor& v, const Tenso void kv_cache_append_prefix_launch(const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& counts, const Tensor& lanes, CyclicKVCacheLayerView cache, - const KVCacheAppendPrefixPlan& plan, cudaStream_t stream) { - launch_cyclic(k, v, positions, counts, lanes, cache, plan, stream); + const KVCacheAppendPrefixPlan& plan, std::uint32_t window, + cudaStream_t stream) { + launch_cyclic(k, v, positions, counts, lanes, cache, plan, window, stream); } } // namespace ninfer::ops::detail diff --git a/src/ops/kv_cache/append/launch.h b/src/ops/kv_cache/append/launch.h index 250b382af3..12c7fbd665 100644 --- a/src/ops/kv_cache/append/launch.h +++ b/src/ops/kv_cache/append/launch.h @@ -28,6 +28,7 @@ void kv_cache_append_prefix_launch(const Tensor& k, const Tensor& v, const Tenso void kv_cache_append_prefix_launch(const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& counts, const Tensor& lanes, CyclicKVCacheLayerView cache, - const KVCacheAppendPrefixPlan& plan, cudaStream_t stream); + const KVCacheAppendPrefixPlan& plan, std::uint32_t window, + cudaStream_t stream); } // namespace ninfer::ops::detail diff --git a/src/ops/launcher/bidirectional_gqa_attention.cu b/src/ops/launcher/bidirectional_gqa_attention.cu new file mode 100644 index 0000000000..41c7831b48 --- /dev/null +++ b/src/ops/launcher/bidirectional_gqa_attention.cu @@ -0,0 +1,185 @@ +#include "ops/launcher/bidirectional_gqa_attention.h" + +#include "core/device.h" +#include "ops/kernel/bidirectional_gqa_attention.cuh" + +#include +#include +#include + +namespace ninfer::ops::detail { +namespace { + +template +void dispatch_token_case(Launch&& launch) { + constexpr int Warps = (Tokens + 3) / 4; + launch.template operator()(); +} + +template +void dispatch_tokens(std::int32_t tokens, Launch&& launch) { + switch (tokens) { +#define NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(TOKENS) \ + case TOKENS: \ + dispatch_token_case(launch); \ + return + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(1); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(2); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(3); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(4); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(5); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(6); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(7); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(8); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(9); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(10); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(11); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(12); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(13); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(14); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(15); + NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE(16); +#undef NINFER_BIDIRECTIONAL_GQA_TOKEN_CASE + default: + throw std::invalid_argument("bidirectional_gqa_attention: unsupported T"); + } +} + +} // namespace + +BidirectionalGqaPlan bidirectional_gqa_resolve_plan(std::int32_t tokens, + GqaContextExecutionEnvelope envelope) { + if (tokens < 1 || tokens > 16) { + throw std::invalid_argument("bidirectional_gqa_attention plan: T must be 1..16"); + } + if (envelope.min_context > envelope.max_context) { + throw std::invalid_argument("bidirectional_gqa_attention plan: invalid envelope"); + } + const std::int32_t warps = (tokens + 3) / 4; + const bool direct = envelope.max_context == 0; + const std::int32_t key_block = + direct || tokens <= 8 || envelope.max_context <= 65536u ? 32 : 64; + std::int32_t split_limit = 32; + if (tokens <= 8) { + split_limit = + envelope.max_context <= 131072u ? 32 : (envelope.max_context <= 196608u ? 48 : 64); + } else if (key_block == 64) { + split_limit = + envelope.max_context <= 131072u ? 32 : (envelope.max_context <= 196608u ? 38 : 40); + } + const std::uint32_t envelope_tiles = + (envelope.max_context + static_cast(key_block) - 1u) / + static_cast(key_block); + const std::int32_t splits = + direct ? 1 : std::min(split_limit, std::max(1, static_cast(envelope_tiles))); + return { + .route = direct ? BidirectionalGqaRoute::Direct : BidirectionalGqaRoute::SplitKv, + .tokens = tokens, + .warps = warps, + .key_block = key_block, + .split_capacity = splits, + }; +} + +const char* bidirectional_gqa_route_name(BidirectionalGqaRoute route) { + switch (route) { + case BidirectionalGqaRoute::Direct: + return "direct"; + case BidirectionalGqaRoute::SplitKv: + return "split_kv"; + } + return "unknown"; +} + +void bidirectional_gqa_attention_launch(const Tensor& q, const Tensor& query_k, + const Tensor& query_v, const Tensor& context_lengths, + const Tensor& valid_columns, const Tensor& table_rows, + float scale, const PagedKVBatchLayerView& context, + const BidirectionalGqaPlan& plan, Tensor& partial_acc, + Tensor& partial_m, Tensor& partial_l, Tensor& out, + cudaStream_t stream) { + dispatch_tokens(q.ne[2], [&]() { + const bool direct = plan.route == BidirectionalGqaRoute::Direct; + if (plan.warps != Warps || plan.split_capacity < 1 || + plan.split_capacity > kBidirectionalGqaMaxSplit) { + throw std::invalid_argument("bidirectional_gqa_attention: inconsistent plan"); + } + if (direct) { + if (plan.key_block != 32 || plan.split_capacity != 1) { + throw std::invalid_argument("bidirectional_gqa_attention: inconsistent plan"); + } + constexpr int KeyBlock = 32; + constexpr std::size_t SmemBytes = + 2u * KeyBlock * kBidirectionalGqaHeadDim * sizeof(__nv_bfloat16); + const dim3 direct_grid(kBidirectionalGqaKVHeads, 1, q.ne[3]); + bidirectional_gqa_split_partial_kernel + <<>>( + static_cast(q.data), + static_cast(query_k.data), + static_cast(query_v.data), + static_cast(context_lengths.data), + static_cast(valid_columns.data), + static_cast(table_rows.data), + static_cast(context.k_pages.data), + static_cast(context.v_pages.data), + static_cast(context.block_tables.data), + context.k_pages.ne[2], context.block_tables.ne[0], + context.block_tables.ne[0] * kPagedKVPageSize, 1, scale, + static_cast<__nv_bfloat16*>(partial_acc.data), + static_cast(partial_m.data), static_cast(partial_l.data), + static_cast<__nv_bfloat16*>(out.data)); + CUDA_CHECK(cudaGetLastError()); + return; + } + + if (plan.route != BidirectionalGqaRoute::SplitKv) { + throw std::invalid_argument("bidirectional_gqa_attention: inconsistent plan"); + } + + const auto launch_split = [&]() { + constexpr std::size_t SmemBytes = + 2u * KeyBlock * kBidirectionalGqaHeadDim * sizeof(__nv_bfloat16); + const dim3 partial_grid(kBidirectionalGqaKVHeads, plan.split_capacity, q.ne[3]); + bidirectional_gqa_split_partial_kernel + <<>>( + static_cast(q.data), + static_cast(query_k.data), + static_cast(query_v.data), + static_cast(context_lengths.data), + static_cast(valid_columns.data), + static_cast(table_rows.data), + static_cast(context.k_pages.data), + static_cast(context.v_pages.data), + static_cast(context.block_tables.data), + context.k_pages.ne[2], context.block_tables.ne[0], + context.block_tables.ne[0] * kPagedKVPageSize, plan.split_capacity, scale, + static_cast<__nv_bfloat16*>(partial_acc.data), + static_cast(partial_m.data), static_cast(partial_l.data), + static_cast<__nv_bfloat16*>(out.data)); + CUDA_CHECK(cudaGetLastError()); + const dim3 reduce_grid(kBidirectionalGqaQHeads, Tokens, q.ne[3]); + bidirectional_gqa_reduce_kernel<<>>( + static_cast(partial_acc.data), + static_cast(partial_m.data), + static_cast(partial_l.data), + static_cast(context_lengths.data), + static_cast(valid_columns.data), + context.block_tables.ne[0] * kPagedKVPageSize, plan.split_capacity, + static_cast<__nv_bfloat16*>(out.data)); + CUDA_CHECK(cudaGetLastError()); + }; + if (plan.key_block == 32) { + launch_split.template operator()<32>(); + return; + } + if constexpr (Tokens > 8) { + if (plan.key_block == 64) { + launch_split.template operator()<64>(); + return; + } + } + throw std::invalid_argument("bidirectional_gqa_attention: inconsistent plan"); + }); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/bidirectional_gqa_attention.h b/src/ops/launcher/bidirectional_gqa_attention.h new file mode 100644 index 0000000000..5828a451e2 --- /dev/null +++ b/src/ops/launcher/bidirectional_gqa_attention.h @@ -0,0 +1,33 @@ +#pragma once + +#include "ninfer/ops/bidirectional_gqa_attention.h" + +namespace ninfer::ops::detail { + +enum class BidirectionalGqaRoute { + Direct, + SplitKv, +}; + +struct BidirectionalGqaPlan { + BidirectionalGqaRoute route = BidirectionalGqaRoute::SplitKv; + std::int32_t tokens = 0; + std::int32_t warps = 0; + std::int32_t key_block = 0; + std::int32_t split_capacity = 0; +}; + +[[nodiscard]] BidirectionalGqaPlan +bidirectional_gqa_resolve_plan(std::int32_t tokens, GqaContextExecutionEnvelope envelope); + +[[nodiscard]] const char* bidirectional_gqa_route_name(BidirectionalGqaRoute route); + +void bidirectional_gqa_attention_launch(const Tensor& q, const Tensor& query_k, + const Tensor& query_v, const Tensor& context_lengths, + const Tensor& valid_columns, const Tensor& table_rows, + float scale, const PagedKVBatchLayerView& context, + const BidirectionalGqaPlan& plan, Tensor& partial_acc, + Tensor& partial_m, Tensor& partial_l, Tensor& out, + cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/dflash2_grouped_conv.cu b/src/ops/launcher/dflash2_grouped_conv.cu new file mode 100644 index 0000000000..eb050cb9c8 --- /dev/null +++ b/src/ops/launcher/dflash2_grouped_conv.cu @@ -0,0 +1,29 @@ +#include "ops/launcher/dflash2_grouped_conv.h" + +#include "core/device.h" +#include "ops/kernel/dflash2_grouped_conv.cuh" + +#include +#include + +namespace ninfer::ops::detail { + +void dflash2_grouped_conv_launch(const Tensor& hidden, const Tensor& delta, const Weight& base, + std::int32_t block_size, std::int32_t group_size, + std::int32_t taps, std::int32_t side, Tensor& out, + cudaStream_t stream) { + constexpr int kBlock = 256; + const std::int32_t hidden_size = hidden.ne[0]; + const std::int32_t tokens = hidden.ne[1]; + const std::int32_t groups = hidden_size / group_size; + const int grid = static_cast( + std::max(1, (static_cast(hidden_size) + kBlock - 1) / kBlock)); + dflash2_grouped_conv_kernel<<>>( + static_cast(hidden.data), + static_cast(delta.data), + static_cast(base.qdata), static_cast<__nv_bfloat16*>(out.data), + hidden_size, tokens, taps, groups, group_size, block_size, side); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/dflash2_grouped_conv.h b/src/ops/launcher/dflash2_grouped_conv.h new file mode 100644 index 0000000000..4b9c9a28be --- /dev/null +++ b/src/ops/launcher/dflash2_grouped_conv.h @@ -0,0 +1,14 @@ +#pragma once + +#include "core/tensor.h" + +#include + +namespace ninfer::ops::detail { + +void dflash2_grouped_conv_launch(const Tensor& hidden, const Tensor& delta, const Weight& base, + std::int32_t block_size, std::int32_t group_size, + std::int32_t taps, std::int32_t side, Tensor& out, + cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/dflash2_selector.cu b/src/ops/launcher/dflash2_selector.cu new file mode 100644 index 0000000000..ea19d055b0 --- /dev/null +++ b/src/ops/launcher/dflash2_selector.cu @@ -0,0 +1,48 @@ +#include "ops/launcher/dflash2_selector.h" + +#include "core/device.h" +#include "ops/kernel/dflash2_selector.cuh" + +#include +#include + +namespace ninfer::ops::detail { + +void dflash2_selector_launch(const Tensor& unary_logits, const Tensor& projected_hidden, + const Weight& predecessor_codebook, + const Weight& successor_codebook, const Tensor& anchors, + Tensor& candidates, Tensor& unary, Tensor& scores, Tensor& drafts, + std::int32_t steps, std::int32_t top_k, cudaStream_t stream) { + if (steps != 7 || top_k != kDflash2SelectorTopK) { + throw std::invalid_argument("dflash2_selector: registered domain is S=7, K=16"); + } + const std::int32_t batch = anchors.ne[0]; + const std::int32_t columns = steps * batch; + constexpr int kTopK = kDflash2SelectorTopK; + + const dim3 topk_grid(columns); + dflash2_selector_topk_kernel<256, kTopK><<>>( + static_cast(unary_logits.data), + static_cast(candidates.data), static_cast(unary.data), + unary_logits.ne[0], batch, steps, columns); + CUDA_CHECK(cudaGetLastError()); + + const dim3 scores_grid(columns); + dflash2_selector_scores_kernel<256><<>>( + static_cast(candidates.data), + static_cast(unary.data), + static_cast(projected_hidden.data), + static_cast(predecessor_codebook.qdata), + static_cast(successor_codebook.qdata), + static_cast(anchors.data), static_cast(scores.data), + unary_logits.ne[0], batch, steps, top_k, columns); + CUDA_CHECK(cudaGetLastError()); + + dflash2_selector_walk_kernel<<>>( + static_cast(candidates.data), + static_cast(scores.data), static_cast(drafts.data), batch, + steps, top_k); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/dflash2_selector.h b/src/ops/launcher/dflash2_selector.h new file mode 100644 index 0000000000..72bdae7cf1 --- /dev/null +++ b/src/ops/launcher/dflash2_selector.h @@ -0,0 +1,20 @@ +#pragma once + +#include "core/tensor.h" + +#include + +#include + +namespace ninfer::ops::detail { + +inline constexpr int kDflash2SelectorRank = 256; +inline constexpr int kDflash2SelectorTopK = 16; + +void dflash2_selector_launch(const Tensor& unary_logits, const Tensor& projected_hidden, + const Weight& predecessor_codebook, + const Weight& successor_codebook, const Tensor& anchors, + Tensor& candidates, Tensor& unary, Tensor& scores, Tensor& drafts, + std::int32_t steps, std::int32_t top_k, cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/swa.cu b/src/ops/launcher/swa.cu new file mode 100644 index 0000000000..96beef1c0a --- /dev/null +++ b/src/ops/launcher/swa.cu @@ -0,0 +1,166 @@ +#include "ops/launcher/swa.h" + +#include "core/device.h" +#include "ops/kernel/bidirectional_gqa_attention.cuh" + +#include +#include +#include + +namespace ninfer::ops::detail { +namespace { + +template +void dispatch_token_case(Launch&& launch) { + constexpr int Warps = (Tokens + 3) / 4; + launch.template operator()(); +} + +template +void dispatch_tokens(std::int32_t tokens, Launch&& launch) { + switch (tokens) { +#define NINFER_SWA_TOKEN_CASE(TOKENS) \ + case TOKENS: \ + dispatch_token_case(launch); \ + return + NINFER_SWA_TOKEN_CASE(1); + NINFER_SWA_TOKEN_CASE(2); + NINFER_SWA_TOKEN_CASE(3); + NINFER_SWA_TOKEN_CASE(4); + NINFER_SWA_TOKEN_CASE(5); + NINFER_SWA_TOKEN_CASE(6); + NINFER_SWA_TOKEN_CASE(7); + NINFER_SWA_TOKEN_CASE(8); + NINFER_SWA_TOKEN_CASE(9); + NINFER_SWA_TOKEN_CASE(10); + NINFER_SWA_TOKEN_CASE(11); + NINFER_SWA_TOKEN_CASE(12); + NINFER_SWA_TOKEN_CASE(13); + NINFER_SWA_TOKEN_CASE(14); + NINFER_SWA_TOKEN_CASE(15); + NINFER_SWA_TOKEN_CASE(16); +#undef NINFER_SWA_TOKEN_CASE + default: + throw std::invalid_argument("swa: unsupported T"); + } +} + +} // namespace + +SwaPlan swa_resolve_plan(std::int32_t tokens, SwaContextExecutionEnvelope envelope, + std::uint32_t window) { + if (tokens < 1 || tokens > 16) { throw std::invalid_argument("swa plan: T must be 1..16"); } + if (window != 2048 && window != 4096) { + throw std::invalid_argument("swa plan: registered windows are 2048 and 4096"); + } + if (envelope.min_context > envelope.max_context) { + throw std::invalid_argument("swa plan: invalid envelope"); + } + // Graph envelopes whose longest context fits three key tiles avoid the second kernel and + // workspace round trip. At four tiles, split-KV is already faster for every qualified T. + constexpr std::uint32_t direct_context_limit = 96; + const bool direct = envelope.max_context <= direct_context_limit; + constexpr std::int32_t key_block = 32; + const std::uint32_t context_rows = std::min(envelope.max_context, window - 1u); + const std::int32_t context_tiles = + static_cast((context_rows + key_block - 1u) / key_block); + constexpr std::int32_t split_limit = 32; + return { + .route = direct ? SwaRoute::Direct : SwaRoute::SplitKv, + .tokens = tokens, + .warps = (tokens + 3) / 4, + .split_capacity = direct ? 1 : std::min(split_limit, std::max(1, context_tiles)), + .max_context = static_cast(envelope.max_context), + .window = window, + }; +} + +const char* swa_route_name(SwaRoute route) { + switch (route) { + case SwaRoute::Direct: + return "direct"; + case SwaRoute::SplitKv: + return "split_kv"; + } + return "unknown"; +} + +void swa_launch(const Tensor& q, const Tensor& query_k, const Tensor& query_v, + const Tensor& positions, const Tensor& valid_columns, const Tensor& lanes, + float scale, const CyclicKVCacheLayerView& context, const SwaPlan& plan, + Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, + cudaStream_t stream) { + dispatch_tokens(q.ne[2], [&]() { + const bool direct = plan.route == SwaRoute::Direct; + if (plan.warps != Warps || plan.split_capacity < 1 || + plan.split_capacity > kSwaMaxCandidateSplit || (direct && plan.split_capacity != 1)) { + throw std::invalid_argument("swa: inconsistent plan"); + } + const auto launch_case = [&]() { + constexpr int KeyBlock = 32; + constexpr std::size_t SmemBytes = + 2u * KeyBlock * kBidirectionalGqaHeadDim * sizeof(__nv_bfloat16); + if (direct) { + const dim3 direct_grid(kBidirectionalGqaKVHeads, 1, q.ne[3]); + swa_split_partial_kernel + <<>>( + static_cast(q.data), + static_cast(query_k.data), + static_cast(query_v.data), + static_cast(positions.data), + static_cast(valid_columns.data), + static_cast(lanes.data), + static_cast(context.k.data), + static_cast(context.v.data), + static_cast(context.padded_capacity), plan.max_context, 1, scale, + static_cast<__nv_bfloat16*>(partial_acc.data), + static_cast(partial_m.data), static_cast(partial_l.data), + static_cast<__nv_bfloat16*>(out.data)); + CUDA_CHECK(cudaGetLastError()); + return; + } + + const dim3 partial_grid(kBidirectionalGqaKVHeads, plan.split_capacity, q.ne[3]); + swa_split_partial_kernel + <<>>( + static_cast(q.data), + static_cast(query_k.data), + static_cast(query_v.data), + static_cast(positions.data), + static_cast(valid_columns.data), + static_cast(lanes.data), + static_cast(context.k.data), + static_cast(context.v.data), + static_cast(context.padded_capacity), plan.max_context, + plan.split_capacity, scale, static_cast<__nv_bfloat16*>(partial_acc.data), + static_cast(partial_m.data), static_cast(partial_l.data), + static_cast<__nv_bfloat16*>(out.data)); + CUDA_CHECK(cudaGetLastError()); + + constexpr int ReduceWarps = 1; + constexpr int ReduceRows = kBidirectionalGqaQHeads * Tokens; + const dim3 reduce_grid((ReduceRows + ReduceWarps - 1) / ReduceWarps, 1, q.ne[3]); + swa_reduce_kernel + <<>>( + static_cast(partial_acc.data), + static_cast(partial_m.data), + static_cast(partial_l.data), + static_cast(positions.data), + static_cast(valid_columns.data), plan.max_context, + plan.split_capacity, static_cast<__nv_bfloat16*>(out.data)); + CUDA_CHECK(cudaGetLastError()); + }; + switch (plan.window) { + case 2048: + launch_case.template operator()<2048>(); + return; + case 4096: + launch_case.template operator()<4096>(); + return; + default: + throw std::invalid_argument("swa: unsupported window"); + } + }); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/launcher/swa.h b/src/ops/launcher/swa.h new file mode 100644 index 0000000000..d413312d92 --- /dev/null +++ b/src/ops/launcher/swa.h @@ -0,0 +1,33 @@ +#pragma once + +#include "ninfer/ops/swa.h" + +namespace ninfer::ops::detail { + +inline constexpr std::int32_t kSwaMaxCandidateSplit = 32; + +enum class SwaRoute { + Direct, + SplitKv, +}; + +struct SwaPlan { + SwaRoute route; + std::int32_t tokens; + std::int32_t warps; + std::int32_t split_capacity; + std::int32_t max_context; + std::uint32_t window; +}; + +[[nodiscard]] SwaPlan swa_resolve_plan(std::int32_t tokens, SwaContextExecutionEnvelope envelope, + std::uint32_t window); +[[nodiscard]] const char* swa_route_name(SwaRoute route); + +void swa_launch(const Tensor& q, const Tensor& query_k, const Tensor& query_v, + const Tensor& positions, const Tensor& valid_columns, const Tensor& lanes, + float scale, const CyclicKVCacheLayerView& context, const SwaPlan& plan, + Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, + cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/wrapper/bidirectional_gqa_attention.cpp b/src/ops/wrapper/bidirectional_gqa_attention.cpp new file mode 100644 index 0000000000..335e1c8ca7 --- /dev/null +++ b/src/ops/wrapper/bidirectional_gqa_attention.cpp @@ -0,0 +1,168 @@ +#include "ninfer/ops/bidirectional_gqa_attention.h" + +#include "core/layout.h" +#include "ops/launcher/bidirectional_gqa_attention.h" + +#include +#include +#include +#include +#include +#include + +namespace ninfer::ops { +namespace { + +constexpr std::int32_t kHeadDim = 128; +constexpr std::int32_t kQHeads = 32; +constexpr std::int32_t kKVHeads = 8; +constexpr float kExpectedScale = 0.08838834764831844055f; + +void require_shape(const Tensor& tensor, std::int32_t n0, std::int32_t n1, std::int32_t n2, + std::int32_t n3, const char* op, const char* name) { + if (tensor.ne[0] != n0 || tensor.ne[1] != n1 || tensor.ne[2] != n2 || tensor.ne[3] != n3) { + throw std::invalid_argument(std::string(op) + ": invalid shape for " + name); + } +} + +void require_contiguous_nonnull(const Tensor& tensor, const char* op, const char* name) { + if (!tensor.is_contiguous()) { + throw std::invalid_argument(std::string(op) + ": " + name + " must be contiguous"); + } + if (tensor.data == nullptr) { + throw std::invalid_argument(std::string(op) + ": " + name + " data must be non-null"); + } +} + +std::uint32_t validate_context(const PagedKVBatchLayerView& context, const char* op) { + if (context.dtype != DType::BF16 || context.quant_group != 0 || + context.num_kv_heads != kKVHeads || context.head_dim != kHeadDim) { + throw std::invalid_argument(std::string(op) + ": invalid context geometry or dtype"); + } + const std::int32_t physical_pages = context.k_pages.ne[2]; + if (physical_pages <= 0 || context.v_pages.ne[2] != physical_pages || + context.block_tables.ne[0] <= 0 || context.block_tables.ne[1] <= 0) { + throw std::invalid_argument(std::string(op) + ": invalid context capacity"); + } + if (context.k_pages.dtype != DType::BF16 || context.v_pages.dtype != DType::BF16) { + throw std::invalid_argument(std::string(op) + ": context K/V must be BF16"); + } + require_shape(context.k_pages, kHeadDim, kPagedKVPageSize, physical_pages, kKVHeads, op, + "context k pages"); + require_shape(context.v_pages, kHeadDim, kPagedKVPageSize, physical_pages, kKVHeads, op, + "context v pages"); + require_contiguous_nonnull(context.k_pages, op, "context k pages"); + require_contiguous_nonnull(context.v_pages, op, "context v pages"); + require_shape(context.block_tables, context.block_tables.ne[0], context.block_tables.ne[1], 1, + 1, op, "context block tables"); + if (context.block_tables.dtype != DType::I32) { + throw std::invalid_argument(std::string(op) + ": context block tables must be I32"); + } + require_contiguous_nonnull(context.block_tables, op, "context block tables"); + if (context.k_scale_pages.data != nullptr || context.v_scale_pages.data != nullptr) { + throw std::invalid_argument(std::string(op) + ": BF16 context must not have scales"); + } + const std::uint64_t logical_capacity = + static_cast(context.block_tables.ne[0]) * kPagedKVPageSize; + if (logical_capacity > static_cast(std::numeric_limits::max())) { + throw std::overflow_error(std::string(op) + ": logical context capacity exceeds int32"); + } + return static_cast(logical_capacity); +} + +struct PartialWorkspace { + Tensor acc; + Tensor m; + Tensor l; +}; + +template +PartialWorkspace allocate_workspace(Allocator& workspace, std::int32_t tokens, std::int32_t splits, + std::int32_t batch_size) { + return { + workspace.alloc(DType::BF16, {kHeadDim, kQHeads, tokens, splits * batch_size}), + workspace.alloc(DType::FP32, {kQHeads, tokens, splits * batch_size}), + workspace.alloc(DType::FP32, {kQHeads, tokens, splits * batch_size}), + }; +} + +} // namespace + +std::size_t bidirectional_gqa_attention_workspace_capacity_bytes( + GqaContextExecutionEnvelope envelope, std::int32_t min_tokens, std::int32_t max_tokens, + std::int32_t batch_size) { + if (min_tokens < 1 || max_tokens < min_tokens || max_tokens > 16 || batch_size < 1 || + batch_size > 8 || envelope.min_context > envelope.max_context || + envelope.max_context > + static_cast(std::numeric_limits::max())) { + throw std::invalid_argument( + "bidirectional_gqa_attention workspace: invalid envelope or token interval"); + } + const auto endpoint_capacity = [&](std::int32_t tokens) { + const auto plan = detail::bidirectional_gqa_resolve_plan(tokens, envelope); + WorkspaceLayoutBuilder layout; + (void)allocate_workspace(layout, tokens, plan.split_capacity, batch_size); + return layout.peak_bytes(1); + }; + + std::size_t maximum = 0; + if (min_tokens <= 8) { maximum = endpoint_capacity(std::min(max_tokens, 8)); } + if (max_tokens >= 9) { maximum = std::max(maximum, endpoint_capacity(max_tokens)); } + return maximum; +} + +void bidirectional_gqa_attention(const Tensor& q, const Tensor& query_k, const Tensor& query_v, + const Tensor& context_lengths, const Tensor& valid_columns, + const Tensor& table_rows, float scale, + const PagedKVBatchLayerView& context, + GqaContextExecutionEnvelope envelope, WorkspaceArena& workspace, + Tensor& out, cudaStream_t stream) { + constexpr const char* op = "bidirectional_gqa_attention"; + if (q.dtype != DType::BF16 || query_k.dtype != DType::BF16 || query_v.dtype != DType::BF16 || + out.dtype != DType::BF16) { + throw std::invalid_argument("bidirectional_gqa_attention: q/k/v/out must be BF16"); + } + if (context_lengths.dtype != DType::I32 || valid_columns.dtype != DType::I32 || + table_rows.dtype != DType::I32) { + throw std::invalid_argument( + "bidirectional_gqa_attention: lengths/valid_columns/table_rows must be I32"); + } + const std::int32_t tokens = q.ne[2]; + const std::int32_t batch = q.ne[3]; + if (tokens < 1 || tokens > 16) { + throw std::invalid_argument("bidirectional_gqa_attention: optimized domain is T=1..16"); + } + if (batch < 1 || batch > 8) { + throw std::invalid_argument("bidirectional_gqa_attention: B must be 1..8"); + } + require_shape(q, kHeadDim, kQHeads, tokens, batch, op, "q"); + require_shape(query_k, kHeadDim, kKVHeads, tokens, batch, op, "query k"); + require_shape(query_v, kHeadDim, kKVHeads, tokens, batch, op, "query v"); + require_shape(context_lengths, batch, 1, 1, 1, op, "context lengths"); + require_shape(valid_columns, batch, 1, 1, 1, op, "valid columns"); + require_shape(table_rows, batch, 1, 1, 1, op, "table rows"); + require_shape(out, kHeadDim, kQHeads, tokens, batch, op, "out"); + require_contiguous_nonnull(q, op, "q"); + require_contiguous_nonnull(query_k, op, "query k"); + require_contiguous_nonnull(query_v, op, "query v"); + require_contiguous_nonnull(context_lengths, op, "context lengths"); + require_contiguous_nonnull(valid_columns, op, "valid columns"); + require_contiguous_nonnull(table_rows, op, "table rows"); + require_contiguous_nonnull(out, op, "out"); + const std::uint32_t logical_capacity = validate_context(context, op); + if (envelope.min_context > envelope.max_context || envelope.max_context > logical_capacity) { + throw std::invalid_argument("bidirectional_gqa_attention: invalid execution envelope"); + } + if (!std::isfinite(scale) || std::abs(scale - kExpectedScale) > 1e-7f) { + throw std::invalid_argument("bidirectional_gqa_attention: scale must be 1/sqrt(128)"); + } + + auto scope = workspace.scope(); + const auto plan = detail::bidirectional_gqa_resolve_plan(tokens, envelope); + PartialWorkspace partial = allocate_workspace(workspace, tokens, plan.split_capacity, batch); + detail::bidirectional_gqa_attention_launch(q, query_k, query_v, context_lengths, valid_columns, + table_rows, scale, context, plan, partial.acc, + partial.m, partial.l, out, stream); +} + +} // namespace ninfer::ops diff --git a/src/ops/wrapper/dflash2_grouped_conv.cpp b/src/ops/wrapper/dflash2_grouped_conv.cpp new file mode 100644 index 0000000000..fb900e4461 --- /dev/null +++ b/src/ops/wrapper/dflash2_grouped_conv.cpp @@ -0,0 +1,45 @@ +#include "ninfer/ops/dflash2_grouped_conv.h" + +#include "ops/launcher/dflash2_grouped_conv.h" + +#include +#include + +namespace ninfer::ops { + +void dflash2_grouped_conv(const Tensor& hidden, const Tensor& delta, const Weight& base, + std::int32_t block_size, std::int32_t group_size, std::int32_t taps, + std::int32_t side, Tensor& out, cudaStream_t stream) { + if (hidden.dtype != DType::BF16 || delta.dtype != DType::BF16 || out.dtype != DType::BF16) { + throw std::invalid_argument("dflash2_grouped_conv: hidden/delta/out must be BF16"); + } + if (!hidden.is_contiguous() || hidden.data == nullptr || !delta.is_contiguous() || + delta.data == nullptr || !out.is_contiguous() || out.data == nullptr) { + throw std::invalid_argument("dflash2_grouped_conv: tensors must be contiguous and non-null"); + } + const std::int32_t hidden_size = hidden.ne[0]; + const std::int32_t tokens = hidden.ne[1]; + if (hidden_size != 5120 || tokens < 1 || tokens > 64 || block_size != 8 || + group_size != 16 || taps != 2 || side < 0 || side > 1 || + (block_size & (block_size - 1)) != 0) { + throw std::invalid_argument( + "dflash2_grouped_conv: registered domain is H=5120, T=1..64, taps=2, G=16, " + "block=8, side in {0,1}"); + } + if (hidden_size % group_size != 0 || out.ne[0] != hidden_size || out.ne[1] != tokens) { + throw std::invalid_argument("dflash2_grouped_conv: invalid hidden/out shapes"); + } + const std::int32_t groups = hidden_size / group_size; + if (delta.ne[0] != 2 * taps * groups || delta.ne[1] != tokens) { + throw std::invalid_argument("dflash2_grouped_conv: invalid delta shape"); + } + if (base.qtype != QType::BF16_CTRL || base.n != taps || base.k != hidden_size || + base.qdata == nullptr) { + throw std::invalid_argument("dflash2_grouped_conv: base must be BF16 [taps,H]"); + } + + detail::dflash2_grouped_conv_launch(hidden, delta, base, block_size, group_size, taps, side, + out, stream); +} + +} // namespace ninfer::ops diff --git a/src/ops/wrapper/dflash2_selector.cpp b/src/ops/wrapper/dflash2_selector.cpp new file mode 100644 index 0000000000..0c4d4aa4ea --- /dev/null +++ b/src/ops/wrapper/dflash2_selector.cpp @@ -0,0 +1,57 @@ +#include "ninfer/ops/dflash2_selector.h" + +#include "ops/launcher/dflash2_selector.h" + +#include +#include + +namespace ninfer::ops { + +void dflash2_selector(const Tensor& unary_logits, const Tensor& projected_hidden, + const Weight& predecessor_codebook, const Weight& successor_codebook, + const Tensor& anchors, Tensor& candidates, Tensor& unary, Tensor& scores, + Tensor& drafts, std::int32_t steps, std::int32_t top_k, cudaStream_t stream) { + const auto invalid = [](const char* message) { throw std::invalid_argument(message); }; + if (unary_logits.dtype != DType::BF16 || projected_hidden.dtype != DType::BF16 || + anchors.dtype != DType::I32 || candidates.dtype != DType::I32 || + unary.dtype != DType::FP32 || scores.dtype != DType::FP32 || drafts.dtype != DType::I32) { + invalid("dflash2_selector: logits/projected must be BF16, ids I32, scores F32"); + } + const auto contiguous = [](const Tensor& tensor) { + return tensor.is_contiguous() && tensor.data != nullptr; + }; + if (!contiguous(unary_logits) || !contiguous(projected_hidden) || !contiguous(anchors) || + !contiguous(candidates) || !contiguous(unary) || !contiguous(scores) || + !contiguous(drafts)) { + invalid("dflash2_selector: tensors must be contiguous and non-null"); + } + const std::int32_t vocab = unary_logits.ne[0]; + const std::int32_t tokens = unary_logits.ne[1]; + const std::int32_t rank = projected_hidden.ne[0]; + const std::int32_t batch = anchors.ne[0]; + if (vocab != 248320 || rank != detail::kDflash2SelectorRank || steps != 7 || + top_k != detail::kDflash2SelectorTopK || batch < 1 || batch > 8 || + tokens != steps * batch || + projected_hidden.ne[1] != tokens) { + invalid("dflash2_selector: registered domain is V=248320, R=256, S=7, B=1..8, K=16"); + } + if (candidates.ne[0] != batch || candidates.ne[1] != steps || candidates.ne[2] != top_k || + unary.ne[0] != batch || unary.ne[1] != steps || unary.ne[2] != top_k || + scores.ne[0] != batch || scores.ne[1] != steps || scores.ne[2] != top_k || + scores.ne[3] != top_k || drafts.ne[0] != tokens) { + invalid("dflash2_selector: scratch or draft shapes are inconsistent"); + } + if (predecessor_codebook.qtype != QType::BF16_CTRL || + successor_codebook.qtype != QType::BF16_CTRL || predecessor_codebook.n != vocab || + predecessor_codebook.k != rank || successor_codebook.n != vocab || + successor_codebook.k != rank || predecessor_codebook.qdata == nullptr || + successor_codebook.qdata == nullptr) { + invalid("dflash2_selector: codebooks must be BF16 [V,256]"); + } + + detail::dflash2_selector_launch(unary_logits, projected_hidden, predecessor_codebook, + successor_codebook, anchors, candidates, unary, scores, drafts, + steps, top_k, stream); +} + +} // namespace ninfer::ops diff --git a/src/ops/wrapper/swa.cpp b/src/ops/wrapper/swa.cpp new file mode 100644 index 0000000000..41df55fd28 --- /dev/null +++ b/src/ops/wrapper/swa.cpp @@ -0,0 +1,144 @@ +#include "ninfer/ops/swa.h" + +#include "core/layout.h" +#include "ops/launcher/swa.h" + +#include +#include +#include +#include +#include +#include + +namespace ninfer::ops { +namespace { + +constexpr std::int32_t kHeadDim = 128; +constexpr std::int32_t kQHeads = 32; +constexpr std::int32_t kKVHeads = 8; +constexpr float kExpectedScale = 0.08838834764831844055f; + +void require_shape(const Tensor& tensor, std::int32_t n0, std::int32_t n1, std::int32_t n2, + std::int32_t n3, const char* op, const char* name) { + if (tensor.ne[0] != n0 || tensor.ne[1] != n1 || tensor.ne[2] != n2 || tensor.ne[3] != n3) { + throw std::invalid_argument(std::string(op) + ": invalid shape for " + name); + } +} + +void require_contiguous_nonnull(const Tensor& tensor, const char* op, const char* name) { + if (!tensor.is_contiguous()) { + throw std::invalid_argument(std::string(op) + ": " + name + " must be contiguous"); + } + if (tensor.data == nullptr) { + throw std::invalid_argument(std::string(op) + ": " + name + " data must be non-null"); + } +} + +void validate_context(const CyclicKVCacheLayerView& context, std::uint32_t window, + const char* op) { + if (context.num_kv_heads != kKVHeads || context.head_dim != kHeadDim || + context.capacity != window || context.padded_capacity < context.capacity || + context.lane_capacity <= 0) { + throw std::invalid_argument(std::string(op) + ": invalid cyclic context"); + } + if (context.padded_capacity > + static_cast(std::numeric_limits::max())) { + throw std::overflow_error(std::string(op) + ": padded capacity exceeds int32"); + } + const auto padded = static_cast(context.padded_capacity); + if (context.k.dtype != DType::BF16 || context.v.dtype != DType::BF16) { + throw std::invalid_argument(std::string(op) + ": context K/V must be BF16"); + } + require_shape(context.k, kHeadDim, padded, kKVHeads, context.lane_capacity, op, "context k"); + require_shape(context.v, kHeadDim, padded, kKVHeads, context.lane_capacity, op, "context v"); + require_contiguous_nonnull(context.k, op, "context k"); + require_contiguous_nonnull(context.v, op, "context v"); +} + +struct PartialWorkspace { + Tensor acc; + Tensor m; + Tensor l; +}; + +template +PartialWorkspace allocate_workspace(Allocator& workspace, std::int32_t tokens, std::int32_t splits, + std::int32_t batch_size) { + return { + workspace.alloc(DType::BF16, {kHeadDim, kQHeads, tokens, splits * batch_size}), + workspace.alloc(DType::FP32, {kQHeads, tokens, splits * batch_size}), + workspace.alloc(DType::FP32, {kQHeads, tokens, splits * batch_size}), + }; +} + +} // namespace + +std::size_t swa_workspace_capacity_bytes(SwaContextExecutionEnvelope envelope, + std::int32_t min_tokens, std::int32_t max_tokens, + std::int32_t batch_size, std::uint32_t window) { + if (min_tokens < 1 || max_tokens < min_tokens || max_tokens > 16 || batch_size < 1 || + batch_size > 8 || envelope.min_context > envelope.max_context || + envelope.max_context > + static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("swa workspace: invalid envelope or token interval"); + } + const auto plan = detail::swa_resolve_plan(max_tokens, envelope, window); + WorkspaceLayoutBuilder layout; + (void)allocate_workspace(layout, max_tokens, plan.split_capacity, batch_size); + return layout.peak_bytes(1); +} + +void swa(const Tensor& q, const Tensor& query_k, const Tensor& query_v, const Tensor& positions, + const Tensor& valid_columns, const Tensor& lanes, float scale, + const CyclicKVCacheLayerView& context, SwaContextExecutionEnvelope envelope, + std::uint32_t window, WorkspaceArena& workspace, Tensor& out, cudaStream_t stream) { + constexpr const char* op = "swa"; + if (window != 2048 && window != 4096) { + throw std::invalid_argument("swa: registered windows are 2048 and 4096"); + } + if (q.dtype != DType::BF16 || query_k.dtype != DType::BF16 || query_v.dtype != DType::BF16 || + out.dtype != DType::BF16) { + throw std::invalid_argument("swa: q/k/v/out must be BF16"); + } + if (positions.dtype != DType::I32 || valid_columns.dtype != DType::I32 || + lanes.dtype != DType::I32) { + throw std::invalid_argument("swa: positions/valid_columns/lanes must be I32"); + } + const std::int32_t tokens = q.ne[2]; + const std::int32_t batch = q.ne[3]; + if (tokens < 1 || tokens > 16) { + throw std::invalid_argument("swa: optimized domain is T=1..16"); + } + if (batch < 1 || batch > 8) { throw std::invalid_argument("swa: B must be 1..8"); } + require_shape(q, kHeadDim, kQHeads, tokens, batch, op, "q"); + require_shape(query_k, kHeadDim, kKVHeads, tokens, batch, op, "query k"); + require_shape(query_v, kHeadDim, kKVHeads, tokens, batch, op, "query v"); + require_shape(positions, tokens, batch, 1, 1, op, "positions"); + require_shape(valid_columns, batch, 1, 1, 1, op, "valid columns"); + require_shape(lanes, batch, 1, 1, 1, op, "lanes"); + require_shape(out, kHeadDim, kQHeads, tokens, batch, op, "out"); + require_contiguous_nonnull(q, op, "q"); + require_contiguous_nonnull(query_k, op, "query k"); + require_contiguous_nonnull(query_v, op, "query v"); + require_contiguous_nonnull(positions, op, "positions"); + require_contiguous_nonnull(valid_columns, op, "valid columns"); + require_contiguous_nonnull(lanes, op, "lanes"); + require_contiguous_nonnull(out, op, "out"); + validate_context(context, window, op); + if (envelope.min_context > envelope.max_context || + envelope.max_context > + static_cast(std::numeric_limits::max())) { + throw std::invalid_argument("swa: invalid execution envelope"); + } + if (!std::isfinite(scale) || std::abs(scale - kExpectedScale) > 1e-7f) { + throw std::invalid_argument("swa: scale must be 1/sqrt(128)"); + } + + auto scope = workspace.scope(); + const auto plan = detail::swa_resolve_plan(tokens, envelope, window); + PartialWorkspace partial = allocate_workspace(workspace, tokens, plan.split_capacity, batch); + detail::swa_launch(q, query_k, query_v, positions, valid_columns, lanes, scale, context, plan, + partial.acc, partial.m, partial.l, out, stream); +} + +} // namespace ninfer::ops diff --git a/src/product/speculative_options.h b/src/product/speculative_options.h index aa307d21b6..946956761a 100644 --- a/src/product/speculative_options.h +++ b/src/product/speculative_options.h @@ -12,6 +12,7 @@ namespace ninfer::product { if (value == "mtp") { return SpeculativeBackend::Mtp; } if (value == "dflash") { return SpeculativeBackend::DFlash; } if (value == "dflash2") { return SpeculativeBackend::DFlash2; } + if (value == "auto") { return SpeculativeBackend::Auto; } throw std::invalid_argument("invalid speculative backend: " + std::string(value)); } @@ -25,6 +26,8 @@ namespace ninfer::product { return "dflash"; case SpeculativeBackend::DFlash2: return "dflash2"; + case SpeculativeBackend::Auto: + return "auto"; } return "unknown"; } @@ -48,8 +51,17 @@ inline void validate_speculative_cli_options(const SpeculativeOptions& options) } return; case SpeculativeBackend::DFlash2: - if (options.draft_tokens == 0 || options.draft_tokens > 15) { - throw std::invalid_argument("--spec dflash2 requires --draft-tokens in [1,15]"); + if (options.draft_tokens != 0 && options.draft_tokens != 7) { + throw std::invalid_argument("--spec dflash2 uses the fixed 7-draft block"); + } + if (options.proposal_head != ProposalHead::Full) { + throw std::invalid_argument("--spec dflash2 requires the full proposal head"); + } + return; + case SpeculativeBackend::Auto: + if (options.draft_tokens != 0 || options.proposal_head != ProposalHead::Full) { + throw std::invalid_argument( + "--draft-tokens and --lm-head-draft require --spec mtp|dflash|dflash2"); } return; } diff --git a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/model_view.h b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/model_view.h index e1d27e463b..919358a9a8 100644 --- a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/model_view.h +++ b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/model_view.h @@ -78,14 +78,43 @@ struct DFlashWeights { Tensor final_norm; }; +struct DFlash2LayerWeights { + Tensor input_norm; + Weight query_key_value; + Weight context_key; + Weight context_value; + Tensor query_norm; + Tensor key_norm; + Weight attention_output; + Weight attention_conv_base; + Weight attention_conv_projection; + Tensor post_attention_norm; + Weight gate_up; + Weight down; + Weight mlp_conv_base; + Weight mlp_conv_projection; +}; + +template +struct DFlash2Weights { + Weight feature_projection; + Tensor context_norm; + std::array layers; + Tensor final_norm; + Weight selector_hidden_projection; + Weight selector_predecessor_codebook; + Weight selector_successor_codebook; +}; + template + class DFlash2Payload, std::size_t FullAttentionLayers, std::size_t GdnLayers> struct ModelView { using FullLayer = FullAttentionWeights; using GdnLayer = GdnWeights; using MtpLayer = MtpWeights; using DFlash = DFlashPayload; + using DFlash2 = DFlash2Payload; DeviceArena* weights_arena = nullptr; Weight token_embedding; @@ -97,6 +126,7 @@ struct ModelView { std::optional optimized_proposal; std::optional mtp; std::optional dflash; + std::optional dflash2; std::optional vision; }; diff --git a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/round_state.h b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/round_state.h index 47b205593f..c3e58c85d6 100644 --- a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/round_state.h +++ b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/round_state.h @@ -89,6 +89,7 @@ struct DFlashDecodeEgress { std::array licensed_tokens{}; std::array licensed_counts{}; std::array accepted_drafts{}; + std::array proposal_extents{}; }; struct OrdinaryDecodeStateLayout { @@ -263,6 +264,7 @@ struct DFlashDecodeState { Tensor licensed_tokens; Tensor licensed_counts; Tensor accepted_drafts; + Tensor egress_proposal_extents; Tensor proposal_ids; Tensor proposal_positions; Tensor append_positions; diff --git a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/startup_features.h b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/startup_features.h index 67eae7f6ea..d5f54ec287 100644 --- a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/startup_features.h +++ b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/startup_features.h @@ -19,6 +19,12 @@ struct StartupFeatures { [[nodiscard]] bool dflash() const noexcept { return speculative == SpeculativeBackend::DFlash; } + [[nodiscard]] bool dflash2() const noexcept { + return speculative == SpeculativeBackend::DFlash2; + } + + [[nodiscard]] bool dflash_like() const noexcept { return dflash() || dflash2(); } + [[nodiscard]] bool optimized_proposal() const noexcept { return speculative_enabled() && proposal_head == ProposalHead::Optimized; } diff --git a/src/targets/qwen3_6/impl/runtime/dflash2_impl.h b/src/targets/qwen3_6/impl/runtime/dflash2_impl.h new file mode 100644 index 0000000000..9ff20f574d --- /dev/null +++ b/src/targets/qwen3_6/impl/runtime/dflash2_impl.h @@ -0,0 +1,475 @@ +#include "targets/qwen3_6/impl/runtime/instance.h" +#include "targets/qwen3_6/impl/runtime/schedule.h" +#include "targets/qwen3_6/impl/runtime/workspace_recipe.h" + +#include "core/dtype.h" +#include "ninfer/ops/argmax.h" +#include "ninfer/ops/dflash2_grouped_conv.h" +#include "ninfer/ops/dflash2_selector.h" +#include "ninfer/ops/embedding.h" +#include "ninfer/ops/kv_cache_append.h" +#include "ninfer/ops/linear.h" +#include "ninfer/ops/prepare_masked_block.h" +#include "ninfer/ops/prepare_ragged_prefix.h" +#include "ninfer/ops/residual_add.h" +#include "ninfer/ops/rmsnorm.h" +#include "ninfer/ops/rope.h" +#include "ninfer/ops/scalar.h" +#include "ninfer/ops/silu_mul.h" +#include "ninfer/ops/speculative_round.h" +#include "ninfer/ops/swa.h" + +#include + +#include +#include +#include +#include + +namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule { +namespace { + +DFlash2PersistentState& dflash2_state(PrefillContext& state) { + if (state.dflash2 == nullptr) { + throw std::logic_error("DFlash2 schedule requires DFlash2 weights and state"); + } + return *state.dflash2; +} + +DFlash2PersistentState& dflash2_state(DFlash2BatchContext& state) { return state.dflash2; } + +DFlash2PersistentState& dflash2_state(DFlash2AppendContext& state) { return state.dflash2; } + +Weight conv_side_weight(const Weight& base, std::int32_t side, std::int32_t taps) { + Weight out = base; + out.qdata = static_cast(base.qdata) + + static_cast(side) * static_cast(taps) * + static_cast(base.k) * dtype_size(DType::BF16); + out.n = taps; + return out; +} + +template +DFlashFeatureSink dflash2_prefill_feature_sink_impl( + PrefillContext& state, DFlashFeatureSink::PrefillConsumer consume_prefill) { + if constexpr (!V::supports_dflash2) { + throw std::logic_error("DFlash2 feature capture is unavailable for this target"); + } else { + using Config = typename V::DFlash2Config; + return DFlashFeatureSink{ + .features = &dflash2_state(state).prefill_features, + .positions = &dflash2_state(state).prefill_positions, + .layers = std::span(Config::target_feature_layers), + .consume_prefill = std::move(consume_prefill), + }; + } +} + +template +DFlashFeatureSink dflash2_batch_feature_sink_impl(DFlash2BatchContext& state, const Tensor& lanes, + const Tensor& valid_columns, + std::int32_t width, std::int32_t batch_size) { + if constexpr (!V::supports_dflash2) { + throw std::logic_error("DFlash2 feature capture is unavailable for this target"); + } else { + using Config = typename V::DFlash2Config; + return DFlashFeatureSink{ + .batch_features = &dflash2_state(state).pending_features, + .batch_lanes = &lanes, + .batch_valid_columns = &valid_columns, + .batch_width = width, + .batch_size = batch_size, + .layers = std::span(Config::target_feature_layers), + }; + } +} + +void append_context_impl(DFlash2AppendContext& state, const Tensor& features, + const Tensor& positions, const Tensor& commit_counts, + const Tensor& lanes, const Tensor& table_rows, + ops::KVCacheAppendPrefixExecutionEnvelope envelope) { + using Config = DFlash2Config; + const std::int32_t width = features.ne[1]; + const std::int32_t batch = features.ne[2]; + const std::int32_t columns = width * batch; + if (width <= 0 || batch <= 0 || features.dtype != DType::BF16 || + features.ne[0] != Config::feature_rows || positions.dtype != DType::I32 || + positions.ne[0] != width || positions.ne[1] != batch || + commit_counts.dtype != DType::I32 || commit_counts.ne[0] != batch || + lanes.dtype != DType::I32 || lanes.ne[0] != batch) { + throw std::invalid_argument("DFlash2 context append inputs are invalid"); + } + if (!state.execution.model.dflash2.has_value()) { + throw std::logic_error("DFlash2 weights are unavailable"); + } + const auto& dflash2 = *state.execution.model.dflash2; + + // The cyclic draft cache retains only the last `local_capacity` absolute + // positions. An oversized prefill chunk can therefore skip its older + // columns and commit just the trailing window. + const bool replace_local_window = batch == 1 && width > Config::local_capacity; + const int local_offset = replace_local_window ? width - Config::local_capacity : 0; + const int local_width = replace_local_window ? Config::local_capacity : width; + Tensor local_counts = commit_counts; + ops::KVCacheAppendPrefixExecutionEnvelope local_envelope = envelope; + if (replace_local_window) { + if (!state.execution.io.dflash_prefill) { + throw std::logic_error("DFlash2 prefill count storage is unavailable"); + } + local_counts = state.execution.io.dflash_prefill->produced_count; + ops::set_i32_scalar(local_counts, Config::local_capacity, + state.execution.device.stream); + local_envelope = {static_cast(Config::local_capacity), + static_cast(Config::local_capacity)}; + } + + auto context_roots = workspace_recipe::dflash_context(state.execution.work, columns); + ops::linear(features.view({Config::feature_rows, columns}), dflash2.feature_projection, + context_roots.projected, state.execution.device.stream); + Tensor context_full = context_roots.normalized; + ops::rmsnorm(context_roots.projected, dflash2.context_norm, Config::rms_epsilon, false, + context_full, state.execution.device.stream); + Tensor context = replace_local_window + ? context_full.slice(1, local_offset, local_width) + : context_full; + Tensor local_positions = replace_local_window + ? positions.slice(0, local_offset, local_width) + : positions; + + for (int layer = 0; layer < Config::layers; ++layer) { + auto layer_scope = state.execution.work.scope(); + const auto& weight = + dflash2.layers.at(static_cast(layer)); + auto layer_roots = + workspace_recipe::dflash_context_layer(state.execution.work, local_width * batch); + Tensor key_raw = + layer_roots.key_raw.view({Config::head_dim, Config::kv_heads, local_width * batch}); + Tensor value = layer_roots.value.view({Config::head_dim, Config::kv_heads, local_width * batch}); + Tensor key_flat = key_raw.view({Config::kv_size, local_width * batch}); + Tensor value_flat = value.view({Config::kv_size, local_width * batch}); + ops::linear(context, weight.context_key, key_flat, state.execution.device.stream); + ops::linear(context, weight.context_value, value_flat, state.execution.device.stream); + Tensor key = layer_roots.key.view({Config::head_dim, Config::kv_heads, local_width * batch}); + ops::rmsnorm(key_raw, weight.key_norm, Config::rms_epsilon, false, key, + state.execution.device.stream); + ops::rope(local_positions.view({local_width * batch}), Config::head_dim, + Config::rope_theta, key, state.execution.device.stream); + Tensor key_batch = key.view({Config::head_dim, Config::kv_heads, local_width, batch}); + Tensor value_batch = value.view({Config::head_dim, Config::kv_heads, local_width, batch}); + Tensor position_batch = local_positions.view({local_width, batch}); + ops::kv_cache_append_prefix(key_batch, value_batch, position_batch, local_counts, lanes, + local_envelope, + dflash2_state(state).local_layer( + static_cast(layer)), + Config::local_window, state.execution.device.stream); + } +} + +void propose_batch_impl(DFlash2BatchContext& state, qwen3_6::DFlashDecodeState& frame, + std::int32_t batch_size, std::uint32_t k, DFlash2Envelopes envelopes) { + using Config = DFlash2Config; + const std::int32_t width = static_cast(k) + 1; + const std::int32_t columns = width * batch_size; + Tensor anchors = frame.anchors.slice(0, 0, batch_size); + Tensor frontiers = frame.execution_frontiers.slice(0, 0, batch_size); + Tensor valid_columns = frame.target_valid_columns.slice(0, 0, batch_size); + Tensor lanes = frame.active_lanes.slice(0, 0, batch_size); + Tensor ids = frame.proposal_ids.slice(1, 0, batch_size); + Tensor positions = frame.proposal_positions.slice(1, 0, batch_size); + Tensor drafts = frame.draft_tokens.slice(1, 0, batch_size); + if (!state.execution.model.dflash2.has_value()) { + throw std::logic_error("DFlash2 weights are unavailable"); + } + const auto& dflash2 = *state.execution.model.dflash2; + + state.execution.work.reset(); + Tensor attention_valid = state.execution.work.alloc(DType::I32, {batch_size}); + ops::set_i32_scalar(attention_valid, static_cast(width), + state.execution.device.stream); + (void)valid_columns; + + ops::prepare_masked_block(anchors, frontiers, attention_valid, Config::mask_token, ids, + positions, state.execution.device.stream); + Tensor residual = state.execution.work.alloc(DType::BF16, {Config::hidden, columns}); + ops::embedding(ids.view({columns}), state.execution.model.token_embedding, residual, + state.execution.device.stream); + + for (int layer = 0; layer < Config::layers; ++layer) { + const auto& weight = dflash2.layers.at(static_cast(layer)); + { + auto attention_scope = state.execution.work.scope(); + auto roots = + workspace_recipe::dflash_attention(state.execution.work, columns); + ops::rmsnorm(residual, weight.input_norm, Config::rms_epsilon, false, roots.hidden, + state.execution.device.stream); + + Tensor conv_coefficients = state.execution.work.alloc(DType::BF16, {1280, columns}); + ops::linear(roots.hidden, weight.attention_conv_projection, conv_coefficients, + state.execution.device.stream); + Tensor conv_hidden = state.execution.work.alloc(DType::BF16, {Config::hidden, columns}); + ops::dflash2_grouped_conv( + roots.hidden, conv_coefficients, + conv_side_weight(weight.attention_conv_base, 0, Config::conv_kernel_size), + static_cast(k + 1), Config::conv_group_size, + Config::conv_kernel_size, 0, conv_hidden, state.execution.device.stream); + + Tensor query_raw = roots.query_raw.view({Config::head_dim, Config::query_heads, columns}); + Tensor key_raw = roots.key_raw.view({Config::head_dim, Config::kv_heads, columns}); + Tensor value = roots.value.view({Config::head_dim, Config::kv_heads, columns}); + Tensor query_flat = query_raw.view({Config::query_size, columns}); + Tensor key_flat = key_raw.view({Config::kv_size, columns}); + Tensor value_flat = value.view({Config::kv_size, columns}); + const Weight qkv_weight = weight.query_key_value; + const std::size_t row_len = static_cast(qkv_weight.k) * dtype_size(DType::BF16); + Weight query_weight = qkv_weight; + query_weight.n = Config::query_size; + Weight key_weight = qkv_weight; + key_weight.n = Config::kv_size; + key_weight.qdata = static_cast(qkv_weight.qdata) + + static_cast(Config::query_size) * row_len; + Weight value_weight = qkv_weight; + value_weight.n = Config::kv_size; + value_weight.qdata = static_cast(qkv_weight.qdata) + + static_cast(Config::query_size + Config::kv_size) * + row_len; + ops::linear(conv_hidden, query_weight, query_flat, state.execution.device.stream); + ops::linear(conv_hidden, key_weight, key_flat, state.execution.device.stream); + ops::linear(conv_hidden, value_weight, value_flat, state.execution.device.stream); + + Tensor query = roots.query.view({Config::head_dim, Config::query_heads, columns}); + Tensor key = roots.key.view({Config::head_dim, Config::kv_heads, columns}); + ops::rmsnorm(query_raw, weight.query_norm, Config::rms_epsilon, false, query, + state.execution.device.stream); + ops::rmsnorm(key_raw, weight.key_norm, Config::rms_epsilon, false, key, + state.execution.device.stream); + ops::rope(positions.view({columns}), Config::head_dim, Config::rope_theta, query, key, + state.execution.device.stream); + Tensor query_batch = + query.view({Config::head_dim, Config::query_heads, width, batch_size}); + Tensor key_batch = + key.view({Config::head_dim, Config::kv_heads, width, batch_size}); + Tensor value_batch = + value.view({Config::head_dim, Config::kv_heads, width, batch_size}); + Tensor attention_batch = roots.attention.view( + {Config::head_dim, Config::query_heads, width, batch_size}); + ops::swa(query_batch, key_batch, value_batch, positions, attention_valid, lanes, + Config::attention_scale, + dflash2_state(state).local_layer(static_cast(layer)), + envelopes.local, Config::local_window, state.execution.work, attention_batch, + state.execution.device.stream); + + Tensor delta = roots.attention_delta.view({Config::hidden, columns}); + ops::linear(roots.attention.view({Config::query_size, columns}), + weight.attention_output, delta, state.execution.device.stream); + Tensor conv_finish = state.execution.work.alloc(DType::BF16, {Config::hidden, columns}); + ops::dflash2_grouped_conv( + delta, conv_coefficients, + conv_side_weight(weight.attention_conv_base, 1, Config::conv_kernel_size), + static_cast(k + 1), Config::conv_group_size, + Config::conv_kernel_size, 1, conv_finish, state.execution.device.stream); + ops::residual_add(conv_finish, residual, state.execution.device.stream); + } + { + auto mlp_scope = state.execution.work.scope(); + auto roots = workspace_recipe::dflash_mlp(state.execution.work, columns); + ops::rmsnorm(residual, weight.post_attention_norm, Config::rms_epsilon, false, + roots.hidden, state.execution.device.stream); + + Tensor conv_coefficients = state.execution.work.alloc(DType::BF16, {1280, columns}); + ops::linear(roots.hidden, weight.mlp_conv_projection, conv_coefficients, + state.execution.device.stream); + Tensor conv_hidden = state.execution.work.alloc(DType::BF16, {Config::hidden, columns}); + ops::dflash2_grouped_conv( + roots.hidden, conv_coefficients, + conv_side_weight(weight.mlp_conv_base, 0, Config::conv_kernel_size), + static_cast(k + 1), Config::conv_group_size, + Config::conv_kernel_size, 0, conv_hidden, state.execution.device.stream); + + Tensor gate_up = roots.gate_up.view({2 * Config::intermediate, columns}); + ops::linear(conv_hidden, weight.gate_up, gate_up, state.execution.device.stream); + ops::silu_mul(gate_up.slice(0, 0, Config::intermediate), + gate_up.slice(0, Config::intermediate, Config::intermediate), + roots.intermediate, state.execution.device.stream); + Tensor delta = roots.delta.view({Config::hidden, columns}); + ops::linear(roots.intermediate, weight.down, delta, state.execution.device.stream); + Tensor conv_finish = state.execution.work.alloc(DType::BF16, {Config::hidden, columns}); + ops::dflash2_grouped_conv( + delta, conv_coefficients, + conv_side_weight(weight.mlp_conv_base, 1, Config::conv_kernel_size), + static_cast(k + 1), Config::conv_group_size, + Config::conv_kernel_size, 1, conv_finish, state.execution.device.stream); + ops::residual_add(conv_finish, residual, state.execution.device.stream); + } + } + + Tensor packed = state.execution.work.alloc( + DType::BF16, {Config::hidden, static_cast(k) * batch_size}); + const std::size_t element_bytes = dtype_size(DType::BF16); + const std::size_t row_bytes = + static_cast(Config::hidden) * static_cast(k) * element_bytes; + const std::size_t source_pitch = + static_cast(Config::hidden) * width * element_bytes; + // DFlash2 predicts the seven masked columns (1..7); the bonus token stays + // at column 0 and is not part of the proposal block. + const auto* source = static_cast(residual.data) + + static_cast(Config::hidden) * element_bytes; + CUDA_CHECK(cudaMemcpy2DAsync(packed.data, row_bytes, source, source_pitch, row_bytes, + static_cast(batch_size), cudaMemcpyDeviceToDevice, + state.execution.device.stream)); + Tensor proposal_hidden = state.execution.work.alloc( + DType::BF16, {Config::hidden, static_cast(k) * batch_size}); + ops::rmsnorm(packed, dflash2.final_norm, Config::rms_epsilon, false, proposal_hidden, + state.execution.device.stream); + + Tensor flat_drafts = drafts.view({static_cast(k) * batch_size}); + Tensor logits = state.execution.work.alloc( + DType::BF16, {TextConfig::output_rows, static_cast(k) * batch_size}); + ops::linear(proposal_hidden, state.execution.model.output_head, logits, + state.execution.device.stream); + Tensor projected = state.execution.work.alloc( + DType::BF16, {Config::selector_rank, static_cast(k) * batch_size}); + ops::linear(proposal_hidden, dflash2.selector_hidden_projection, projected, + state.execution.device.stream); + Tensor candidates = state.execution.work.alloc( + DType::I32, {batch_size, Config::block_drafts, Config::selector_top_k}); + Tensor unary = state.execution.work.alloc( + DType::FP32, {batch_size, Config::block_drafts, Config::selector_top_k}); + Tensor scores = state.execution.work.alloc( + DType::FP32, + {batch_size, Config::block_drafts, Config::selector_top_k, Config::selector_top_k}); + ops::dflash2_selector(logits, projected, dflash2.selector_predecessor_codebook, + dflash2.selector_successor_codebook, anchors, candidates, unary, scores, + flat_drafts, Config::block_drafts, Config::selector_top_k, + state.execution.device.stream); + state.execution.work.reset(); +} + +auto dflash2_decode_batch_body(DFlash2BatchContext& state, std::int32_t batch_size, + std::uint32_t k, DFlash2Envelopes envelopes, + ops::CausalAttentionExecutionEnvelope target_envelope) { + return [&state, batch_size, k, envelopes, target_envelope] { + if (batch_size <= 0 || batch_size > static_cast(kMaximumConcurrency) || + k == 0 || k > kDFlashDecodeMaximumDrafts || k != DFlash2Config::block_drafts) { + throw std::logic_error("DFlash2 decode batch state is incomplete"); + } + qwen3_6::DFlashDecodeState& frame = state.frame; + const std::int32_t width = static_cast(k) + 1; + CUDA_CHECK(cudaMemcpyAsync(frame.ingress.data, &state.host_ingress, + sizeof(qwen3_6::DFlashDecodeIngress), cudaMemcpyHostToDevice, + state.execution.device.stream)); + + Tensor anchors = frame.anchors.slice(0, 0, batch_size); + Tensor frontiers = frame.execution_frontiers.slice(0, 0, batch_size); + Tensor context_starts = frame.context_frontiers.slice(0, 0, batch_size); + Tensor extents = frame.proposal_extents.slice(0, 0, batch_size); + Tensor valid_columns = frame.target_valid_columns.slice(0, 0, batch_size); + Tensor text_rows = frame.text_kv_table_rows.slice(0, 0, batch_size); + Tensor lanes = frame.active_lanes.slice(0, 0, batch_size); + Tensor append_positions = frame.append_positions.slice(1, 0, batch_size); + Tensor append_counts = frame.append_counts.slice(0, 0, batch_size); + Tensor drafts = frame.draft_tokens.slice(1, 0, batch_size); + Tensor verify_ids = frame.verify_ids.slice(1, 0, batch_size); + Tensor target_positions = frame.proposal_positions.slice(1, 0, batch_size); + Tensor target_tokens = frame.target_argmax.slice(1, 0, batch_size); + Tensor target_logits = frame.target_logits.slice(2, 0, batch_size); + Tensor target_hidden = frame.target_hidden.slice(2, 0, batch_size); + Tensor selected_hidden = frame.target_continuation_hidden.slice(1, 0, batch_size); + Tensor licensed_tokens = frame.licensed_tokens.slice(1, 0, batch_size); + Tensor licensed_counts = frame.licensed_counts.slice(0, 0, batch_size); + Tensor accepted = frame.accepted_drafts.slice(0, 0, batch_size); + + state.execution.work.reset(); + Tensor compact_features = state.execution.work.alloc( + DType::BF16, {DFlash2Config::feature_rows, width, batch_size}); + ops::prepare_ragged_prefix(dflash2_state(state).pending_features, lanes, context_starts, + frontiers, compact_features, append_positions, append_counts, + state.execution.device.stream); + DFlash2AppendContext append_state{.execution = state.execution, + .dflash2 = state.dflash2}; + append_context_impl(append_state, compact_features, append_positions, append_counts, lanes, + text_rows, envelopes.append); + + propose_batch_impl(state, frame, batch_size, k, envelopes); + ops::speculative_prepare_verify_ids(anchors, drafts, extents, verify_ids, + state.execution.device.stream); + + TextContext card(state.execution.device, state.execution.model, state.execution.work, {}, + state.execution.linear_attention, state.execution.io, + state.execution.prefill_hidden, state.execution.prefill_chunk, 0, {}, + &state.text_cache); + DFlashFeatureSink sink = dflash2_batch_feature_sink_impl( + state, lanes, valid_columns, width, batch_size); + target_verify_accept(state.execution, state.continuation_hidden_store, card, + TargetVerifyFrameView{ + .ids = verify_ids, + .cache_positions = target_positions, + .rope_positions = target_positions, + .valid_columns = valid_columns, + .kv_table_rows = text_rows, + .target_hidden = target_hidden, + .target_logits = target_logits, + .target_tokens = target_tokens, + .drafts = drafts, + .current_extents = extents, + .frontiers = frontiers, + .anchors = anchors, + .licensed_tokens = licensed_tokens, + .licensed_counts = licensed_counts, + .accepted_drafts = accepted, + .selected_hidden = selected_hidden, + .replay_records = state.execution.replay_records, + .sampling = frame.sampling, + .feature_sink = &sink, + }, + target_envelope); + CUDA_CHECK(cudaMemcpyAsync( + frame.egress_proposal_extents.slice(0, 0, batch_size).data, extents.data, + static_cast(batch_size) * sizeof(std::int32_t), cudaMemcpyDeviceToDevice, + state.execution.device.stream)); + CUDA_CHECK(cudaMemcpyAsync(&state.host_egress, frame.egress.data, + sizeof(qwen3_6::DFlashDecodeEgress), cudaMemcpyDeviceToHost, + state.execution.device.stream)); + }; +} + +} // namespace + +DFlashFeatureSink dflash2_feature_sink(PrefillContext& state, + DFlashFeatureSink::PrefillConsumer consume_prefill) { + return dflash2_prefill_feature_sink_impl(state, std::move(consume_prefill)); +} + +void dflash2_append_context(DFlash2AppendContext& state, const Tensor& features, + const Tensor& positions, const Tensor& commit_counts, + const Tensor& lanes, const Tensor& table_rows, + ops::KVCacheAppendPrefixExecutionEnvelope envelope) { + append_context_impl(state, features, positions, commit_counts, lanes, table_rows, envelope); +} + +void dflash2_append_context(PrefillContext& state, const Tensor& features, const Tensor& positions, + const Tensor& commit_counts, const Tensor& lanes, + const Tensor& table_rows, + ops::KVCacheAppendPrefixExecutionEnvelope envelope) { + if (state.dflash2 == nullptr) { + throw std::logic_error("DFlash2 context append requires DFlash2 state"); + } + DFlash2AppendContext context{.execution = state.execution, .dflash2 = *state.dflash2}; + append_context_impl(context, features, positions, commit_counts, lanes, table_rows, envelope); +} + +void capture_dflash2_decode_batch(DFlash2BatchContext& state, std::int32_t batch_size, + std::uint32_t k, DFlash2Envelopes envelopes, + ops::CausalAttentionExecutionEnvelope target_envelope, + DecodeGraphDefinition& definition) { + auto body = dflash2_decode_batch_body(state, batch_size, k, envelopes, target_envelope); + capture_graph(state, definition, body); +} + +void dflash2_decode_batch(DFlash2BatchContext& state, std::int32_t batch_size, std::uint32_t k, + DFlash2Envelopes envelopes, ops::CausalAttentionExecutionEnvelope target_envelope, + DecodeGraphExecutable* executable) { + auto body = dflash2_decode_batch_body(state, batch_size, k, envelopes, target_envelope); + run_prepared(state, executable, body); +} + +} // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule diff --git a/src/targets/qwen3_6/impl/runtime/dflash_context.h b/src/targets/qwen3_6/impl/runtime/dflash_context.h index ecb1dd47ce..f393024b4e 100644 --- a/src/targets/qwen3_6/impl/runtime/dflash_context.h +++ b/src/targets/qwen3_6/impl/runtime/dflash_context.h @@ -26,4 +26,18 @@ struct DFlashPersistentState { cudaStream_t stream); }; +struct DFlash2PersistentState { + CyclicKVCache local; + CyclicKVCache rewrite_checkpoint_local; + Tensor prefill_features; + Tensor prefill_positions; + Tensor pending_features; + + DFlash2PersistentState(DeviceSpan backing, const DFlash2PersistentLayout& layout); + + [[nodiscard]] CyclicKVCacheLayerView local_layer(std::uint32_t layer) const; + void save_rewrite_checkpoint(std::int32_t lane, cudaStream_t stream); + void restore_rewrite_checkpoint(std::int32_t lane, cudaStream_t stream); +}; + } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS diff --git a/src/targets/qwen3_6/impl/runtime/dflash_context_impl.h b/src/targets/qwen3_6/impl/runtime/dflash_context_impl.h index 18b1935d42..c9c97f4283 100644 --- a/src/targets/qwen3_6/impl/runtime/dflash_context_impl.h +++ b/src/targets/qwen3_6/impl/runtime/dflash_context_impl.h @@ -38,4 +38,36 @@ void DFlashPersistentState::save_rewrite_checkpoint(std::int32_t source_slot, local.copy_slot_from(local, source_slot, destination_slot, stream); } +DFlash2PersistentState::DFlash2PersistentState(DeviceSpan backing, + const DFlash2PersistentLayout& layout) + : local(backing, layout.local), + rewrite_checkpoint_local(backing, layout.rewrite_checkpoint_local), + prefill_features(layout.prefill_features.bind(backing)), + prefill_positions(layout.prefill_positions.bind(backing)), + pending_features(layout.pending_features.bind(backing)) { + if (local.layer_count() != DFlash2Config::local_layers || + rewrite_checkpoint_local.layer_count() != DFlash2Config::local_layers || + local.capacity() != DFlash2Config::local_capacity || + rewrite_checkpoint_local.capacity() != DFlash2Config::local_capacity || + local.num_kv_heads() != DFlash2Config::kv_heads || + rewrite_checkpoint_local.num_kv_heads() != DFlash2Config::kv_heads || + local.head_dim() != DFlash2Config::head_dim || + rewrite_checkpoint_local.head_dim() != DFlash2Config::head_dim || + local.lane_capacity() != rewrite_checkpoint_local.lane_capacity()) { + throw std::invalid_argument("DFlash2 persistent cache layout is invalid"); + } +} + +CyclicKVCacheLayerView DFlash2PersistentState::local_layer(std::uint32_t layer) const { + return local.layer_view(layer); +} + +void DFlash2PersistentState::save_rewrite_checkpoint(std::int32_t lane, cudaStream_t stream) { + rewrite_checkpoint_local.copy_slot_from(local, lane, lane, stream); +} + +void DFlash2PersistentState::restore_rewrite_checkpoint(std::int32_t lane, cudaStream_t stream) { + local.copy_slot_from(rewrite_checkpoint_local, lane, lane, stream); +} + } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS diff --git a/src/targets/qwen3_6/impl/runtime/dflash_impl.h b/src/targets/qwen3_6/impl/runtime/dflash_impl.h index cffb1e1700..0c50960b30 100644 --- a/src/targets/qwen3_6/impl/runtime/dflash_impl.h +++ b/src/targets/qwen3_6/impl/runtime/dflash_impl.h @@ -172,7 +172,7 @@ void append_context_impl(Context& state, const Tensor& features, const Tensor& p ops::kv_cache_append_prefix( key_batch, value_batch, position_batch, local_counts, lanes, local_envelope, dflash_state(state).local_layer(static_cast(layer)), - state.execution.device.stream); + Config::local_capacity, state.execution.device.stream); } else { ops::kv_cache_append_prefix( key_batch, value_batch, position_batch, commit_counts, table_rows, envelope, diff --git a/src/targets/qwen3_6/impl/runtime/instance.h b/src/targets/qwen3_6/impl/runtime/instance.h index 3e3a7153bd..d13a9f80f7 100644 --- a/src/targets/qwen3_6/impl/runtime/instance.h +++ b/src/targets/qwen3_6/impl/runtime/instance.h @@ -16,12 +16,14 @@ using WeightsProfile = typename Variant::WeightsProfile; using TextConfig = typename Variant::TextConfig; using VisionConfig = typename Variant::VisionConfig; using DFlashConfig = typename Variant::DFlashConfig; +using DFlash2Config = typename Variant::DFlash2Config; using LoadedModelData = typename Variant::ModelView; using FullAttentionWeights = typename LoadedModelData::FullLayer; using GdnWeights = typename LoadedModelData::GdnLayer; using MlpWeights = typename Variant::PostMixerWeights; using MtpWeights = typename LoadedModelData::MtpLayer; using DFlashWeights = typename LoadedModelData::DFlash; +using DFlash2Weights = typename LoadedModelData::DFlash2; using FullAttentionProjectionWeights = typename Variant::FullAttentionProjectionWeights; using GdnProjectionWeights = typename Variant::GdnProjectionWeights; using VisionWeights = typename Variant::VisionWeights; @@ -74,4 +76,10 @@ inline std::vector dflash_graph_profiles(std::uint32_t ca return Variant::dflash_graph_profiles(capacity, draft_window, batch_size); } +inline std::vector dflash2_graph_profiles(std::uint32_t capacity, + std::uint32_t draft_window, + std::uint32_t batch_size) { + return Variant::dflash2_graph_profiles(capacity, draft_window, batch_size); +} + } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS diff --git a/src/targets/qwen3_6/impl/runtime/instantiate.h b/src/targets/qwen3_6/impl/runtime/instantiate.h index ab76f098b5..ff667049c4 100644 --- a/src/targets/qwen3_6/impl/runtime/instantiate.h +++ b/src/targets/qwen3_6/impl/runtime/instantiate.h @@ -21,6 +21,7 @@ #include "targets/qwen3_6/impl/runtime/graph_impl.h" #include "targets/qwen3_6/impl/runtime/speculative_target_impl.h" #include "targets/qwen3_6/impl/runtime/dflash_impl.h" +#include "targets/qwen3_6/impl/runtime/dflash2_impl.h" #include "targets/qwen3_6/impl/runtime/decode_impl.h" #include "targets/qwen3_6/impl/runtime/mtp_impl.h" #include "targets/qwen3_6/impl/runtime/request_plan_impl.h" diff --git a/src/targets/qwen3_6/impl/runtime/layouts.h b/src/targets/qwen3_6/impl/runtime/layouts.h index 859f4f5d7e..c331d1f53a 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts.h +++ b/src/targets/qwen3_6/impl/runtime/layouts.h @@ -31,11 +31,24 @@ struct DFlashPersistentLayout { [[nodiscard]] std::size_t kv_payload_bytes() const noexcept { return full.payload_bytes(); } }; +struct DFlash2PersistentLayout { + CyclicKVCacheLayout local; + CyclicKVCacheLayout rewrite_checkpoint_local; + TensorLayout prefill_features; + TensorLayout prefill_positions; + TensorLayout pending_features; + + [[nodiscard]] std::size_t kv_payload_bytes() const noexcept { + return local.payload_bytes() + rewrite_checkpoint_local.payload_bytes(); + } +}; + struct PersistentLayout { qwen3_6::DecoderStateLayout decoder; qwen3_6::StateImageDeviceLayout state_images; std::optional replay_records; std::optional dflash; + std::optional dflash2; qwen3_6::RoundStateLayout round; TensorLayout prefill_hidden; std::optional score_hidden; @@ -61,6 +74,8 @@ struct WorkspacePlan { std::size_t mtp_round = 0; std::size_t dflash_context = 0; std::size_t dflash_round = 0; + std::size_t dflash2_context = 0; + std::size_t dflash2_round = 0; std::size_t causal_score = 0; std::size_t general_capacity = 0; std::optional vision; diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index b21165ac86..8c5eee0dff 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/src/targets/qwen3_6/impl/runtime/layouts_impl.h @@ -220,6 +220,30 @@ PersistentLayout persistent_layout(const SequencePlanImpl& plan) { "DFlash pending target features"); } } + if constexpr (Variant::supports_dflash2) { + if (plan.features.dflash2()) { + DFlash2PersistentLayout& dflash2 = out.dflash2.emplace(); + dflash2.local = plan_cyclic_kv_cache(builder, DFlash2Config::local_layers, + DFlash2Config::local_capacity, + DFlash2Config::kv_heads, DFlash2Config::head_dim, + static_cast(plan.max_concurrency)); + dflash2.rewrite_checkpoint_local = plan_cyclic_kv_cache( + builder, DFlash2Config::local_layers, DFlash2Config::local_capacity, + DFlash2Config::kv_heads, DFlash2Config::head_dim, + static_cast(plan.max_concurrency)); + dflash2.prefill_features = add_tensor( + builder, DType::BF16, {DFlash2Config::feature_rows, effective_prefill_chunk}, + "DFlash2 prefill target features"); + dflash2.prefill_positions = add_tensor(builder, DType::I32, {effective_prefill_chunk}, + "DFlash2 prefill target positions"); + dflash2.pending_features = add_tensor( + builder, DType::BF16, + {DFlash2Config::feature_rows, + static_cast(plan.draft_window + 1U), + static_cast(plan.max_concurrency)}, + "DFlash2 pending target features"); + } + } out.round = qwen3_6::begin_round_state_layout( builder, qwen3_6::RoundStateSpec{.hidden = TextConfig::hidden, @@ -227,7 +251,7 @@ PersistentLayout persistent_layout(const SequencePlanImpl& plan) { .batch_capacity = plan.max_concurrency, .draft_window = plan.draft_window, .enable_mtp = plan.features.mtp(), - .enable_dflash = plan.features.dflash()}); + .enable_dflash = plan.features.dflash_like()}); out.prefill_hidden = add_tensor( builder, DType::BF16, {TextConfig::hidden, effective_prefill_chunk}, "step prefill hidden"); if (plan.causal_scoring) { @@ -572,9 +596,110 @@ WorkspacePlan build_workspace_plan(const SequencePlanImpl& plan) { } } + if (plan.features.dflash2()) { + if constexpr (!Variant::supports_dflash2) { + throw std::logic_error("unsupported target reached DFlash2 scratch planning"); + } else { + const auto dflash2_context_capacity = [&](std::int32_t tokens, bool compact_input) { + WorkspaceLayoutBuilder layout; + if (compact_input) { + matrix(layout, DType::BF16, DFlash2Config::feature_rows, tokens); + } + (void)workspace_recipe::dflash_context(layout, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, DFlash2Config::hidden, DFlash2Config::feature_rows, + ops::LinearPolicy::A16Only, tokens, tokens); + { + auto layer = layout.scope(); + (void)workspace_recipe::dflash_context_layer(layout, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, DFlash2Config::kv_size, DFlash2Config::hidden, + ops::LinearPolicy::A16Only, tokens, tokens); + } + return finish(layout); + }; + const auto dflash2_proposal_capacity = [&](std::int32_t width, std::int32_t batch) { + WorkspaceLayoutBuilder layout; + const std::int32_t tokens = width * batch; + matrix(layout, DType::BF16, DFlash2Config::hidden, tokens); + { + auto attention = layout.scope(); + (void)workspace_recipe::dflash_attention(layout, tokens); + matrix(layout, DType::BF16, 1280, tokens); + matrix(layout, DType::BF16, DFlash2Config::hidden, tokens); + matrix(layout, DType::BF16, DFlash2Config::hidden, tokens); + scratch(layout, ops::swa_workspace_capacity_bytes( + {0, plan.capacity}, width, width, batch, + DFlash2Config::local_window)); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, DFlash2Config::query_size + 2 * DFlash2Config::kv_size, + DFlash2Config::hidden, ops::LinearPolicy::A16Only, tokens, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, DFlash2Config::hidden, DFlash2Config::query_size, + ops::LinearPolicy::A16Only, tokens, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, 1280, DFlash2Config::hidden, ops::LinearPolicy::A16Only, + tokens, tokens); + } + { + auto mlp = layout.scope(); + (void)workspace_recipe::dflash_mlp(layout, tokens); + matrix(layout, DType::BF16, 1280, tokens); + matrix(layout, DType::BF16, DFlash2Config::hidden, tokens); + matrix(layout, DType::BF16, DFlash2Config::hidden, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, 2 * DFlash2Config::intermediate, DFlash2Config::hidden, + ops::LinearPolicy::A16Only, tokens, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, DFlash2Config::hidden, DFlash2Config::intermediate, + ops::LinearPolicy::A16Only, tokens, tokens); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, 1280, DFlash2Config::hidden, ops::LinearPolicy::A16Only, + tokens, tokens); + } + matrix(layout, DType::BF16, DFlash2Config::hidden, drafts * batch); + matrix(layout, DType::BF16, DFlash2Config::hidden, drafts * batch); + matrix(layout, DType::BF16, TextConfig::output_rows, drafts * batch); + matrix(layout, DType::BF16, DFlash2Config::selector_rank, drafts * batch); + (void)ops::linear_workspace_capacity_bytes( + QType::BF16_CTRL, DFlash2Config::selector_rank, DFlash2Config::hidden, + ops::LinearPolicy::A16Only, drafts * batch, drafts * batch); + matrix(layout, DType::I32, batch * DFlash2Config::block_drafts * + DFlash2Config::selector_top_k, + 1); + matrix(layout, DType::FP32, batch * DFlash2Config::block_drafts * + DFlash2Config::selector_top_k, + 1); + matrix(layout, DType::FP32, batch * DFlash2Config::block_drafts * + DFlash2Config::selector_top_k * + DFlash2Config::selector_top_k, + 1); + return finish(layout); + }; + + out.dflash2_context = dflash2_context_capacity(chunk, false); + for (std::int32_t batch = 1; batch <= static_cast(plan.max_concurrency); + ++batch) { + const std::int32_t aggregate = verify * batch; + WorkspaceLayoutBuilder target; + matrix(target, DType::BF16, TextConfig::hidden, aggregate); + target_body(target, aggregate, aggregate, qwen3_6::TextPhase::Verify, + GdnWorkspacePath::ReplayRecord, batch, verify, verify, text_envelope); + const std::size_t accept = + ops::speculative_accept_greedy_drafts_workspace_capacity_bytes( + TextConfig::token_domain, drafts, drafts, batch, batch); + const std::size_t proposal = dflash2_proposal_capacity(verify, batch); + out.dflash2_round = std::max({out.dflash2_round, finish(target), accept, + dflash2_context_capacity(aggregate, true), + proposal}); + } + } + } + out.general_capacity = std::max({out.text_prefill, out.ordinary_round, out.mtp_prefill, out.mtp_round, - out.dflash_context, out.dflash_round, out.causal_score}); + out.dflash_context, out.dflash_round, out.dflash2_context, out.dflash2_round, + out.causal_score}); out.capacity = out.general_capacity; if (plan.features.vision) { const std::uint32_t merged = static_cast( diff --git a/src/targets/qwen3_6/impl/runtime/program.h b/src/targets/qwen3_6/impl/runtime/program.h index 6c0f8d33a0..413285eab0 100644 --- a/src/targets/qwen3_6/impl/runtime/program.h +++ b/src/targets/qwen3_6/impl/runtime/program.h @@ -655,6 +655,7 @@ class ProgramImplCore { std::optional replay_records; std::optional replay_fold; std::optional dflash; + std::optional dflash2; qwen3_6::RoundState io; Tensor prefill_hidden; std::optional score_hidden; @@ -672,6 +673,7 @@ class ProgramImplCore { DecodeGraphFamily ordinary_graphs; DecodeGraphFamily mtp_graphs; DecodeGraphFamily dflash_graphs; + DecodeGraphFamily dflash2_graphs; PinnedHostBuffer round_host; std::optional score_logprobs_host; @@ -685,6 +687,9 @@ class ProgramImplCore { std::optional dflash_host; qwen3_6::DFlashDecodeIngress* dflash_host_ingress = nullptr; qwen3_6::DFlashDecodeEgress* dflash_host_egress = nullptr; + std::optional dflash2_host; + qwen3_6::DFlashDecodeIngress* dflash2_host_ingress = nullptr; + qwen3_6::DFlashDecodeEgress* dflash2_host_egress = nullptr; std::size_t workspace_logical_peak_bytes = 0; std::size_t vision_handoff_peak_bytes = 0; @@ -1162,6 +1167,10 @@ class ProgramImplCore { decode_dflash_batch(std::span lanes, std::span budgets, runtime::ExecutionTiming* failed_timing); + [[nodiscard]] runtime::BatchedGeneratedRound + decode_dflash2_batch(std::span lanes, + std::span budgets, + runtime::ExecutionTiming* failed_timing); void resize_sequence_kv_entitlement(SequenceState& sequence, std::uint32_t text_pages, std::uint32_t backend_pages); void bind_sequence_kv(SequenceState& sequence); diff --git a/src/targets/qwen3_6/impl/runtime/program_impl.h b/src/targets/qwen3_6/impl/runtime/program_impl.h index 36f04c8bdb..713e5d0416 100644 --- a/src/targets/qwen3_6/impl/runtime/program_impl.h +++ b/src/targets/qwen3_6/impl/runtime/program_impl.h @@ -4,6 +4,7 @@ #include "core/nvtx.h" #include "targets/qwen3_6/impl/runtime/schedule.h" +#include "targets/qwen3_6/impl/runtime/spec_decision.h" #include "ninfer/ops/gdn_replay.h" #include "ninfer/ops/linear.h" #include "ninfer/ops/prepare_ragged_prefix.h" @@ -603,6 +604,15 @@ schedule::DFlashEnvelopes dflash_envelopes(std::uint32_t min_frontier, std::uint }; } +schedule::DFlash2Envelopes dflash2_envelopes(std::uint32_t min_frontier, + std::uint32_t max_frontier, std::uint32_t k) { + (void)min_frontier; + return schedule::DFlash2Envelopes{ + .local = {0, max_frontier}, + .append = {0, k + 1}, + }; +} + DecodeGraphProfile& select_graph_profile(DecodeGraphFamily& family, std::uint32_t batch_size, std::uint32_t frontier, const char* label) { const auto it = std::find_if( @@ -750,6 +760,10 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence ? std::make_optional(sizeof(qwen3_6::DFlashDecodeIngress) + sizeof(qwen3_6::DFlashDecodeEgress)) : std::nullopt), + dflash2_host(plan.speculative_backend == SpeculativeBackend::DFlash2 + ? std::make_optional(sizeof(qwen3_6::DFlashDecodeIngress) + + sizeof(qwen3_6::DFlashDecodeEgress)) + : std::nullopt), context_source_ready_(device_in), context_completion_(device_in), context_transfer_timers_{CudaEventTimer(device_in, device_in.transfer_stream), CudaEventTimer(device_in, device_in.transfer_stream), @@ -759,15 +773,19 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence } if (model.features != plan.features || model.mtp.has_value() != plan.features.mtp() || model.dflash.has_value() != plan.features.dflash() || + model.dflash2.has_value() != plan.features.dflash2() || model.optimized_proposal.has_value() != plan.features.optimized_proposal() || model.vision.has_value() != plan.features.vision) { throw std::invalid_argument( "Qwen3.6 loaded weights do not match the frozen startup features"); } - if (model.mtp.has_value() && model.dflash.has_value()) { + if (model.mtp.has_value() && (model.dflash.has_value() || model.dflash2.has_value())) { throw std::invalid_argument("MTP and DFlash model views are mutually exclusive"); } - if (model.dflash.has_value() && model.vision.has_value()) { + if (model.dflash.has_value() && model.dflash2.has_value()) { + throw std::invalid_argument("DFlash and DFlash2 model views are mutually exclusive"); + } + if ((model.dflash.has_value() || model.dflash2.has_value()) && model.vision.has_value()) { throw std::invalid_argument("DFlash and Vision model views are mutually exclusive"); } if (workspace_plan.general_capacity == 0 || @@ -847,6 +865,10 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence if (dflash.has_value() != plan.features.dflash()) { throw std::logic_error("DFlash state does not match the frozen sequence plan"); } + if (plan.persistent.dflash2) { dflash2.emplace(backing, *plan.persistent.dflash2); } + if (dflash2.has_value() != plan.features.dflash2()) { + throw std::logic_error("DFlash2 state does not match the frozen sequence plan"); + } if (qwen3_6::PagedKVCache* backend = backend_kv_cache()) { backend_host_kv_page_stride = plan_host_kv_page_layout(backend->page_pool().geometry()).page_stride; @@ -892,10 +914,14 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence if (io.ordinary.has_value() != (speculative_backend == SpeculativeBackend::None)) { throw std::logic_error("ordinary decode frame does not match the sequence plan"); } - if (io.dflash_prefill.has_value() != (speculative_backend == SpeculativeBackend::DFlash)) { + if (io.dflash_prefill.has_value() != + (speculative_backend == SpeculativeBackend::DFlash || + speculative_backend == SpeculativeBackend::DFlash2)) { throw std::logic_error("DFlash prefill scratch does not match the sequence plan"); } - if (io.dflash_decode.has_value() != (speculative_backend == SpeculativeBackend::DFlash)) { + if (io.dflash_decode.has_value() != + (speculative_backend == SpeculativeBackend::DFlash || + speculative_backend == SpeculativeBackend::DFlash2)) { throw std::logic_error("DFlash decode frame does not match the sequence plan"); } prefill_hidden = plan.persistent.prefill_hidden.bind(backing); @@ -946,6 +972,14 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence *dflash_host_ingress = {}; *dflash_host_egress = {}; } + if (dflash2_host) { + dflash2_host_ingress = static_cast(dflash2_host->data()); + dflash2_host_egress = reinterpret_cast( + static_cast(dflash2_host->data()) + + sizeof(qwen3_6::DFlashDecodeIngress)); + *dflash2_host_ingress = {}; + *dflash2_host_egress = {}; + } if (io.dflash_prefill) { CUDA_CHECK(cudaMemsetAsync(io.dflash_prefill->produced_count.data, 0, io.dflash_prefill->produced_count.bytes(), device.stream)); @@ -1067,6 +1101,7 @@ std::vector ProgramImplCore::causal_score(PreparedPromptData&& prompt, decoder->text_kv, nullptr, nullptr, + nullptr, cursor, nullptr, nullptr, @@ -8494,6 +8529,7 @@ runtime::ExecutionTiming ProgramImplCore::append_forced_tokens( decoder->text_kv, decoder->mtp_cache(), dflash ? &*dflash : nullptr, + dflash2 ? &*dflash2 : nullptr, cursor, nullptr, nullptr, @@ -8504,7 +8540,8 @@ runtime::ExecutionTiming ProgramImplCore::append_forced_tokens( mark_workspace_usage(speculative_backend == SpeculativeBackend::Mtp ? workspace_plan.mtp_prefill : workspace_plan.text_prefill); - if (speculative_backend == SpeculativeBackend::DFlash) { + if (speculative_backend == SpeculativeBackend::DFlash || + speculative_backend == SpeculativeBackend::DFlash2) { mark_workspace_usage(workspace_plan.dflash_context); } const schedule::PrefillChunkResult result = schedule::prefill_text_chunk( @@ -9430,7 +9467,8 @@ void ProgramImplCore::start_sequence(std::uint32_t lane, SequenceState& sequence ? std::min(capacity, prompt_tokens + (initial_mtp_extent == 0 ? 0U : initial_mtp_extent - 1U)) : speculative_backend == SpeculativeBackend::DFlash ? prompt_tokens - : 0U; + : speculative_backend == SpeculativeBackend::DFlash2 ? prompt_tokens + : 0U; materialize_sequence_kv(sequence, prompt_tokens, backend_materialized); install_sampling(sequence, request, request_plan.sampling); sequence.rope_delta = staged.prompt.rope_delta; @@ -9461,6 +9499,16 @@ void ProgramImplCore::start_sequence(std::uint32_t lane, SequenceState& sequence CUDA_CHECK(cudaMemcpyAsync(io.dflash_decode->ingress.data, dflash_host_ingress, sizeof(qwen3_6::DFlashDecodeIngress), cudaMemcpyHostToDevice, device.stream)); + } else if (speculative_backend == SpeculativeBackend::DFlash2) { + if (!dflash2 || !io.dflash_decode) { + throw std::logic_error("DFlash2 prefill state is incomplete"); + } + *dflash2_host_ingress = {}; + dflash2_host_ingress->active_lanes[0] = static_cast(sequence.lane); + dflash2_host_ingress->dflash_kv_table_rows[0] = 0; + CUDA_CHECK(cudaMemcpyAsync(io.dflash_decode->ingress.data, dflash2_host_ingress, + sizeof(qwen3_6::DFlashDecodeIngress), cudaMemcpyHostToDevice, + device.stream)); } staged.elapsed_seconds += std::chrono::duration(Clock::now() - started).count(); @@ -10449,6 +10497,53 @@ void ProgramImplCore::prepare_graphs() { } } } + if (speculative_backend == SpeculativeBackend::DFlash2) { + const auto batch_one_profiles = dflash2_graph_profiles(capacity, draft_window, 1); + validate_graph_profiles(batch_one_profiles, capacity - 1, "DFlash2"); + schedule::DFlash2BatchContext dflash2_state{execution_core(), + decoder->text_kv, + *dflash2, + *io.dflash_decode, + *dflash2_host_ingress, + *dflash2_host_egress, + state_images->continuation_hidden_store()}; + const GraphExecutionProfile code_warm = batch_one_profiles.front(); + const ops::CausalAttentionExecutionEnvelope code_warm_target{ + 1, static_cast(std::min( + capacity, static_cast(code_warm.max) + draft_window + 1ULL))}; + prepare_representative(code_warm.min, 1); + device.synchronize(); + schedule::dflash2_decode_batch(dflash2_state, 1, draft_window, + dflash2_envelopes(code_warm.min, code_warm.max, draft_window), + code_warm_target, nullptr); + device.synchronize(); + + dflash2_graphs.profiles.reserve(batch_one_profiles.size() * max_concurrency); + for (std::uint32_t batch_size = 1; batch_size <= max_concurrency; ++batch_size) { + const auto planned_profiles = + batch_size == 1 ? batch_one_profiles + : dflash2_graph_profiles(capacity, draft_window, batch_size); + validate_graph_profiles(planned_profiles, capacity - 1, "DFlash2"); + for (const GraphExecutionProfile planned : planned_profiles) { + dflash2_graphs.profiles.emplace_back(); + DecodeGraphProfile& profile = dflash2_graphs.profiles.back(); + profile.batch_size = batch_size; + profile.min_execution_frontier = planned.min; + profile.max_execution_frontier = planned.max; + profile.topology_class = + planned.topology_class * max_concurrency + (batch_size - 1U); + const ops::CausalAttentionExecutionEnvelope target_envelope{ + 1, + static_cast(std::min( + capacity, static_cast(planned.max) + draft_window + 1ULL))}; + + schedule::capture_dflash2_decode_batch( + dflash2_state, static_cast(batch_size), draft_window, + dflash2_envelopes(planned.min, planned.max, draft_window), target_envelope, + profile.definition); + } + } + } if (!ordinary_graphs.profiles.empty()) { instantiate_graph_family(ordinary_graphs, "ordinary", device, prepare_representative); @@ -10459,6 +10554,9 @@ void ProgramImplCore::prepare_graphs() { if (speculative_backend == SpeculativeBackend::DFlash) { instantiate_graph_family(dflash_graphs, "DFlash", device, prepare_representative); } + if (speculative_backend == SpeculativeBackend::DFlash2) { + instantiate_graph_family(dflash2_graphs, "DFlash2", device, prepare_representative); + } clear_stable_controls(); state_images->zero_all(device.stream); @@ -10470,6 +10568,14 @@ void ProgramImplCore::prepare_graphs() { CUDA_CHECK(cudaMemsetAsync(dflash->pending_features.data, 0, dflash->pending_features.bytes(), device.stream)); } + if (dflash2) { + CUDA_CHECK(cudaMemsetAsync(dflash2->prefill_features.data, 0, + dflash2->prefill_features.bytes(), device.stream)); + CUDA_CHECK(cudaMemsetAsync(dflash2->prefill_positions.data, 0, + dflash2->prefill_positions.bytes(), device.stream)); + CUDA_CHECK(cudaMemsetAsync(dflash2->pending_features.data, 0, + dflash2->pending_features.bytes(), device.stream)); + } CUDA_CHECK(cudaMemsetAsync(token_counts.data, 0, token_counts.bytes(), device.stream)); device.synchronize(); for (std::uint32_t row = 0; row < max_concurrency; ++row) { @@ -10667,6 +10773,7 @@ ProgramImplCore::advance_prefill(SequenceState& sequence, RequestControl& reques decoder->text_kv, decoder->mtp_cache(), dflash ? &*dflash : nullptr, + dflash2 ? &*dflash2 : nullptr, staged.cursor, static_cast( sampling_config.slice(1, static_cast(sequence.lane), 1).data), @@ -11426,9 +11533,202 @@ ProgramImplCore::decode_raw(std::span lanes, if (speculative_backend == SpeculativeBackend::Mtp) { return decode_mtp_batch(lanes, budgets, failed_timing); } + if (speculative_backend == SpeculativeBackend::DFlash2) { + return decode_dflash2_batch(lanes, budgets, failed_timing); + } return decode_dflash_batch(lanes, budgets, failed_timing); } +runtime::BatchedGeneratedRound +ProgramImplCore::decode_dflash2_batch(std::span lanes, + std::span budgets, + runtime::ExecutionTiming* failed_timing) { + if (speculative_backend != SpeculativeBackend::DFlash2 || !io.dflash_decode || !dflash2) { + throw std::logic_error("DFlash2 batch execution requires the DFlash2 backend"); + } + if (lanes.empty() || lanes.size() > max_concurrency || budgets.size() != lanes.size()) { + throw std::invalid_argument("DFlash2 batch membership is invalid"); + } + + const std::uint32_t width = draft_window + 1U; + std::uint32_t maximum_frontier = 0; + std::uint32_t maximum_target_tokens = 1; + for (std::size_t row = 0; row < lanes.size(); ++row) { + const std::uint32_t lane = lanes[row]; + if (lane >= max_concurrency || + std::find(lanes.begin(), lanes.begin() + static_cast(row), lane) != + lanes.begin() + static_cast(row)) { + throw std::invalid_argument("DFlash2 batch contains an invalid or duplicate lane"); + } + const SequenceState& sequence = active_sequence(lane); + const RequestControl& request = requests[lane]; + if (request.lifecycle != Lifecycle::Active || + budgets[row].generated_tokens_remaining == 0 || !sequence.kv || + text_kv_addresses->bound_row(sequence.kv->text) < 0 || sequence.execution_frontier >= capacity || + sequence.text_kv_valid != sequence.execution_frontier || + sequence.dflash_context_frontier > sequence.execution_frontier || + sequence.execution_frontier - sequence.dflash_context_frontier > width || + sequence.ledger_frontier != sequence.execution_frontier + 1 || + sequence.ledger.size() != sequence.ledger_frontier || + sequence.prefix_identity.size() != sequence.ledger_frontier) { + throw std::logic_error("DFlash2 batch row is not decode-ready"); + } + const std::uint32_t max_by_budget = budgets[row].generated_tokens_remaining > 1 + ? budgets[row].generated_tokens_remaining - 1U + : 0U; + // Length demotion (measured on the 32GiB 5090, threshold 20480): + // past ~17-20k prompt tokens the DFlash2 advantage over flat MTP + // disappears (~+2.6% at 15k), so lanes beyond the threshold draft + // nothing this round (dense single-token round through the same + // kernel/graph machinery as the spec-degrade fallback). The target + // KV cache is shared, so the row stays decode-ready; a fresh request + // below the threshold simply drafts again. + const bool lane_beyond_demote = + sequence.execution_frontier > qwen3_6::kSpecDemoteTokens; + const std::uint32_t extent = + lane_beyond_demote + ? 0U + : std::min({draft_window, max_by_budget, + capacity - sequence.execution_frontier - 1U}); + maximum_frontier = std::max(maximum_frontier, sequence.execution_frontier); + maximum_target_tokens = + std::max(maximum_target_tokens, sequence.execution_frontier + extent + 1U); + } + + runtime::ExecutionTimingRecorder timing(runtime::ExecutionTimingPhase::Submit, failed_timing); + const auto started = Clock::now(); + try { + std::optional submit_range; + submit_range.emplace(nvtx::Name::DecodeDFlashSubmit, nvtx::Category::DFlash, + static_cast(lanes.size())); + DecodeGraphExecutable* executable = nullptr; + schedule::DFlash2Envelopes envelopes = dflash2_envelopes(0, maximum_frontier, draft_window); + ops::CausalAttentionExecutionEnvelope target_envelope{1, maximum_target_tokens}; + if (use_cuda_graph) { + DecodeGraphProfile& profile = + select_graph_profile(dflash2_graphs, static_cast(lanes.size()), + maximum_frontier, "DFlash2 batch"); + executable = &install_graph_profile(dflash2_graphs, profile, "DFlash2 batch"); + envelopes = dflash2_envelopes(profile.min_execution_frontier, + profile.max_execution_frontier, draft_window); + target_envelope = { + 1, static_cast(std::min( + capacity, static_cast(profile.max_execution_frontier) + + draft_window + 1ULL))}; + } + + for (std::size_t row = 0; row < lanes.size(); ++row) { + SequenceState& sequence = active_sequence(lanes[row]); + const RequestControl& request = requests[lanes[row]]; + const std::uint32_t frontier = sequence.execution_frontier; + const std::uint32_t max_by_budget = budgets[row].generated_tokens_remaining > 1 + ? budgets[row].generated_tokens_remaining - 1U + : 0U; + const std::uint32_t extent = + std::min({draft_window, max_by_budget, capacity - frontier - 1U}); + dflash2_host_ingress->anchors[row] = sequence.ledger.back(); + dflash2_host_ingress->execution_frontiers[row] = + checked_i32(frontier, "DFlash2 batch frontier"); + dflash2_host_ingress->context_frontiers[row] = + checked_i32(sequence.dflash_context_frontier, "DFlash2 context frontier"); + dflash2_host_ingress->proposal_extents[row] = static_cast(extent); + dflash2_host_ingress->target_valid_columns[row] = static_cast(extent + 1U); + dflash2_host_ingress->text_kv_table_rows[row] = + text_kv_addresses->bound_row(sequence.kv->text); + dflash2_host_ingress->dflash_kv_table_rows[row] = 0; + dflash2_host_ingress->active_lanes[row] = static_cast(sequence.lane); + dflash2_host_ingress->sampling[row] = request.sampling_host; + materialize_sequence_kv(sequence, frontier + extent + 1U, frontier); + } + + schedule::DFlash2BatchContext schedule_state{ + {device, model, work, state_images->linear(), + replay_records ? &*replay_records : nullptr, io, prefill_hidden, prefill_chunk, + proposal_head}, + decoder->text_kv, + *dflash2, + *io.dflash_decode, + *dflash2_host_ingress, + *dflash2_host_egress, + state_images->continuation_hidden_store()}; + + mark_workspace_usage(workspace_plan.dflash2_round); + schedule::dflash2_decode_batch(schedule_state, static_cast(lanes.size()), + draft_window, envelopes, target_envelope, executable); + submit_range.reset(); + timing.begin_wait(); + { + nvtx::ScopedRange wait_range(nvtx::Name::DecodeDFlashWait, nvtx::Category::Control, + static_cast(lanes.size())); + device.synchronize(); + } + timing.end_wait(); + + const double seconds = std::chrono::duration(Clock::now() - started).count(); + for (std::size_t row = 0; row < lanes.size(); ++row) { + SequenceState& sequence = active_sequence(lanes[row]); + RequestControl& request = requests[lanes[row]]; + const std::uint32_t base_E = sequence.execution_frontier; + const std::uint32_t base_S = sequence.ledger_frontier; + const std::int32_t count_i = dflash2_host_egress->licensed_counts[row]; + const std::int32_t accepted_i = dflash2_host_egress->accepted_drafts[row]; + const std::uint32_t extent = + static_cast(dflash2_host_egress->proposal_extents[row]); + if (count_i <= 0 || count_i > static_cast(width) || accepted_i < 0 || + accepted_i + 1 != count_i || accepted_i > static_cast(extent) || + extent > width || + extent > static_cast(dflash2_host_ingress->proposal_extents[row]) || + static_cast(count_i) > budgets[row].generated_tokens_remaining || + static_cast(base_E) + static_cast(count_i) > + capacity) { + throw std::runtime_error("DFlash2 batch returned invalid row metadata"); + } + const std::span row_tokens(dflash2_host_egress->licensed_tokens.data() + + row * width, + static_cast(count_i)); + validate_licensed_tokens(row_tokens); + if (extent == 0) { + request.speculative_stats.fallback_steps += 1; + } else { + request.speculative_stats.rounds += 1; + request.speculative_stats.drafted_tokens += extent; + request.speculative_stats.accepted_tokens += static_cast(accepted_i); + for (std::int32_t i = 0; i < accepted_i; ++i) { + request.speculative_stats.accepted_per_position[static_cast(i)] += + 1; + } + } + sequence.dflash_context_frontier = base_E; + request.pending = PendingCandidate{ + .kind = PendingKind::Speculative, + .base_E = base_E, + .base_S = base_S, + .prompt_tokens = 0, + .produced = static_cast(count_i), + }; + request.lifecycle = Lifecycle::Pending; + request.timings.decode_seconds += seconds; + } + return runtime::BatchedGeneratedRound{ + .tokens = std::span(dflash2_host_egress->licensed_tokens.data(), + lanes.size() * width), + .row_counts = std::span(dflash2_host_egress->licensed_counts.data(), + lanes.size()), + .row_stride = width, + }; + } catch (...) { + timing.begin_wait(); + try { + nvtx::ScopedRange wait_range(nvtx::Name::DecodeDFlashWait, nvtx::Category::Control, + static_cast(lanes.size())); + device.synchronize(); + } catch (...) {} + timing.end_wait(); + clear_execution_failure_lanes(lanes); + throw; + } +} + runtime::ExecutionTiming ProgramImplCore::resolve_non_speculative_pending(SequenceState& sequence, RequestControl& request, std::uint32_t accepted_tokens, bool terminal, diff --git a/src/targets/qwen3_6/impl/runtime/request_plan_impl.h b/src/targets/qwen3_6/impl/runtime/request_plan_impl.h index 504725148e..549084dae0 100644 --- a/src/targets/qwen3_6/impl/runtime/request_plan_impl.h +++ b/src/targets/qwen3_6/impl/runtime/request_plan_impl.h @@ -571,6 +571,15 @@ std::optional ProgramImplCore::inspect_lane( shared_source->backend_frontier < plan->reuse_base)))) { throw std::logic_error("published DFlash checkpoint is not materializable"); } + if ((is_rewrite_checkpoint_restore(plan->reuse) || + plan->reuse == ReusePath::PrivateLongAnchor || + plan->reuse == ReusePath::SharedStablePrefix) && + speculative_backend == SpeculativeBackend::DFlash2 && + (!dflash2 || + (source != nullptr && source->dflash_context_frontier < plan->reuse_base) || + (shared_source != nullptr && shared_source->backend_frontier < plan->reuse_base))) { + throw std::logic_error("published DFlash2 checkpoint is not materializable"); + } const std::optional& desired = base.rewrite_checkpoint; const bool can_retain_rewrite = diff --git a/src/targets/qwen3_6/impl/runtime/schedule.h b/src/targets/qwen3_6/impl/runtime/schedule.h index c9dc6eebec..d2f62c31b8 100644 --- a/src/targets/qwen3_6/impl/runtime/schedule.h +++ b/src/targets/qwen3_6/impl/runtime/schedule.h @@ -9,6 +9,7 @@ #include "ninfer/ops/sampling.h" #include "ninfer/ops/sliding_window_attention.h" #include "ninfer/ops/softmax_attention.h" +#include "ninfer/ops/swa.h" #include "targets/qwen3_6/impl/runtime/dflash_context.h" #include "targets/qwen3_6/impl/runtime/text_context.h" #include "targets/qwen3_6/impl/runtime/vision_context.h" @@ -47,6 +48,7 @@ struct PrefillContext { const qwen3_6::PagedKVCache& text_cache; const qwen3_6::PagedKVCache* mtp_cache; DFlashPersistentState* dflash; + DFlash2PersistentState* dflash2; std::uint32_t text_kv_base; const ops::SamplingConfig* sampling; Tensor* rewrite_checkpoint_hidden; @@ -90,6 +92,21 @@ struct DFlashAppendContext { DFlashPersistentState& dflash; }; +struct DFlash2BatchContext { + ExecutionCore execution; + const qwen3_6::PagedKVCache& text_cache; + DFlash2PersistentState& dflash2; + qwen3_6::DFlashDecodeState& frame; + const qwen3_6::DFlashDecodeIngress& host_ingress; + qwen3_6::DFlashDecodeEgress& host_egress; + Tensor& continuation_hidden_store; +}; + +struct DFlash2AppendContext { + ExecutionCore execution; + DFlash2PersistentState& dflash2; +}; + struct MtpCausalAttentionEnvelopes { ops::CausalAttentionExecutionEnvelope target_verify; ops::CausalAttentionExecutionEnvelope batch; @@ -102,6 +119,14 @@ struct DFlashEnvelopes { ops::KVCacheAppendPrefixExecutionEnvelope append; }; +// DFlash2 is a pure sliding-window draft: its local attention uses the swa +// kernel (symmetric window, no full-context stage), so only the local and +// append envelopes exist. +struct DFlash2Envelopes { + ops::SwaContextExecutionEnvelope local; + ops::KVCacheAppendPrefixExecutionEnvelope append; +}; + struct TargetVerifyFrameView { Tensor ids; Tensor cache_positions; @@ -196,4 +221,23 @@ void dflash_decode_batch(DFlashBatchContext& state, std::int32_t batch_size, std ops::CausalAttentionExecutionEnvelope target_envelope, DecodeGraphExecutable* executable); +[[nodiscard]] DFlashFeatureSink +dflash2_feature_sink(PrefillContext& state, DFlashFeatureSink::PrefillConsumer consume_prefill = {}); +void dflash2_append_context(DFlash2AppendContext& state, const Tensor& features, + const Tensor& positions, const Tensor& commit_counts, + const Tensor& lanes, const Tensor& table_rows, + ops::KVCacheAppendPrefixExecutionEnvelope envelope); +void dflash2_append_context(PrefillContext& state, const Tensor& features, const Tensor& positions, + const Tensor& commit_counts, const Tensor& lanes, + const Tensor& table_rows, + ops::KVCacheAppendPrefixExecutionEnvelope envelope); +void capture_dflash2_decode_batch(DFlash2BatchContext& state, std::int32_t batch_size, + std::uint32_t k, DFlash2Envelopes envelopes, + ops::CausalAttentionExecutionEnvelope target_envelope, + DecodeGraphDefinition& definition); +void dflash2_decode_batch(DFlash2BatchContext& state, std::int32_t batch_size, std::uint32_t k, + DFlash2Envelopes envelopes, + ops::CausalAttentionExecutionEnvelope target_envelope, + DecodeGraphExecutable* executable); + } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule diff --git a/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h b/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h index 40ccda9d3f..115c20905b 100644 --- a/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h @@ -17,6 +17,23 @@ DFlashFeatureSink make_dflash_prefill_sink(PrefillContext& state) { if (!state.execution.io.dflash_decode || state.dflash_host_ingress == nullptr) { throw std::logic_error("DFlash prefill controls are unavailable"); } + if (state.execution.model.features.dflash2()) { + return dflash2_feature_sink( + state, [&state](const Tensor& features, const Tensor& positions, + bool rewrite_checkpoint) { + auto& frame = *state.execution.io.dflash_decode; + Tensor count = frame.append_counts.slice(0, 0, 1); + Tensor lane = frame.state_destination_slots.slice(0, 0, 1); + Tensor row = frame.dflash_kv_table_rows.slice(0, 0, 1); + ops::set_i32_scalar(count, features.ne[1], state.execution.device.stream); + const auto exact = static_cast(features.ne[1]); + dflash2_append_context(state, features, positions, count, lane, row, {exact, exact}); + if (rewrite_checkpoint) { + state.dflash2->save_rewrite_checkpoint(state.dflash_host_ingress->active_lanes[0], + state.execution.device.stream); + } + }); + } return dflash_feature_sink( state, [&state](const Tensor& features, const Tensor& positions, bool rewrite_checkpoint) { auto& frame = *state.execution.io.dflash_decode; diff --git a/src/targets/qwen3_6/impl/runtime/workspace_recipe.h b/src/targets/qwen3_6/impl/runtime/workspace_recipe.h index 576d049439..551fe49561 100644 --- a/src/targets/qwen3_6/impl/runtime/workspace_recipe.h +++ b/src/targets/qwen3_6/impl/runtime/workspace_recipe.h @@ -252,10 +252,15 @@ struct DFlashAttentionRoots { Tensor query; Tensor key; Tensor attention; + Tensor attention_delta; }; template DFlashAttentionRoots dflash_attention(Allocator& allocator, std::int32_t tokens) { + Tensor attention_delta; + if constexpr (Config::bf16_weights) { + attention_delta = matrix(allocator, DType::BF16, Config::hidden, tokens); + } return { matrix(allocator, DType::BF16, Config::hidden, tokens), matrix(allocator, DType::BF16, Config::query_size, tokens), @@ -264,19 +269,30 @@ DFlashAttentionRoots dflash_attention(Allocator& allocator, std::int32_t tokens) matrix(allocator, DType::BF16, Config::query_size, tokens), matrix(allocator, DType::BF16, Config::kv_size, tokens), matrix(allocator, DType::BF16, Config::query_size, tokens), + attention_delta, }; } struct DFlashMlpRoots { Tensor hidden; + Tensor gate_up; Tensor intermediate; + Tensor delta; }; template DFlashMlpRoots dflash_mlp(Allocator& allocator, std::int32_t tokens) { + Tensor gate_up; + Tensor delta; + if constexpr (Config::bf16_weights) { + gate_up = matrix(allocator, DType::BF16, 2 * Config::intermediate, tokens); + delta = matrix(allocator, DType::BF16, Config::hidden, tokens); + } return { matrix(allocator, DType::BF16, Config::hidden, tokens), + gate_up, matrix(allocator, DType::BF16, Config::intermediate, tokens), + delta, }; } diff --git a/src/targets/qwen3_6/impl/state/round_state.cpp b/src/targets/qwen3_6/impl/state/round_state.cpp index b8a719a9fd..a8d441e188 100644 --- a/src/targets/qwen3_6/impl/state/round_state.cpp +++ b/src/targets/qwen3_6/impl/state/round_state.cpp @@ -338,6 +338,8 @@ DFlashDecodeState::DFlashDecodeState(DeviceSpan backing, const DFlashDecodeState egress_tensor(offsetof(DFlashDecodeEgress, licensed_counts), DType::I32, {batch}); accepted_drafts = egress_tensor(offsetof(DFlashDecodeEgress, accepted_drafts), DType::I32, {batch}); + egress_proposal_extents = + egress_tensor(offsetof(DFlashDecodeEgress, proposal_extents), DType::I32, {batch}); proposal_ids = layout.proposal_ids.bind(backing); proposal_positions = layout.proposal_positions.bind(backing); append_positions = layout.append_positions.bind(backing); diff --git a/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h b/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h index b9f1a33799..cf2b7bdf1b 100644 --- a/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h +++ b/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h @@ -33,6 +33,7 @@ enum class WeightsProfile : std::uint8_t { Qwen38GroupwiseInt, Qwen36Nvfp4, Qwen38Nvfp4, + Qwen38Nvfp4DFlash2, }; using Frontend = qwen3_6::Frontend; diff --git a/src/targets/qwen3_6_27b/impl/config.h b/src/targets/qwen3_6_27b/impl/config.h index 5c30aa1896..dcce9cdba3 100644 --- a/src/targets/qwen3_6_27b/impl/config.h +++ b/src/targets/qwen3_6_27b/impl/config.h @@ -1,94 +1,148 @@ -#pragma once - -#include -#include -#include - -#include - -namespace ninfer::targets::qwen3_6_27b::detail { - -struct TextConfig { - static constexpr int hidden = 5120; - static constexpr int layers = 64; - static constexpr int intermediate = 17408; - - // The output matrix is padded for the selected kernels. Only token IDs in - // [0, token_domain) are tokenizer-addressable and valid sampling results. - static constexpr int output_rows = 248320; - static constexpr int token_domain = static_cast(qwen3_6::kTokenDomain); - - static constexpr int gdn_conv_kernel = 4; - static constexpr int gdn_conv_state_width = gdn_conv_kernel - 1; - static constexpr int gdn_key_heads = 16; - static constexpr int gdn_key_head_dim = 128; - static constexpr int gdn_value_heads = 48; - static constexpr int gdn_value_head_dim = 128; - - static constexpr int query_heads = 24; - static constexpr int kv_heads = 4; - static constexpr int head_dim = 256; - static constexpr int rotary_dim = 64; - - static constexpr int full_attention_interval = qwen3_6::kHybridAttentionInterval; - static constexpr float rms_epsilon = 1.0e-6F; - static constexpr float rope_theta = 1.0e7F; - - static constexpr int key_dim = gdn_key_heads * gdn_key_head_dim; - static constexpr int value_dim = gdn_value_heads * gdn_value_head_dim; - static constexpr int convolution_dim = 2 * key_dim + value_dim; - static constexpr int query_size = query_heads * head_dim; - static constexpr int kv_size = kv_heads * head_dim; - static constexpr int query_projection_rows = 2 * query_size; - - static constexpr int mtp_layers = 1; - static constexpr int mtp_input_rows = 2 * hidden; - static constexpr int mtp_attention_input_rows = 2 * query_size + 2 * kv_size; - static constexpr int mtp_mlp_gate_up_rows = 2 * intermediate; - - [[nodiscard]] static constexpr bool is_full_attention(int layer) { - return qwen3_6::is_full_attention_layer(layer); - } - - [[nodiscard]] static constexpr int full_attention_layers() { - return qwen3_6::full_attention_layers(layers); - } - - [[nodiscard]] static constexpr int gdn_layers() { return qwen3_6::gdn_layers(layers); } - - [[nodiscard]] static constexpr int full_attention_index(int layer) { - return qwen3_6::full_attention_index(layer); - } - - [[nodiscard]] static constexpr int gdn_index(int layer) { return qwen3_6::gdn_index(layer); } -}; - -static_assert(TextConfig::full_attention_layers() == 16); -static_assert(TextConfig::gdn_layers() == 48); - -struct VisionConfig : qwen3_6::VisionBackboneConfig { - static constexpr int output_hidden = TextConfig::hidden; -}; - -struct DFlashConfig { - static constexpr bool supported = false; - static constexpr int local_layers = 0; - static constexpr int local_capacity = 0; - static constexpr int query_heads = 0; - static constexpr int kv_heads = 0; - static constexpr int head_dim = 0; - static constexpr int feature_rows = 0; - static constexpr int hidden = 0; - static constexpr int intermediate = 0; - static constexpr int query_size = 0; - static constexpr int kv_size = 0; -}; - -inline constexpr float kAttentionScale = 0.0625F; -inline constexpr float kGdnScale = 0.08838834764831845F; -inline constexpr std::uint32_t kPrefillChunkAlignment = 128; -inline constexpr std::uint32_t kMaximumMtpDraftTokens = 5; -inline constexpr std::uint32_t kMaximumDFlashDraftTokens = 0; -inline constexpr std::uint32_t kNativeContext = 262144; - -} // namespace ninfer::targets::qwen3_6_27b::detail +#pragma once + +#include +#include +#include + +#include +#include + +namespace ninfer::targets::qwen3_6_27b::detail { + +struct TextConfig { + static constexpr int hidden = 5120; + static constexpr int layers = 64; + static constexpr int intermediate = 17408; + + // The output matrix is padded for the selected kernels. Only token IDs in + // [0, token_domain) are tokenizer-addressable and valid sampling results. + static constexpr int output_rows = 248320; + static constexpr int token_domain = static_cast(qwen3_6::kTokenDomain); + + static constexpr int gdn_conv_kernel = 4; + static constexpr int gdn_conv_state_width = gdn_conv_kernel - 1; + static constexpr int gdn_key_heads = 16; + static constexpr int gdn_key_head_dim = 128; + static constexpr int gdn_value_heads = 48; + static constexpr int gdn_value_head_dim = 128; + + static constexpr int query_heads = 24; + static constexpr int kv_heads = 4; + static constexpr int head_dim = 256; + static constexpr int rotary_dim = 64; + + static constexpr int full_attention_interval = qwen3_6::kHybridAttentionInterval; + static constexpr float rms_epsilon = 1.0e-6F; + static constexpr float rope_theta = 1.0e7F; + + static constexpr int key_dim = gdn_key_heads * gdn_key_head_dim; + static constexpr int value_dim = gdn_value_heads * gdn_value_head_dim; + static constexpr int convolution_dim = 2 * key_dim + value_dim; + static constexpr int query_size = query_heads * head_dim; + static constexpr int kv_size = kv_heads * head_dim; + static constexpr int query_projection_rows = 2 * query_size; + + static constexpr int mtp_layers = 1; + static constexpr int mtp_input_rows = 2 * hidden; + static constexpr int mtp_attention_input_rows = 2 * query_size + 2 * kv_size; + static constexpr int mtp_mlp_gate_up_rows = 2 * intermediate; + + [[nodiscard]] static constexpr bool is_full_attention(int layer) { + return qwen3_6::is_full_attention_layer(layer); + } + + [[nodiscard]] static constexpr int full_attention_layers() { + return qwen3_6::full_attention_layers(layers); + } + + [[nodiscard]] static constexpr int gdn_layers() { return qwen3_6::gdn_layers(layers); } + + [[nodiscard]] static constexpr int full_attention_index(int layer) { + return qwen3_6::full_attention_index(layer); + } + + [[nodiscard]] static constexpr int gdn_index(int layer) { return qwen3_6::gdn_index(layer); } +}; + +static_assert(TextConfig::full_attention_layers() == 16); +static_assert(TextConfig::gdn_layers() == 48); + +struct VisionConfig : qwen3_6::VisionBackboneConfig { + static constexpr int output_hidden = TextConfig::hidden; +}; + +struct DFlashConfig { + static constexpr bool supported = true; + // DSpark draft weights ship as BF16 (not W8 like the 35B DFlash draft), + // so the runtime uses the generic BF16 MMA GEMM path for projections. + static constexpr bool bf16_weights = true; + static constexpr int layers = 5; + // DSpark has no sliding-window local layers: every layer uses the full + // dual-source (target-feature context + noise rows) attention path. + // A single legacy local cyclic layer is still allocated because the + // cyclic-cache layout requires layers >= 1; full_only keeps it unused. + static constexpr bool full_only = true; + static constexpr int local_layers = 1; + static constexpr int full_layers = full_only ? layers : layers - local_layers; + static constexpr int feature_layers = 5; + static constexpr int feature_rows = feature_layers * TextConfig::hidden; + static constexpr int hidden = TextConfig::hidden; + static constexpr int intermediate = 10240; + static constexpr int query_heads = 40; + static constexpr int kv_heads = 8; + static constexpr int head_dim = 128; + static constexpr int query_size = query_heads * head_dim; + static constexpr int kv_size = kv_heads * head_dim; + static constexpr int local_capacity = 4096; + static constexpr std::uint32_t local_window = 4096; + static constexpr int mask_token = 248077; + static constexpr float rms_epsilon = 1.0e-6F; + static constexpr float rope_theta = 1.0e7F; + static constexpr float attention_scale = 0.08838834764831845F; + // SVIP self-verification length policy: stop drafting after the first + // position whose base-logit entropy exceeds threshold^2. + static constexpr float svip_entropy_threshold = 2.5F; + static constexpr std::array target_feature_layers{4, 16, 28, 40, 52}; +}; + +// DFlash2 block-diffusion draft: five sliding-window (2048) non-causal layers, +// grouped convolutions, and a rank-256 candidate selector with a 16-way path +// walk. The block always drafts seven tokens after the bonus token. +struct DFlash2Config { + static constexpr bool supported = true; + static constexpr bool bf16_weights = true; + static constexpr int layers = 5; + static constexpr bool full_only = false; + static constexpr int local_layers = 5; + static constexpr int full_layers = full_only ? layers : layers - local_layers; + static constexpr int feature_layers = 5; + static constexpr int feature_rows = feature_layers * TextConfig::hidden; + static constexpr int hidden = TextConfig::hidden; + static constexpr int intermediate = 17408; + static constexpr int query_heads = 32; + static constexpr int kv_heads = 8; + static constexpr int head_dim = 128; + static constexpr int query_size = query_heads * head_dim; + static constexpr int kv_size = kv_heads * head_dim; + static constexpr int local_capacity = 2048; + static constexpr std::uint32_t local_window = 2048; + static constexpr int mask_token = 248070; + static constexpr float rms_epsilon = 1.0e-6F; + static constexpr float rope_theta = 1.0e7F; + static constexpr float attention_scale = 0.08838834764831845F; + static constexpr int block_drafts = 7; + static constexpr int conv_group_size = 16; + static constexpr int conv_kernel_size = 2; + static constexpr int selector_rank = 256; + static constexpr int selector_top_k = 16; + static constexpr std::array target_feature_layers{5, 19, 33, 47, 61}; +}; + +inline constexpr float kAttentionScale = 0.0625F; +inline constexpr float kGdnScale = 0.08838834764831845F; +inline constexpr std::uint32_t kPrefillChunkAlignment = 128; +inline constexpr std::uint32_t kMaximumMtpDraftTokens = 5; +inline constexpr std::uint32_t kMaximumDFlashDraftTokens = 7; +inline constexpr std::uint32_t kNativeContext = 262144; + +} // namespace ninfer::targets::qwen3_6_27b::detail diff --git a/src/targets/qwen3_6_27b/impl/load/bindings.cpp b/src/targets/qwen3_6_27b/impl/load/bindings.cpp index f942a7976e..267245a9a0 100644 --- a/src/targets/qwen3_6_27b/impl/load/bindings.cpp +++ b/src/targets/qwen3_6_27b/impl/load/bindings.cpp @@ -40,6 +40,7 @@ NumericFormat endpoint_format(WeightsProfile weights_profile) { case WeightsProfile::Qwen36Nvfp4: return NumericFormat::W8G32_F16S; case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return NumericFormat::FP8_E4M3FN_ROW_BF16S; } throw std::invalid_argument("qwen3_6_27b: invalid weights profile"); @@ -431,6 +432,9 @@ ArtifactLoadPlan bind_artifact(artifact::Binder& binder, WeightsProfile weights_ case WeightsProfile::Qwen38Nvfp4: bind_qwen38_nvfp4_text_layers(binder, out); break; + case WeightsProfile::Qwen38Nvfp4DFlash2: + bind_qwen38_nvfp4_text_layers(binder, out); + break; default: throw std::invalid_argument("qwen3_6_27b: invalid weights profile"); } @@ -474,6 +478,66 @@ ArtifactLoadPlan bind_artifact(artifact::Binder& binder, WeightsProfile weights_ .format = NumericFormat::W8G32_F16S}; out.mtp.final_norm = bind_mtp("mtp/final_norm", NumericFormat::BF16, {5120}); + if (weights_profile == WeightsProfile::Qwen38Nvfp4DFlash2) { + const artifact::TensorPlacement dflash2_placement = + features.dflash2() ? artifact::TensorPlacement::Device + : artifact::TensorPlacement::ValidateOnly; + const auto bind_dflash2 = [&](std::string_view name, NumericFormat format, + std::initializer_list shape) { + return artifact::bind_tensor(binder, name, format, shape, dflash2_placement); + }; + const auto bind_dflash2_weight = [&](std::string_view name, NumericFormat format, + std::initializer_list shape) { + return WeightPlan{.object = bind_dflash2(name, format, shape), .format = format}; + }; + out.dflash2.feature_projection = bind_dflash2_weight("dflash2/feature_projection", + NumericFormat::BF16, {5120, 25600}); + out.dflash2.context_norm = + bind_dflash2("dflash2/context_norm", NumericFormat::BF16, {5120}); + for (std::size_t layer = 0; layer < DFlash2Config::layers; ++layer) { + DFlash2LayerPlan& target = out.dflash2.layers[layer]; + const std::string prefix = "dflash2/layers/" + std::to_string(layer) + "/"; + target.input_norm = bind_dflash2(prefix + "input_norm", NumericFormat::BF16, {5120}); + target.query_key_value = bind_dflash2_weight( + prefix + "attention/query_key_value", NumericFormat::BF16, {6144, 5120}); + target.context_key = bind_dflash2_weight(prefix + "attention/context_key", + NumericFormat::BF16, {1024, 5120}); + target.context_value = bind_dflash2_weight(prefix + "attention/context_value", + NumericFormat::BF16, {1024, 5120}); + target.query_norm = + bind_dflash2(prefix + "attention/query_norm", NumericFormat::BF16, {128}); + target.key_norm = + bind_dflash2(prefix + "attention/key_norm", NumericFormat::BF16, {128}); + target.attention_output = bind_dflash2_weight(prefix + "attention/output", + NumericFormat::BF16, {5120, 4096}); + target.attention_conv_base = WeightPlan{ + .object = bind_dflash2(prefix + "attention_conv/base_kernel", NumericFormat::BF16, + {2, 2, 5120}), + .format = NumericFormat::BF16}; + target.attention_conv_projection = bind_dflash2_weight( + prefix + "attention_conv/kernel_projection", NumericFormat::BF16, {1280, 5120}); + target.post_attention_norm = + bind_dflash2(prefix + "post_attention_norm", NumericFormat::BF16, {5120}); + target.mlp.gate_up = bind_dflash2_weight(prefix + "mlp/gate_up", NumericFormat::BF16, + {34816, 5120}); + target.mlp.down = + bind_dflash2_weight(prefix + "mlp/down", NumericFormat::BF16, {5120, 17408}); + target.mlp_conv_base = WeightPlan{ + .object = bind_dflash2(prefix + "mlp_conv/base_kernel", NumericFormat::BF16, + {2, 2, 5120}), + .format = NumericFormat::BF16}; + target.mlp_conv_projection = bind_dflash2_weight( + prefix + "mlp_conv/kernel_projection", NumericFormat::BF16, {1280, 5120}); + } + out.dflash2.final_norm = bind_dflash2("dflash2/final_norm", NumericFormat::BF16, {5120}); + out.dflash2.selector_hidden_projection = bind_dflash2_weight( + "dflash2/candidate_selector/hidden_projection", NumericFormat::BF16, {256, 5120}); + out.dflash2.selector_predecessor_codebook = bind_dflash2_weight( + "dflash2/candidate_selector/predecessor_codebook", NumericFormat::BF16, {248320, 256}); + out.dflash2.selector_successor_codebook = bind_dflash2_weight( + "dflash2/candidate_selector/successor_codebook", NumericFormat::BF16, {248320, 256}); + } + const artifact::TensorPlacement vision_placement = features.vision ? artifact::TensorPlacement::Device : artifact::TensorPlacement::ValidateOnly; @@ -582,6 +646,50 @@ LoadedModelData::LoadedModelData(BindingPlan plan, artifact::MaterializedArtifac NumericFormat::BF16, {5120}); } + if (plan.features.dflash2()) { + DFlash2Weights& target = runtime.dflash2.emplace(); + target.feature_projection = materialized_weight(backing, plan.dflash2.feature_projection, + 5120, 25600); + target.context_norm = artifact::materialized_tensor(backing, plan.dflash2.context_norm, + NumericFormat::BF16, {5120}); + for (std::size_t layer = 0; layer < DFlash2Config::layers; ++layer) { + const DFlash2LayerPlan& source = plan.dflash2.layers[layer]; + DFlash2LayerWeights& weights = target.layers[layer]; + weights.input_norm = artifact::materialized_tensor(backing, source.input_norm, + NumericFormat::BF16, {5120}); + weights.query_key_value = materialized_weight(backing, source.query_key_value, 6144, 5120); + weights.context_key = materialized_weight(backing, source.context_key, 1024, 5120); + weights.context_value = + materialized_weight(backing, source.context_value, 1024, 5120); + weights.query_norm = artifact::materialized_tensor(backing, source.query_norm, + NumericFormat::BF16, {128}); + weights.key_norm = + artifact::materialized_tensor(backing, source.key_norm, NumericFormat::BF16, {128}); + weights.attention_output = + materialized_weight(backing, source.attention_output, 5120, 4096); + weights.attention_conv_base = artifact::materialized_weight( + backing, source.attention_conv_base.object, NumericFormat::BF16, 4, 5120); + weights.attention_conv_projection = + materialized_weight(backing, source.attention_conv_projection, 1280, 5120); + weights.post_attention_norm = artifact::materialized_tensor( + backing, source.post_attention_norm, NumericFormat::BF16, {5120}); + weights.gate_up = materialized_weight(backing, source.mlp.gate_up, 34816, 5120); + weights.down = materialized_weight(backing, source.mlp.down, 5120, 17408); + weights.mlp_conv_base = artifact::materialized_weight( + backing, source.mlp_conv_base.object, NumericFormat::BF16, 4, 5120); + weights.mlp_conv_projection = + materialized_weight(backing, source.mlp_conv_projection, 1280, 5120); + } + target.final_norm = artifact::materialized_tensor(backing, plan.dflash2.final_norm, + NumericFormat::BF16, {5120}); + target.selector_hidden_projection = materialized_weight( + backing, plan.dflash2.selector_hidden_projection, 256, 5120); + target.selector_predecessor_codebook = materialized_weight( + backing, plan.dflash2.selector_predecessor_codebook, 248320, 256); + target.selector_successor_codebook = materialized_weight( + backing, plan.dflash2.selector_successor_codebook, 248320, 256); + } + if (plan.features.vision) { auto& vision = runtime.vision.emplace(); vision.common = qwen3_6::materialize_vision_common( diff --git a/src/targets/qwen3_6_27b/impl/load/bindings.h b/src/targets/qwen3_6_27b/impl/load/bindings.h index 70a14140fe..85a0394683 100644 --- a/src/targets/qwen3_6_27b/impl/load/bindings.h +++ b/src/targets/qwen3_6_27b/impl/load/bindings.h @@ -1,5 +1,6 @@ #pragma once +#include "targets/qwen3_6_27b/impl/config.h" #include #include #include @@ -104,6 +105,32 @@ struct MtpPlan { artifact::ObjectHandle final_norm; }; +struct DFlash2LayerPlan { + artifact::ObjectHandle input_norm; + WeightPlan query_key_value; + WeightPlan context_key; + WeightPlan context_value; + artifact::ObjectHandle query_norm; + artifact::ObjectHandle key_norm; + WeightPlan attention_output; + WeightPlan attention_conv_base; + WeightPlan attention_conv_projection; + artifact::ObjectHandle post_attention_norm; + MlpPlan mlp; + WeightPlan mlp_conv_base; + WeightPlan mlp_conv_projection; +}; + +struct DFlash2Plan { + WeightPlan feature_projection; + artifact::ObjectHandle context_norm; + std::array layers; + artifact::ObjectHandle final_norm; + WeightPlan selector_hidden_projection; + WeightPlan selector_predecessor_codebook; + WeightPlan selector_successor_codebook; +}; + struct BindingPlan { qwen3_6::FrontendResourcePlan frontend; qwen3_6::StartupFeatures features; @@ -115,6 +142,7 @@ struct BindingPlan { artifact::ObjectHandle draft_head; artifact::ObjectHandle draft_head_token_ids; MtpPlan mtp; + DFlash2Plan dflash2; qwen3_6::VisionBackbonePlan vision_backbone; qwen3_6::VisionMergerInputPlan vision_merger_input; @@ -190,10 +218,13 @@ struct MtpAttentionPayload { using RuntimeModelView = qwen3_6::ModelView, - kFullAttentionLayers, kGdnLayers>; + qwen3_6::DFlash2Weights, kFullAttentionLayers, + kGdnLayers>; using FullAttentionWeights = RuntimeModelView::FullLayer; using GdnWeights = RuntimeModelView::GdnLayer; using MtpWeights = RuntimeModelView::MtpLayer; +using DFlash2Weights = RuntimeModelView::DFlash2; +using DFlash2LayerWeights = qwen3_6::DFlash2LayerWeights; class LoadedModelData { public: diff --git a/src/targets/qwen3_6_27b/impl/package.cpp b/src/targets/qwen3_6_27b/impl/package.cpp index 67647994fe..9076657ce6 100644 --- a/src/targets/qwen3_6_27b/impl/package.cpp +++ b/src/targets/qwen3_6_27b/impl/package.cpp @@ -95,15 +95,48 @@ Package::WeightsProfile Package::resolve_weights(const artifact::ArtifactIdentit if (identity.model_id == qwen3_8_model_id && identity.weights_id == "nvfp4") { return WeightsProfile::Qwen38Nvfp4; } + if (identity.model_id == qwen3_8_model_id && identity.weights_id == "nvfp4-dflash2") { + return WeightsProfile::Qwen38Nvfp4DFlash2; + } throw std::runtime_error("artifact identity '" + identity.model_id + "/" + identity.weights_id + "' is not supported by target '" + std::string(target_key) + "'"); } +namespace { +EngineOptions resolved_auto_speculative(const EngineOptions& options, + detail::WeightsProfile weights_profile) { + EngineOptions resolved = options; + if (options.speculative.backend != SpeculativeBackend::Auto) { return resolved; } + const DType kv_dtype = + options.kv_cache == KvCacheStorage::BFloat16 + ? DType::BF16 + : (options.kv_cache == KvCacheStorage::Int8Group64 + ? DType::I8 + : (options.kv_cache == KvCacheStorage::Fp8E4M3Row256 + ? DType::FP8_E4M3FN + : DType::BF16)); + const std::uint32_t draft_capacity = + kv_dtype == DType::FP8_E4M3FN ? 8192U : (kv_dtype == DType::I8 ? 4096U : 2048U); + // The artifact's frozen startup features enforce this limit; the engine + // check makes the fallback graceful (selects MTP) instead of a load error. + if (weights_profile == detail::WeightsProfile::Qwen38Nvfp4DFlash2 && !options.enable_vision && + options.max_context <= draft_capacity) { + resolved.speculative.backend = SpeculativeBackend::DFlash2; + if (resolved.speculative.draft_tokens == 0) { resolved.speculative.draft_tokens = 7; } + } else { + resolved.speculative.backend = SpeculativeBackend::Mtp; + if (resolved.speculative.draft_tokens == 0) { resolved.speculative.draft_tokens = 3; } + } + return resolved; +} +} // namespace + Package::LoadPlan Package::plan_load(artifact::Binder& binder, const EngineOptions& options, WeightsProfile weights_profile) { + const EngineOptions resolved = resolved_auto_speculative(options, weights_profile); return LoadPlan(std::make_unique( weights_profile, - detail::bind_artifact(binder, weights_profile, qwen3_6::startup_features(options)))); + detail::bind_artifact(binder, weights_profile, qwen3_6::startup_features(resolved)))); } std::unique_ptr diff --git a/src/targets/qwen3_6_27b/impl/variant.cpp b/src/targets/qwen3_6_27b/impl/variant.cpp index c2036d9145..8634e090b9 100644 --- a/src/targets/qwen3_6_27b/impl/variant.cpp +++ b/src/targets/qwen3_6_27b/impl/variant.cpp @@ -113,6 +113,37 @@ std::size_t post_mixer_workspace_bytes(QType gate_up_qtype, QType down_qtype, return layout.peak_bytes(1); } +std::vector +dflash_base_profiles(std::uint32_t capacity, std::uint32_t draft_window) { + if (draft_window == 0 || capacity == 0) { return {}; } + const std::uint32_t block = draft_window + 1; + const std::uint32_t max_frontier = capacity - 1; + std::vector ends{ + 96U, 127U, 511U, 1023U, 2047U, 4095U, 8191U, 16383U, 32767U, 65536U, 131072U, 196608U, + }; + const auto add_target_boundary = [&](std::uint32_t visible_end) { + if (visible_end >= block) { ends.push_back(visible_end - block); } + }; + for (const std::uint32_t visible_end : {128U, 512U, 2048U, 4096U, 8198U, 16390U, 32768U}) { + add_target_boundary(visible_end); + } + if (draft_window >= 6 && draft_window <= 15) { + add_target_boundary(draft_window <= 11 ? 512U : 1024U); + } + std::sort(ends.begin(), ends.end()); + ends.erase(std::unique(ends.begin(), ends.end()), ends.end()); + return graph_profiles_through(max_frontier, ends); +} + +bool dflash_target_uses_chunked_small_t(std::uint32_t draft_window, std::uint32_t batch_size, + std::uint32_t max_visible_keys) { + const std::uint32_t tokens = draft_window + 1; + if (tokens <= 6) { return false; } + if (batch_size > 1) { return true; } + const std::uint32_t prompt_visible_limit = tokens <= 12 ? 512U : 1024U; + return max_visible_keys > prompt_visible_limit; +} + } // namespace std::vector Variant::ordinary_graph_profiles(std::uint32_t capacity) { @@ -155,6 +186,21 @@ std::vector Variant::dflash_graph_profiles(std::uint32_t, return {}; } +std::vector Variant::dflash2_graph_profiles(std::uint32_t capacity, + std::uint32_t draft_window, + std::uint32_t batch_size) { + std::vector profiles = dflash_base_profiles(capacity, draft_window); + for (GraphExecutionProfile& profile : profiles) { + const std::uint32_t target_max = static_cast(std::min( + capacity, static_cast(profile.max) + draft_window + 1ULL)); + const bool split_swa = profile.max > 96U; + const bool chunked_target = + dflash_target_uses_chunked_small_t(draft_window, batch_size, target_max); + profile.topology_class = (chunked_target ? 2U : 0U) | (split_swa ? 1U : 0U); + } + return profiles; +} + void Variant::attention_projection(const Tensor& hidden, const FullAttentionProjectionWeights& weights, Tensor& query, Tensor& gate, Tensor& key, Tensor& value, qwen3_6::TextPhase, @@ -349,6 +395,7 @@ std::size_t Variant::attention_projection_workspace_capacity_bytes(WeightsProfil return ops::attn_input_proj_workspace_capacity_bytes( QType::NVFP4, 14336, TextConfig::hidden, kNvfp4TextPolicy, first, last); case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return ops::attn_input_proj_workspace_capacity_bytes( QType::FP8_E4M3FN_ROW_BF16S, 14336, TextConfig::hidden, kFp8TextPolicy, first, last); } @@ -369,6 +416,7 @@ std::size_t Variant::attention_output_projection_workspace_capacity_bytes( TextConfig::query_size, kNvfp4TextPolicy, first, last); case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return ops::linear_add_workspace_capacity_bytes(QType::FP8_E4M3FN_ROW_BF16S, TextConfig::hidden, TextConfig::query_size, kFp8TextPolicy, first, last); @@ -389,6 +437,7 @@ std::size_t Variant::gdn_input_projection_workspace_capacity_bytes(WeightsProfil return ops::gdn_input_proj_workspace_capacity_bytes(QType::NVFP4, 16384, TextConfig::hidden, kNvfp4TextPolicy, first, last); case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return ops::gdn_input_proj_workspace_capacity_bytes( QType::FP8_E4M3FN_ROW_BF16S, 16384, TextConfig::hidden, kFp8TextPolicy, first, last); } @@ -412,6 +461,7 @@ std::size_t Variant::gdn_input_projection_snapshot_workspace_capacity_bytes( QType::NVFP4, 16384, TextConfig::hidden, kNvfp4TextPolicy, batch_size, first, last)); case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return std::max(kMinimumLeafWorkspaceBytes, ops::gdn_input_proj_conv_snapshot_workspace_capacity_bytes( QType::FP8_E4M3FN_ROW_BF16S, 16384, TextConfig::hidden, kFp8TextPolicy, @@ -437,6 +487,7 @@ std::size_t Variant::gdn_input_projection_record_workspace_capacity_bytes( QType::NVFP4, 16384, TextConfig::hidden, kNvfp4TextPolicy, batch_size, first, last)); case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return std::max(kMinimumLeafWorkspaceBytes, ops::gdn_input_proj_conv_record_workspace_capacity_bytes( QType::FP8_E4M3FN_ROW_BF16S, 16384, TextConfig::hidden, kFp8TextPolicy, @@ -460,6 +511,7 @@ std::size_t Variant::gdn_output_projection_workspace_capacity_bytes(WeightsProfi return ops::linear_add_workspace_capacity_bytes( QType::NVFP4, TextConfig::hidden, TextConfig::value_dim, kNvfp4TextPolicy, first, last); case WeightsProfile::Qwen38Nvfp4: + case WeightsProfile::Qwen38Nvfp4DFlash2: return ops::linear_add_workspace_capacity_bytes(QType::FP8_E4M3FN_ROW_BF16S, TextConfig::hidden, TextConfig::value_dim, kFp8TextPolicy, first, last); diff --git a/src/targets/qwen3_6_27b/impl/variant.h b/src/targets/qwen3_6_27b/impl/variant.h index 75332671ad..353d18fc54 100644 --- a/src/targets/qwen3_6_27b/impl/variant.h +++ b/src/targets/qwen3_6_27b/impl/variant.h @@ -19,6 +19,7 @@ struct Variant { using TextConfig = detail::TextConfig; using VisionConfig = detail::VisionConfig; using DFlashConfig = detail::DFlashConfig; + using DFlash2Config = detail::DFlash2Config; using ModelView = detail::RuntimeModelView; using FullAttentionProjectionWeights = detail::FullAttentionProjectionPayload; using GdnProjectionWeights = detail::GdnProjectionPayload; @@ -35,8 +36,13 @@ struct Variant { static constexpr std::uint32_t maximum_dflash_draft_tokens = kMaximumDFlashDraftTokens; static constexpr std::uint32_t maximum_context = kNativeContext; static constexpr bool supports_dflash = DFlashConfig::supported; + static constexpr bool supports_dflash2 = DFlash2Config::supported; static constexpr std::int32_t draft_head_rows = 131072; + [[nodiscard]] static constexpr bool dflash2_weights(WeightsProfile profile) { + return profile == WeightsProfile::Qwen38Nvfp4DFlash2; + } + static void attention_projection(const Tensor& hidden, const FullAttentionProjectionWeights& weights, Tensor& query, Tensor& gate, Tensor& key, Tensor& value, @@ -125,6 +131,9 @@ struct Variant { [[nodiscard]] static std::vector dflash_graph_profiles(std::uint32_t capacity, std::uint32_t draft_window, std::uint32_t batch_size); + [[nodiscard]] static std::vector + dflash2_graph_profiles(std::uint32_t capacity, std::uint32_t draft_window, + std::uint32_t batch_size); }; } // namespace ninfer::targets::qwen3_6_27b::detail diff --git a/src/targets/qwen3_6_35b_a3b/impl/config.h b/src/targets/qwen3_6_35b_a3b/impl/config.h index 7429be4744..af530d08c5 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/config.h +++ b/src/targets/qwen3_6_35b_a3b/impl/config.h @@ -70,6 +70,7 @@ struct VisionConfig : qwen3_6::VisionBackboneConfig { struct DFlashConfig { static constexpr bool supported = true; + static constexpr bool bf16_weights = false; static constexpr int layers = 6; static constexpr int local_layers = 5; static constexpr int feature_layers = 8; @@ -90,6 +91,39 @@ struct DFlashConfig { 22, 27, 32, 37}; }; +// The 35B-A3B package registers no DFlash2 draft; the family runtime still +// needs a concrete config type for its compile-time aliases and recipe +// instantiations, so it mirrors the 27B DFlash2 shape. +struct DFlash2Config { + static constexpr bool supported = false; + static constexpr bool bf16_weights = true; + static constexpr int layers = 5; + static constexpr bool full_only = false; + static constexpr int local_layers = 5; + static constexpr int full_layers = 0; + static constexpr int feature_layers = 5; + static constexpr int feature_rows = feature_layers * TextConfig::hidden; + static constexpr int hidden = TextConfig::hidden; + static constexpr int intermediate = 17408; + static constexpr int query_heads = 32; + static constexpr int kv_heads = 8; + static constexpr int head_dim = 128; + static constexpr int query_size = query_heads * head_dim; + static constexpr int kv_size = kv_heads * head_dim; + static constexpr int local_capacity = 2048; + static constexpr std::uint32_t local_window = 2048; + static constexpr int mask_token = 248070; + static constexpr float rms_epsilon = 1.0e-6F; + static constexpr float rope_theta = 1.0e7F; + static constexpr float attention_scale = 0.08838834764831845F; + static constexpr int block_drafts = 7; + static constexpr int conv_group_size = 16; + static constexpr int conv_kernel_size = 2; + static constexpr int selector_rank = 256; + static constexpr int selector_top_k = 16; + static constexpr std::array target_feature_layers{5, 19, 33, 47, 61}; +}; + inline constexpr float kAttentionScale = 0.0625F; inline constexpr float kGdnScale = 0.08838834764831845F; inline constexpr std::uint32_t kPrefillChunkAlignment = 128; diff --git a/src/targets/qwen3_6_35b_a3b/impl/load/bindings.h b/src/targets/qwen3_6_35b_a3b/impl/load/bindings.h index 6015467b7c..9560558a9b 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/load/bindings.h +++ b/src/targets/qwen3_6_35b_a3b/impl/load/bindings.h @@ -129,7 +129,7 @@ struct GdnProjectionPayload { using RuntimeModelView = qwen3_6::ModelView, kFullAttentionLayers, kGdnLayers>; + qwen3_6::DFlashWeights, qwen3_6::DFlash2Weights<0>, kFullAttentionLayers, kGdnLayers>; using FullAttentionWeights = RuntimeModelView::FullLayer; using GdnWeights = RuntimeModelView::GdnLayer; using MtpWeights = RuntimeModelView::MtpLayer; diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp index 96ab3ac8e6..ac611cab74 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp @@ -123,6 +123,10 @@ std::vector Variant::dflash_graph_profiles(std::uint32_t } return profiles; } +std::vector Variant::dflash2_graph_profiles(std::uint32_t, std::uint32_t, std::uint32_t) { + return {}; +} + void Variant::attention_projection(const Tensor& hidden, const FullAttentionProjectionWeights& weights, Tensor& query, diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.h b/src/targets/qwen3_6_35b_a3b/impl/variant.h index c6802e43f4..4c85ccce8f 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.h +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.h @@ -17,6 +17,7 @@ struct Variant { using TextConfig = detail::TextConfig; using VisionConfig = detail::VisionConfig; using DFlashConfig = detail::DFlashConfig; + using DFlash2Config = detail::DFlash2Config; using ModelView = detail::RuntimeModelView; using FullAttentionProjectionWeights = detail::AttentionProjectionPayload; using GdnProjectionWeights = detail::GdnProjectionPayload; @@ -33,6 +34,7 @@ struct Variant { static constexpr std::uint32_t maximum_dflash_draft_tokens = kMaximumDFlashDraftTokens; static constexpr std::uint32_t maximum_context = kNativeContext; static constexpr bool supports_dflash = DFlashConfig::supported; + static constexpr bool supports_dflash2 = DFlash2Config::supported; static constexpr std::int32_t draft_head_rows = 131072; [[nodiscard]] static std::vector @@ -42,6 +44,9 @@ struct Variant { [[nodiscard]] static std::vector dflash_graph_profiles(std::uint32_t capacity, std::uint32_t draft_window, std::uint32_t batch_size); + [[nodiscard]] static std::vector + dflash2_graph_profiles(std::uint32_t capacity, std::uint32_t draft_window, + std::uint32_t batch_size); static void attention_projection(const Tensor& hidden, const FullAttentionProjectionWeights& weights, Tensor& query, From e97e57114db3749a18cf725d6961f5281c3c555c Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 08:55:29 +0800 Subject: [PATCH 3/4] fix(spec): DFlash2 smoke-test fixes --- src/ops/linear/bf16/bf16_dispatch.cpp | 15 ++++++- src/ops/linear/bf16/bf16_gemm_mma.cu | 36 +++++++++++++++++ src/serve/generation_service.cpp | 16 ++++++-- .../qwen3_6/impl/runtime/dflash2_impl.h | 4 ++ .../qwen3_6/impl/runtime/program_impl.h | 40 +++++++++++-------- src/targets/qwen3_6/impl/runtime/schedule.h | 3 +- .../qwen3_6/impl/runtime/text_prefill_impl.h | 18 +++++---- src/targets/qwen3_6_27b/impl/variant.cpp | 1 + 8 files changed, 103 insertions(+), 30 deletions(-) diff --git a/src/ops/linear/bf16/bf16_dispatch.cpp b/src/ops/linear/bf16/bf16_dispatch.cpp index da0c705a8b..6170648087 100644 --- a/src/ops/linear/bf16/bf16_dispatch.cpp +++ b/src/ops/linear/bf16/bf16_dispatch.cpp @@ -9,10 +9,21 @@ namespace ninfer::ops::detail { Bf16Launch select_bf16_a16_launch(std::int32_t n, std::int32_t k, std::int32_t t) { - const bool supported_problem = (n == 14336 && k == 5120) || (n == 5120 && k == 6144); - if (!supported_problem || t <= 0) { + const bool legacy_problem = (n == 14336 && k == 5120) || (n == 5120 && k == 6144); + const bool dflash2_problem = (n == 5120 && k == 25600) || (n == 6144 && k == 5120) || + (n == 4096 && k == 5120) || (n == 1024 && k == 5120) || + (n == 5120 && k == 4096) || (n == 1280 && k == 5120) || + (n == 34816 && k == 5120) || (n == 5120 && k == 17408) || + (n == 256 && k == 5120); + if ((!legacy_problem && !dflash2_problem) || t <= 0) { throw std::invalid_argument("bf16 linear: unsupported shape or T"); } + if (dflash2_problem) { + // The generic MMA core already tiles arbitrary admitted n/k with the + // 64x128x64 production schedule; the DFlash2 draft never needs the + // fixed-shape decode/small-T specializations (T = verify width 8..64). + return launch_bf16_mma; + } if (t == 1) { return launch_bf16_decode; } const std::int32_t small_t_end = n == 5120 ? kBf16SmallTMaxTokens : kBf16LinearSmallTDispatchEnd; diff --git a/src/ops/linear/bf16/bf16_gemm_mma.cu b/src/ops/linear/bf16/bf16_gemm_mma.cu index 33df64ab2e..f4a84d336e 100644 --- a/src/ops/linear/bf16/bf16_gemm_mma.cu +++ b/src/ops/linear/bf16/bf16_gemm_mma.cu @@ -57,6 +57,42 @@ void launch_bf16_mma(const Tensor& x, const Weight& weight, Tensor& out, cudaStr launch_geometry>(x, weight, out, stream); return; } + if (weight.n == 5120 && weight.k == 25600) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 6144 && weight.k == 5120) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 4096 && weight.k == 5120) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 1024 && weight.k == 5120) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 5120 && weight.k == 4096) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 1280 && weight.k == 5120) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 34816 && weight.k == 5120) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 5120 && weight.k == 17408) { + launch_geometry>(x, weight, out, stream); + return; + } + if (weight.n == 256 && weight.k == 5120) { + launch_geometry>(x, weight, out, stream); + return; + } throw std::invalid_argument("bf16 linear MMA: unsupported exact problem"); } diff --git a/src/serve/generation_service.cpp b/src/serve/generation_service.cpp index 08d0bf7227..6627d61bb5 100644 --- a/src/serve/generation_service.cpp +++ b/src/serve/generation_service.cpp @@ -1,6 +1,7 @@ #include "serve/generation_service.h" #include "product/media_acquire/acquire.h" +#include "serve/console_log.h" #include "serve/translate.h" #include @@ -450,10 +451,17 @@ void GenerationService::warmup() { turn.content.push_back(std::move(content)); request.messages.push_back(std::move(turn)); request.max_tokens = 4; - PreparedRequest prepared = - prepare_impl(request, GenerationConsumerMode::Aggregate, {}, {}, - CacheParticipation::Disabled, DeadlinePolicy::UnboundedStartup); - run(prepared, nullptr); + // The warmup is advisory: a failed round must not take the server down, + // the first real request surfaces the real error instead. + try { + PreparedRequest prepared = + prepare_impl(request, GenerationConsumerMode::Aggregate, {}, {}, + CacheParticipation::Disabled, DeadlinePolicy::UnboundedStartup); + run(prepared, nullptr); + } catch (const std::exception& exception) { + write_console_log(ConsoleLogLevel::Warning, + std::string("warmup failed (continuing): ") + exception.what()); + } } } // namespace ninfer::serve diff --git a/src/targets/qwen3_6/impl/runtime/dflash2_impl.h b/src/targets/qwen3_6/impl/runtime/dflash2_impl.h index 9ff20f574d..5df847ab29 100644 --- a/src/targets/qwen3_6/impl/runtime/dflash2_impl.h +++ b/src/targets/qwen3_6/impl/runtime/dflash2_impl.h @@ -365,6 +365,8 @@ auto dflash2_decode_batch_body(DFlash2BatchContext& state, std::int32_t batch_si Tensor valid_columns = frame.target_valid_columns.slice(0, 0, batch_size); Tensor text_rows = frame.text_kv_table_rows.slice(0, 0, batch_size); Tensor lanes = frame.active_lanes.slice(0, 0, batch_size); + Tensor state_sources = frame.state_source_slots.slice(0, 0, batch_size); + Tensor state_destinations = frame.state_destination_slots.slice(0, 0, batch_size); Tensor append_positions = frame.append_positions.slice(1, 0, batch_size); Tensor append_counts = frame.append_counts.slice(0, 0, batch_size); Tensor drafts = frame.draft_tokens.slice(1, 0, batch_size); @@ -406,6 +408,8 @@ auto dflash2_decode_batch_body(DFlash2BatchContext& state, std::int32_t batch_si .rope_positions = target_positions, .valid_columns = valid_columns, .kv_table_rows = text_rows, + .state_source_slots = state_sources, + .state_destination_slots = state_destinations, .target_hidden = target_hidden, .target_logits = target_logits, .target_tokens = target_tokens, diff --git a/src/targets/qwen3_6/impl/runtime/program_impl.h b/src/targets/qwen3_6/impl/runtime/program_impl.h index 713e5d0416..a4ef0abeeb 100644 --- a/src/targets/qwen3_6/impl/runtime/program_impl.h +++ b/src/targets/qwen3_6/impl/runtime/program_impl.h @@ -8536,7 +8536,8 @@ runtime::ExecutionTiming ProgramImplCore::append_forced_tokens( selectors.source, selectors.destination, 0, - dflash_host_ingress}; + dflash_host_ingress, + dflash2_host_ingress}; mark_workspace_usage(speculative_backend == SpeculativeBackend::Mtp ? workspace_plan.mtp_prefill : workspace_plan.text_prefill); @@ -9467,7 +9468,7 @@ void ProgramImplCore::start_sequence(std::uint32_t lane, SequenceState& sequence ? std::min(capacity, prompt_tokens + (initial_mtp_extent == 0 ? 0U : initial_mtp_extent - 1U)) : speculative_backend == SpeculativeBackend::DFlash ? prompt_tokens - : speculative_backend == SpeculativeBackend::DFlash2 ? prompt_tokens + : speculative_backend == SpeculativeBackend::DFlash2 ? 0U : 0U; materialize_sequence_kv(sequence, prompt_tokens, backend_materialized); install_sampling(sequence, request, request_plan.sampling); @@ -9640,7 +9641,7 @@ runtime::ExecutionTiming ProgramImplCore::resolve_pending_raw( timing.resume_submit(); replay_fold->execute(std::span(fold_rows.data(), lanes.size()), device.stream); - + if (needs_hidden_correction) { const auto batch = static_cast(lanes.size()); Tensor selector_tensor; @@ -9653,7 +9654,9 @@ runtime::ExecutionTiming ProgramImplCore::resolve_pending_raw( hidden = frame.target_hidden.slice(2, 0, batch); selected = frame.target_continuation_hidden.slice(1, 0, batch); destinations = frame.state_destination_slots.slice(0, 0, batch); - } else if (speculative_backend == SpeculativeBackend::DFlash && io.dflash_decode) { + } else if ((speculative_backend == SpeculativeBackend::DFlash || + speculative_backend == SpeculativeBackend::DFlash2) && + io.dflash_decode) { qwen3_6::DFlashDecodeState& frame = *io.dflash_decode; selector_tensor = frame.proposal_extents.slice(0, 0, batch); hidden = frame.target_hidden.slice(2, 0, batch); @@ -9671,7 +9674,7 @@ runtime::ExecutionTiming ProgramImplCore::resolve_pending_raw( device.stream); } - if (speculative_backend == SpeculativeBackend::DFlash) { + if (speculative_backend == SpeculativeBackend::DFlash) { std::array append_lanes{}; std::array append_starts{}; std::array append_counts{}; @@ -9692,7 +9695,7 @@ runtime::ExecutionTiming ProgramImplCore::resolve_pending_raw( } } - timing.begin_wait(); + timing.begin_wait(); device.synchronize(); timing.end_wait(); work.reset(); @@ -9722,7 +9725,9 @@ runtime::ExecutionTiming ProgramImplCore::resolve_pending_raw( const TokenId* token_base = speculative_backend == SpeculativeBackend::Mtp ? mtp_host_egress->licensed_tokens.data() + row * width - : dflash_host_egress->licensed_tokens.data() + row * width; + : speculative_backend == SpeculativeBackend::DFlash2 + ? dflash2_host_egress->licensed_tokens.data() + row * width + : dflash_host_egress->licensed_tokens.data() + row * width; sequence.ledger.insert(sequence.ledger.end(), token_base, token_base + committed); sequence.prefix_identity.append_generated(committed, sequence.rope_delta); sequence.prefix_digests.append_generated( @@ -10781,7 +10786,8 @@ ProgramImplCore::advance_prefill(SequenceState& sequence, RequestControl& reques selectors.source, selectors.destination, staged.initial_mtp_extent, - dflash_host_ingress}; + dflash_host_ingress, + dflash2_host_ingress}; if (staged.mtp_bridge == MtpBridgeMode::BeforeSuffix) { if (staged.cursor != staged.base || staged.base == 0 || @@ -10876,7 +10882,8 @@ ProgramImplCore::advance_prefill(SequenceState& sequence, RequestControl& reques final_chunk_tokens = result.processed_tokens; sequence.text_kv_valid = staged.cursor; if (staged.prepare_mtp) { sequence.mtp_kv_valid = staged.cursor; } - if (speculative_backend == SpeculativeBackend::DFlash) { + if (speculative_backend == SpeculativeBackend::DFlash || + speculative_backend == SpeculativeBackend::DFlash2) { sequence.dflash_context_frontier = staged.cursor; } commit_sequence_kv(sequence, sequence.text_kv_valid, backend_kv_valid(sequence)); @@ -11511,7 +11518,7 @@ ProgramImplCore::decode_dflash_batch(std::span lanes, .timing = timing.finish(), }; } catch (...) { - timing.begin_wait(); + timing.begin_wait(); try { nvtx::ScopedRange wait_range(nvtx::Name::DecodeDFlashWait, nvtx::Category::Control, static_cast(lanes.size())); @@ -11638,7 +11645,8 @@ ProgramImplCore::decode_dflash2_batch(std::span lanes, dflash2_host_ingress->dflash_kv_table_rows[row] = 0; dflash2_host_ingress->active_lanes[row] = static_cast(sequence.lane); dflash2_host_ingress->sampling[row] = request.sampling_host; - materialize_sequence_kv(sequence, frontier + extent + 1U, frontier); + // DFlash2 owns no backend KV; only the shared text cache materializes. + materialize_sequence_kv(sequence, frontier + extent + 1U, 0); } schedule::DFlash2BatchContext schedule_state{ @@ -11670,11 +11678,11 @@ ProgramImplCore::decode_dflash2_batch(std::span lanes, RequestControl& request = requests[lanes[row]]; const std::uint32_t base_E = sequence.execution_frontier; const std::uint32_t base_S = sequence.ledger_frontier; - const std::int32_t count_i = dflash2_host_egress->licensed_counts[row]; + const std::int32_t count_i = dflash2_host_egress->licensed_counts[row]; const std::int32_t accepted_i = dflash2_host_egress->accepted_drafts[row]; - const std::uint32_t extent = + const std::uint32_t extent = static_cast(dflash2_host_egress->proposal_extents[row]); - if (count_i <= 0 || count_i > static_cast(width) || accepted_i < 0 || + if (count_i <= 0 || count_i > static_cast(width) || accepted_i < 0 || accepted_i + 1 != count_i || accepted_i > static_cast(extent) || extent > width || extent > static_cast(dflash2_host_ingress->proposal_extents[row]) || @@ -11709,7 +11717,7 @@ ProgramImplCore::decode_dflash2_batch(std::span lanes, request.lifecycle = Lifecycle::Pending; request.timings.decode_seconds += seconds; } - return runtime::BatchedGeneratedRound{ + return runtime::BatchedGeneratedRound{ .tokens = std::span(dflash2_host_egress->licensed_tokens.data(), lanes.size() * width), .row_counts = std::span(dflash2_host_egress->licensed_counts.data(), @@ -11717,7 +11725,7 @@ ProgramImplCore::decode_dflash2_batch(std::span lanes, .row_stride = width, }; } catch (...) { - timing.begin_wait(); + timing.begin_wait(); try { nvtx::ScopedRange wait_range(nvtx::Name::DecodeDFlashWait, nvtx::Category::Control, static_cast(lanes.size())); diff --git a/src/targets/qwen3_6/impl/runtime/schedule.h b/src/targets/qwen3_6/impl/runtime/schedule.h index d2f62c31b8..b00a3bba5e 100644 --- a/src/targets/qwen3_6/impl/runtime/schedule.h +++ b/src/targets/qwen3_6/impl/runtime/schedule.h @@ -55,7 +55,8 @@ struct PrefillContext { std::int32_t state_source_slot = 0; std::int32_t state_destination_slot = 0; std::uint32_t mtp_proposal_extent = 0; - const qwen3_6::DFlashDecodeIngress* dflash_host_ingress = nullptr; + const qwen3_6::DFlashDecodeIngress* dflash_host_ingress = nullptr; + const qwen3_6::DFlashDecodeIngress* dflash2_host_ingress = nullptr; }; struct OrdinaryBatchContext { diff --git a/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h b/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h index 115c20905b..4df3fca3e2 100644 --- a/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_prefill_impl.h @@ -14,10 +14,10 @@ namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule { namespace { DFlashFeatureSink make_dflash_prefill_sink(PrefillContext& state) { - if (!state.execution.io.dflash_decode || state.dflash_host_ingress == nullptr) { - throw std::logic_error("DFlash prefill controls are unavailable"); - } if (state.execution.model.features.dflash2()) { + if (!state.execution.io.dflash_decode || state.dflash2_host_ingress == nullptr) { + throw std::logic_error("DFlash2 prefill controls are unavailable"); + } return dflash2_feature_sink( state, [&state](const Tensor& features, const Tensor& positions, bool rewrite_checkpoint) { @@ -29,11 +29,15 @@ DFlashFeatureSink make_dflash_prefill_sink(PrefillContext& state) { const auto exact = static_cast(features.ne[1]); dflash2_append_context(state, features, positions, count, lane, row, {exact, exact}); if (rewrite_checkpoint) { - state.dflash2->save_rewrite_checkpoint(state.dflash_host_ingress->active_lanes[0], - state.execution.device.stream); + state.dflash2->save_rewrite_checkpoint( + state.dflash2_host_ingress->active_lanes[0], + state.execution.device.stream); } }); } + if (!state.execution.io.dflash_decode || state.dflash_host_ingress == nullptr) { + throw std::logic_error("DFlash prefill controls are unavailable"); + } return dflash_feature_sink( state, [&state](const Tensor& features, const Tensor& positions, bool rewrite_checkpoint) { auto& frame = *state.execution.io.dflash_decode; @@ -80,7 +84,7 @@ PrefillChunkResult prefill_text_chunk(PrefillContext& state, std::span(*split_frontier) : -1); const std::span prompt(ids.data(), ids.size()); - if (state.dflash != nullptr) { + if (state.dflash != nullptr || state.dflash2 != nullptr) { DFlashFeatureSink sink = make_dflash_prefill_sink(state); return card.prefill_chunk(prompt, state.text_kv_base, nominal_length, finalize_at_end, sink); @@ -93,7 +97,7 @@ PrefillChunkResult prefill_multimodal_chunk(PrefillContext& state, const Prepare std::uint32_t nominal_length, std::optional split_frontier, bool finalize_at_end) { - if (state.dflash != nullptr) { + if (state.dflash != nullptr || state.dflash2 != nullptr) { throw std::logic_error("DFlash staged multimodal prefill is unavailable"); } TextContext card(state.execution.device, state.execution.model, state.execution.work, diff --git a/src/targets/qwen3_6_27b/impl/variant.cpp b/src/targets/qwen3_6_27b/impl/variant.cpp index 8634e090b9..712fb098f1 100644 --- a/src/targets/qwen3_6_27b/impl/variant.cpp +++ b/src/targets/qwen3_6_27b/impl/variant.cpp @@ -538,6 +538,7 @@ std::size_t Variant::post_mixer_workspace_capacity_bytes(WeightsProfile weights_ return post_mixer_workspace_bytes(QType::NVFP4, QType::NVFP4, kNvfp4TextPolicy, first, last); case WeightsProfile::Qwen38Nvfp4: { + case WeightsProfile::Qwen38Nvfp4DFlash2: const std::size_t nvfp4 = post_mixer_workspace_bytes(QType::NVFP4, QType::NVFP4, kNvfp4TextPolicy, first, last); const std::size_t fp8 = post_mixer_workspace_bytes( From e1f12d90fafce14679ec40a13fbb96beda33c9e3 Mon Sep 17 00:00:00 2001 From: NInfer Agent Date: Mon, 31 Aug 2026 19:18:51 +0800 Subject: [PATCH 4/4] fix(spec): DFlash2 graph-prep segfault + auto resolution + MTP fallback Applied to the PR3 branch from the integration fixes (9bf5135, d4f6e0d): 1. CUDA-graph preparation segfaulted with --spec dflash2: the graph representative wrote dflash_host_ingress unconditionally, but DFlash2 keeps separate dflash2_host_* buffers and dflash_host_ingress is null when only DFlash2 is active. Select ingress/egress by backend and zero DFlash2's pending_features too. 2. --spec auto failed to load because only plan_load resolved it; the planner and loaded model saw SpeculativeBackend::Auto and the startup features mismatch check rejected the weights. registry now resolves auto once up front (resolved_auto_speculative is a public Package static; 35b gets the same hook, auto -> MTP). 3. auto picks the backend from the artifact weights: DFlash2 artifact always DFlash2 (no MTP head), any other artifact defaults to MTP. Verified: dflash2 CLI/serve generation, dflash2 + auto -> speculative= dflash2, base + auto -> MTP. --- .../qwen3_6/impl/runtime/program_impl.h | 44 +++++++++++++------ .../ninfer/targets/qwen3_6_27b/package.h | 4 ++ src/targets/qwen3_6_27b/impl/package.cpp | 29 +++++------- .../ninfer/targets/qwen3_6_35b_a3b/package.h | 3 ++ src/targets/qwen3_6_35b_a3b/impl/package.cpp | 11 +++++ src/targets/registry.cpp | 18 +++++--- 6 files changed, 71 insertions(+), 38 deletions(-) diff --git a/src/targets/qwen3_6/impl/runtime/program_impl.h b/src/targets/qwen3_6/impl/runtime/program_impl.h index a4ef0abeeb..27573ebe86 100644 --- a/src/targets/qwen3_6/impl/runtime/program_impl.h +++ b/src/targets/qwen3_6/impl/runtime/program_impl.h @@ -10300,6 +10300,11 @@ void ProgramImplCore::prepare_graphs() { dflash->pending_features.slice(2, static_cast(row), 1); CUDA_CHECK(cudaMemsetAsync(pending.data, 0, pending.bytes(), device.stream)); } + if (dflash2) { + const Tensor pending = + dflash2->pending_features.slice(2, static_cast(row), 1); + CUDA_CHECK(cudaMemsetAsync(pending.data, 0, pending.bytes(), device.stream)); + } } set_device_i32(io.pos, checked_i32(frontier, "graph representative position")); set_device_i32(io.rope_pos, checked_i32(frontier, "graph representative rope position")); @@ -10308,24 +10313,33 @@ void ProgramImplCore::prepare_graphs() { checked_i32(frontier, "graph representative MTP position")); } if (io.dflash_decode) { - *dflash_host_ingress = {}; - *dflash_host_egress = {}; + // DFlash and DFlash2 share the decode frame but keep separate + // host ingress/egress buffers; only the active backend's is + // allocated (dflash_host_* vs dflash2_host_*). + qwen3_6::DFlashDecodeIngress* ingress = + speculative_backend == SpeculativeBackend::DFlash2 ? dflash2_host_ingress + : dflash_host_ingress; + qwen3_6::DFlashDecodeEgress* egress = + speculative_backend == SpeculativeBackend::DFlash2 ? dflash2_host_egress + : dflash_host_egress; + *ingress = {}; + *egress = {}; const std::uint32_t extent = std::min(draft_window, capacity - frontier - 1U); for (std::uint32_t row = 0; row < batch_size; ++row) { - dflash_host_ingress->anchors[row] = 0; - dflash_host_ingress->execution_frontiers[row] = + ingress->anchors[row] = 0; + ingress->execution_frontiers[row] = checked_i32(frontier, "graph representative DFlash frontier"); - dflash_host_ingress->context_frontiers[row] = + ingress->context_frontiers[row] = checked_i32(frontier, "graph representative DFlash context frontier"); - dflash_host_ingress->proposal_extents[row] = static_cast(extent); - dflash_host_ingress->target_valid_columns[row] = + ingress->proposal_extents[row] = static_cast(extent); + ingress->target_valid_columns[row] = static_cast(extent + 1U); - dflash_host_ingress->text_kv_table_rows[row] = static_cast(row); - dflash_host_ingress->dflash_kv_table_rows[row] = static_cast(row); - dflash_host_ingress->active_lanes[row] = static_cast(row); - dflash_host_ingress->state_source_slots[row] = capture_state_slot(row); - dflash_host_ingress->state_destination_slots[row] = capture_state_slot(row); - dflash_host_ingress->sampling[row] = {}; + ingress->text_kv_table_rows[row] = static_cast(row); + ingress->dflash_kv_table_rows[row] = static_cast(row); + ingress->active_lanes[row] = static_cast(row); + ingress->state_source_slots[row] = capture_state_slot(row); + ingress->state_destination_slots[row] = capture_state_slot(row); + ingress->sampling[row] = {}; } } if (io.mtp_decode) { @@ -10516,11 +10530,15 @@ void ProgramImplCore::prepare_graphs() { const ops::CausalAttentionExecutionEnvelope code_warm_target{ 1, static_cast(std::min( capacity, static_cast(code_warm.max) + draft_window + 1ULL))}; + std::fprintf(stderr, "[df2diag] prepare_representative min=%u\n", code_warm.min); prepare_representative(code_warm.min, 1); + std::fprintf(stderr, "[df2diag] after prepare_representative\n"); device.synchronize(); + std::fprintf(stderr, "[df2diag] calling dflash2_decode_batch\n"); schedule::dflash2_decode_batch(dflash2_state, 1, draft_window, dflash2_envelopes(code_warm.min, code_warm.max, draft_window), code_warm_target, nullptr); + std::fprintf(stderr, "[df2diag] after dflash2_decode_batch\n"); device.synchronize(); dflash2_graphs.profiles.reserve(batch_one_profiles.size() * max_concurrency); diff --git a/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h b/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h index cf2b7bdf1b..314699696f 100644 --- a/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h +++ b/src/targets/qwen3_6_27b/export/ninfer/targets/qwen3_6_27b/package.h @@ -124,6 +124,10 @@ struct Package { [[nodiscard]] static ModelSamplingDefaults sampling_defaults(std::string_view model); [[nodiscard]] static WeightsProfile resolve_weights(const artifact::ArtifactIdentity& identity); + // Resolves --spec auto to a concrete backend (MTP or DFlash2) using the + // artifact weights profile and context budget. + [[nodiscard]] static EngineOptions resolved_auto_speculative(const EngineOptions& options, + WeightsProfile weights_profile); [[nodiscard]] static LoadPlan plan_load(artifact::Binder& binder, const EngineOptions& options, WeightsProfile weights_profile); [[nodiscard]] static std::unique_ptr diff --git a/src/targets/qwen3_6_27b/impl/package.cpp b/src/targets/qwen3_6_27b/impl/package.cpp index 9076657ce6..75e6a1704a 100644 --- a/src/targets/qwen3_6_27b/impl/package.cpp +++ b/src/targets/qwen3_6_27b/impl/package.cpp @@ -102,25 +102,17 @@ Package::WeightsProfile Package::resolve_weights(const artifact::ArtifactIdentit "' is not supported by target '" + std::string(target_key) + "'"); } -namespace { -EngineOptions resolved_auto_speculative(const EngineOptions& options, - detail::WeightsProfile weights_profile) { +EngineOptions Package::resolved_auto_speculative(const EngineOptions& options, + WeightsProfile weights_profile) { EngineOptions resolved = options; if (options.speculative.backend != SpeculativeBackend::Auto) { return resolved; } - const DType kv_dtype = - options.kv_cache == KvCacheStorage::BFloat16 - ? DType::BF16 - : (options.kv_cache == KvCacheStorage::Int8Group64 - ? DType::I8 - : (options.kv_cache == KvCacheStorage::Fp8E4M3Row256 - ? DType::FP8_E4M3FN - : DType::BF16)); - const std::uint32_t draft_capacity = - kv_dtype == DType::FP8_E4M3FN ? 8192U : (kv_dtype == DType::I8 ? 4096U : 2048U); - // The artifact's frozen startup features enforce this limit; the engine - // check makes the fallback graceful (selects MTP) instead of a load error. - if (weights_profile == detail::WeightsProfile::Qwen38Nvfp4DFlash2 && !options.enable_vision && - options.max_context <= draft_capacity) { + // The artifact weights decide the backend, not the context length: a + // DFlash2 artifact has no MTP draft head (the two are mutually + // exclusive), so auto must pick DFlash2 even at long contexts - a memory + // shortfall then surfaces as a clear reservation error instead of a + // confusing weights mismatch. A non-DFlash2 artifact defaults to MTP. + if (weights_profile == detail::WeightsProfile::Qwen38Nvfp4DFlash2 && + !options.enable_vision) { resolved.speculative.backend = SpeculativeBackend::DFlash2; if (resolved.speculative.draft_tokens == 0) { resolved.speculative.draft_tokens = 7; } } else { @@ -129,11 +121,10 @@ EngineOptions resolved_auto_speculative(const EngineOptions& options, } return resolved; } -} // namespace Package::LoadPlan Package::plan_load(artifact::Binder& binder, const EngineOptions& options, WeightsProfile weights_profile) { - const EngineOptions resolved = resolved_auto_speculative(options, weights_profile); + const EngineOptions resolved = Package::resolved_auto_speculative(options, weights_profile); return LoadPlan(std::make_unique( weights_profile, detail::bind_artifact(binder, weights_profile, qwen3_6::startup_features(resolved)))); diff --git a/src/targets/qwen3_6_35b_a3b/export/ninfer/targets/qwen3_6_35b_a3b/package.h b/src/targets/qwen3_6_35b_a3b/export/ninfer/targets/qwen3_6_35b_a3b/package.h index a049706dae..f079d442d9 100644 --- a/src/targets/qwen3_6_35b_a3b/export/ninfer/targets/qwen3_6_35b_a3b/package.h +++ b/src/targets/qwen3_6_35b_a3b/export/ninfer/targets/qwen3_6_35b_a3b/package.h @@ -118,6 +118,9 @@ struct Package { [[nodiscard]] static ModelSamplingDefaults sampling_defaults(std::string_view model); [[nodiscard]] static WeightsProfile resolve_weights(const artifact::ArtifactIdentity& identity); + // Resolves --spec auto (this target has no DFlash2 weights; auto -> MTP). + [[nodiscard]] static EngineOptions resolved_auto_speculative(const EngineOptions& options, + WeightsProfile weights_profile); [[nodiscard]] static LoadPlan plan_load(artifact::Binder& binder, const EngineOptions& options, WeightsProfile weights_profile); [[nodiscard]] static std::unique_ptr diff --git a/src/targets/qwen3_6_35b_a3b/impl/package.cpp b/src/targets/qwen3_6_35b_a3b/impl/package.cpp index 1bed0142e0..2d274ab63a 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/package.cpp +++ b/src/targets/qwen3_6_35b_a3b/impl/package.cpp @@ -72,6 +72,17 @@ Package::WeightsProfile Package::resolve_weights(const artifact::ArtifactIdentit "' is not supported by target '" + std::string(target_key) + "'"); } +EngineOptions Package::resolved_auto_speculative(const EngineOptions& options, + WeightsProfile) { + EngineOptions resolved = options; + if (options.speculative.backend == SpeculativeBackend::Auto) { + // This target carries no DFlash2 weights; auto falls back to MTP. + resolved.speculative.backend = SpeculativeBackend::Mtp; + if (resolved.speculative.draft_tokens == 0) { resolved.speculative.draft_tokens = 3; } + } + return resolved; +} + Package::LoadPlan Package::plan_load(artifact::Binder& binder, const EngineOptions& options, WeightsProfile weights_profile) { return LoadPlan(std::make_unique( diff --git a/src/targets/registry.cpp b/src/targets/registry.cpp index 17548ebe61..dd1479595f 100644 --- a/src/targets/registry.cpp +++ b/src/targets/registry.cpp @@ -93,6 +93,12 @@ ConstructedTarget construct_registered(const EngineOptions& options, DeviceConte std::string_view target_key) { const auto& identity = reader.identity(); const auto weights_profile = Target::resolve_weights(identity); + // Resolve --spec auto up front so the planner, the load plan and the + // program all see the same concrete backend (previously only plan_load + // resolved it, so the planner built an Auto plan and the frozen startup + // features mismatch check rejected the loaded weights). + const EngineOptions resolved_options = + Target::resolved_auto_speculative(options, weights_profile); const ModelSamplingDefaults sampling_defaults = Target::sampling_defaults(identity.model_id); const runtime::ContextCostIdentity context_cost_identity{ .hardware_class = runtime::context_cost_hardware_class( @@ -104,14 +110,14 @@ ConstructedTarget construct_registered(const EngineOptions& options, DeviceConte context_cost_identity, options.context_cost.preset_path); artifact::Binder binder(reader); - auto load_plan = Target::plan_load(binder, options, weights_profile); - auto sequence_planner = Target::make_sequence_planner(device, options, weights_profile); + auto load_plan = Target::plan_load(binder, resolved_options, weights_profile); + auto sequence_planner = Target::make_sequence_planner(device, resolved_options, weights_profile); const runtime::SequenceCapacityCurve curve = sequence_planner.capacity_curve(); const std::size_t preflight_runtime_bytes = runtime_bytes_after_planned_weights(load_plan.materialization().device_capacity_bytes); - (void)runtime::resolve_kv_capacity(options.kv_capacity, curve, preflight_runtime_bytes); + (void)runtime::resolve_kv_capacity(resolved_options.kv_capacity, curve, preflight_runtime_bytes); - auto progress = artifact_progress(options.load_progress); + auto progress = artifact_progress(resolved_options.load_progress); auto materialized = artifact::materialize(reader, load_plan.materialization(), device, progress.callback ? &progress : nullptr); const artifact::MaterializationStats stats = materialized.stats(); @@ -119,13 +125,13 @@ ConstructedTarget construct_registered(const EngineOptions& options, DeviceConte auto model = Target::construct_loaded_model(std::move(load_plan), std::move(materialized)); device.synchronize(); runtime::KvCapacityResolution capacity_resolution = - runtime::resolve_kv_capacity(options.kv_capacity, curve, current_free_device_bytes()); + runtime::resolve_kv_capacity(resolved_options.kv_capacity, curve, current_free_device_bytes()); auto sequence_plan = std::move(sequence_planner).finalize(capacity_resolution.main_page_groups); if (sequence_plan.device_reservation_bytes() != capacity_resolution.runtime_reservation_bytes || sequence_plan.kv_capacity() != capacity_resolution.resolved_tokens) { throw std::logic_error("resolved KV capacity does not match the finalized target plan"); } - auto loaded = std::make_unique(std::move(model), options); + auto loaded = std::make_unique(std::move(model), resolved_options); auto instance = std::make_unique(std::move(loaded), capacity_resolution, std::move(sequence_plan), device); device.synchronize();