PR3: DFlash2 draft backend + length-based speculative switching

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 <cuda_runtime.h>
+
+#include <cstddef>
+#include <cstdint>
+
+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

include/ninfer/ops/dflash2_grouped_conv.h

@@ -0,0 +1,38 @@
+#pragma once
+
+#include "core/tensor.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+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

include/ninfer/ops/dflash2_selector.h

@@ -0,0 +1,47 @@
+#pragma once
+
+#include "core/tensor.h"
+
+#include <cuda_runtime.h>
+
+#include <cstdint>
+
+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

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

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 <cuda_runtime.h>
+
+#include <cstddef>
+#include <cstdint>
+
+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; lanes[b] selects its cyclic-cache lane. Columns 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)<window. Thus
+ * distance window-1 is included, distance window is excluded, and every query row sees every
+ * live temporary query row from the same batch row. scale is 1/sqrt(128).
+ *
+ * Context and query K/V are unchanged. out is the only observable mutation and is completely
+ * overwritten. The current optimized implementation domain is T=1..16 on sm_120a.
+ *
+ * The caller guarantees min_context <= L <= max_context, sequential nonnegative positions, and
+ * that the cyclic context contains the declared live interval. The execution envelope may affect
+ * finite launch selection and workspace capacity, never the admitted key set.
+ */
+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);
+
+/**
+ * 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 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);
+
+} // namespace ninfer::ops

include/ninfer/types.h

@@ -69,6 +69,8 @@ enum class SpeculativeBackend : std::uint8_t {
     None,
     Mtp,
     DFlash,
+    DFlash2,
+    Auto,
 };
 
 struct SpeculativeOptions {

src/CMakeLists.txt

@@ -65,8 +65,11 @@ target_link_libraries(ninfer_nvfp4_tma PRIVATE ninfer_core CUDA::cudart CUDA::cu
 add_library(ninfer_ops STATIC
   ops/launcher/add_bias.cu
   ops/launcher/argmax.cu
+  ops/launcher/bidirectional_gqa_attention.cu
   ops/launcher/cast.cu
   ops/launcher/causal_conv1d.cu
+  ops/launcher/dflash2_grouped_conv.cu
+  ops/launcher/dflash2_selector.cu
   ops/launcher/embed_gather.cu
   ops/launcher/gdn_gating.cu
   ops/launcher/gelu.cu
@@ -83,6 +86,7 @@ add_library(ninfer_ops STATIC
   ops/launcher/rope.cu
   ops/launcher/sampling.cu
   ops/launcher/scalar.cu
+  ops/launcher/swa.cu
   ops/launcher/scatter.cu
   ops/launcher/sigmoid_gate_mul.cu
   ops/launcher/silu_and_mul.cu
@@ -238,8 +242,11 @@ add_library(ninfer_ops STATIC
   ops/linear/linear.cpp
   ops/wrapper/add_bias.cpp
   ops/wrapper/argmax.cpp
+  ops/wrapper/bidirectional_gqa_attention.cpp
   ops/wrapper/cast.cpp
   ops/wrapper/causal_conv1d_silu.cpp
+  ops/wrapper/dflash2_grouped_conv.cpp
+  ops/wrapper/dflash2_selector.cpp
   ops/wrapper/embedding.cpp
   ops/wrapper/gdn_gating.cpp
   ops/wrapper/gdn_gating_proj.cpp
@@ -261,6 +268,7 @@ add_library(ninfer_ops STATIC
   ops/wrapper/rmsnorm.cpp
   ops/wrapper/rope.cpp
   ops/wrapper/sampling.cpp
+  ops/wrapper/swa.cpp
   ops/wrapper/scalar.cpp
   ops/wrapper/scatter.cpp
   ops/wrapper/sigmoid_mul.cpp