feat(dflash2): execute the NVFP4-encoded DFlash2 module on the nvfp4full artifact

README.md

@@ -17,13 +17,14 @@ runtime:
 | [Qwen3.6-27B NVFP4](https://huggingface.co/neroued/Qwen3.6-27B-nvfp4-NInfer) | `nvfp4` | `qwen3_6_27b_nvfp4.ninfer` | 18,324,064,000 bytes (17.07 GiB) | `bce5f00d066c0f20f1317bf1fdcb458264cf95837c3b1f3fbec163694627893a` |
 | [Qwen3.8-27B](https://huggingface.co/neroued/Qwen3.8-27B-NInfer) | `groupwise-int` | `qwen3_8_27b.ninfer` | 20,437,336,576 bytes (19.03 GiB) | `0634abb07024221de141456cf04a42ab74b18bc38e1b781c6eb2e062a467eec3` |
 | [Qwen3.8-27B NVFP4](https://huggingface.co/neroued/Qwen3.8-27B-nvfp4-NInfer) | `nvfp4` | `qwen3_8_27b_nvfp4.ninfer` | 23,719,496,192 bytes (22.09 GiB) | `552c374c685dce302603b95fbe940fb04243c0cd44c083efc644ad3d980d462c` |
-| [Qwen3.8-27B NVFP4F](https://huggingface.co/cometkim/Qwen3.8-27B-nvfp4full-NInfer) | `nvfp4full` | `qwen3_8_27b_nvfp4full.ninfer` | 18,324,059,648 bytes (17.07 GiB) | `2f59cc27d67cb7acba0ba8a0e0881ac89c1db2b267a60119a696fefa12faf4e7` |
+| [Qwen3.8-27B NVFP4F](https://huggingface.co/cometkim/Qwen3.8-27B-nvfp4full-NInfer) | `nvfp4full` | `qwen3_8_27b_nvfp4full.ninfer` | 19,406,942,468 bytes (18.07 GiB) | `abb1e120d5f1f32d61689604d238227ff579ab76cbd9319628f3b3904fffd9af` |
 | [Qwen3.6-35B-A3B](https://huggingface.co/neroued/Qwen3.6-35B-A3B-NInfer) | `groupwise-int` | `qwen3_6_35b_a3b.ninfer` | 22,783,246,080 bytes (21.22 GiB) | `1fb9ea0b5b8561e49d9604115ec89e5d9f2b6f6434e32c37c57fffd480a325d2` |
 
-The current Qwen3.8 `groupwise-int` and `nvfp4` artifacts include DFlash2 companion weights;
-select `--spec dflash2 --draft-tokens 7 --lm-head-draft` in a current source build (portable
-v0.6.1 predates this backend). The `nvfp4full` (Qwen3.8-27B NVFP4F) artifact does not include
-DFlash2 companion weights, so `--spec dflash2` is currently unsupported on it. Older Qwen3.8
+The current Qwen3.8 `groupwise-int`, `nvfp4`, and `nvfp4full` artifacts include DFlash2 companion
+weights; select `--spec dflash2 --draft-tokens 7 --lm-head-draft` in a current source build (portable
+v0.6.1 predates this backend). The `nvfp4full` artifact stores its DFlash2 module in a weight-only
+NVFP4 encoding (matrices only; the upstream `W8G32_F16S` schema is not used), which this build's
+unified module binder executes directly. Older Qwen3.8
 artifacts remain usable for Text, Vision and MTP in the current build, but cannot enable
 DFlash2. See [DFlash2 on Windows](docs/windows.md#dflash2) for launch and validation commands.
 

bench/ops/candidate_selector_bench.cu

@@ -140,14 +140,31 @@ struct Fixture {
     }
 
     void tensors(std::int32_t batch_size, Tensor& ids, Tensor& unary, Tensor& hidden,
-                 Tensor& anchor, Tensor& predecessor, Tensor& successor, Tensor& positions,
+                 Tensor& anchor, Weight& predecessor, Weight& successor, Tensor& positions,
                  Tensor& draft, Tensor& q) {
         ids         = Tensor(candidate_ids.p, DType::I32, {kCandidates, kSteps, batch_size});
         unary       = Tensor(unary_scores.p, DType::FP32, {kCandidates, kSteps, batch_size});
         hidden      = Tensor(projected_hidden.p, DType::BF16, {kRank, kSteps, batch_size});
         anchor      = Tensor(anchors.p, DType::I32, {batch_size});
-        predecessor = Tensor(predecessor_codebook.p, DType::BF16, {kRank, kCodebookRows});
-        successor   = Tensor(successor_codebook.p, DType::BF16, {kRank, kCodebookRows});
+        const auto codebook_weight = [&](void* data) {
+            Weight weight{};
+            weight.payload      = data;
+            weight.payload_bytes =
+                static_cast<std::uint64_t>(kCodebookRows) * kRank * sizeof(std::uint16_t);
+            weight.qtype            = QType::BF16_CTRL;
+            weight.ndim             = 2;
+            weight.qdata            = data;
+            weight.n                = kCodebookRows;
+            weight.k                = kRank;
+            weight.shape[0]         = kCodebookRows;
+            weight.shape[1]         = kRank;
+            weight.padded_shape[0]  = kCodebookRows;
+            weight.padded_shape[1]  = kRank;
+            weight.layout           = QuantLayout::Contiguous;
+            return weight;
+        };
+        predecessor = codebook_weight(predecessor_codebook.p);
+        successor   = codebook_weight(successor_codebook.p);
         positions   = Tensor(base_positions.p, DType::I32, {batch_size});
         draft       = Tensor(drafts.p, DType::I32, {kSteps, batch_size});
         q           = Tensor(proposal_q.p, DType::FP32, {kCandidates, kSteps, batch_size});
@@ -161,8 +178,8 @@ void run(std::int32_t batch_size, Mode mode, const Options& options, Fixture& fi
     Tensor unary;
     Tensor hidden;
     Tensor anchor;
-    Tensor predecessor;
-    Tensor successor;
+    Weight predecessor;
+    Weight successor;
     Tensor positions;
     Tensor draft;
     Tensor q;

include/ninfer/ops/attn_input_proj.h

@@ -84,13 +84,15 @@ void attn_input_proj(const Tensor& x, const Weight& query_key_gate_value_weight,
                      Tensor& gate, Tensor& k, Tensor& v, cudaStream_t stream);
 
 /**
- * Three-output W8 specialization. The W8G32_F16S RowSplit parent stores rows in order
- * [query 4096, key 1024, value 1024]. Registered parent forms are [6144,2048] with BF16
- * x [2048,T] for the Qwen3.6 companion and [6144,5120] with BF16 x [5120,T] for DFlash2.
- * q is contiguous BF16 [4096,T], and k/v are contiguous BF16 [1024,T]. Every route writes the
- * three independent final allocations directly; no parent output or transient workspace is
- * materialized. T may be any positive value. Q and K remain raw projection outputs: this Op does
- * not normalize or rotate either tensor.
+ * Three-output specialization. The parent stores rows in physical order [query, key, value].
+ * Registered parent forms are the W8G32_F16S RowSplit matrices [6144,2048] with BF16 x [2048,T]
+ * for the Qwen3.6 companion, [6144,5120] with BF16 x [5120,T] for DFlash2, and the weight-only
+ * NVFP4 BlockScaleK16M128x4 matrix [6144,5120] with BF16 x [5120,T] for the fork-format DFlash2
+ * module. q is contiguous BF16 [4096,T], and k/v are contiguous BF16 [1024,T]. Every route
+ * writes the three independent final allocations directly; no parent output or transient
+ * workspace is materialized. T may be any positive value (the NVFP4 route serves extents above
+ * its fused small-T family in 32-token chunks). Q and K remain raw projection outputs: this Op
+ * does not normalize or rotate either tensor.
  */
 void attn_input_proj(const Tensor& x, const Weight& query_key_value_weight, Tensor& q, Tensor& k,
                      Tensor& v, cudaStream_t stream);

include/ninfer/ops/candidate_selector.h

@@ -19,10 +19,12 @@ namespace ninfer::ops {
  *
  * For K in [1,15] and B in [1,8], the inputs are contiguous candidate_ids I32 [16,K,B],
  * unary_scores FP32 [16,K,B], projected_hidden BF16 [256,K,B], anchors I32 [B],
- * predecessor_codebook and successor_codebook BF16 [256,248320], base_positions I32 [B], and a
- * device-resident SamplingConfig[B]. Candidate rank is the fastest axis. The 16 candidate ids in
- * each row are distinct, and all candidate and anchor token ids lie in [0,248077); the registered
- * vocabulary, artifact binding, and linear_topk producer establish that trusted value contract.
+ * predecessor_codebook and successor_codebook weights of logical shape [256,248320] in either
+ * BF16_CTRL Contiguous or weight-only NVFP4 BlockScaleK16M128x4 form, base_positions I32 [B],
+ * and a device-resident SamplingConfig[B]. Candidate rank is the fastest axis. The 16 candidate
+ * ids in each row are distinct, and all candidate and anchor token ids lie in [0,248077); the
+ * registered vocabulary, artifact binding, and linear_topk producer establish that trusted value
+ * contract.
  *
  * Starting with predecessor=anchors[b], each position i in [0,K) computes:
  *
@@ -31,6 +33,10 @@ namespace ninfer::ops {
  *                   * projected_hidden[r,i,b]
  *                   * successor_codebook[r,candidate_ids[c,i,b]].
  *
+ * The oracle evaluates that formula in FP64 from the represented (decoded) codebook values; the
+ * NVFP4 production route decodes each gathered row element with its exact stored scale in FP32
+ * before the products.
+ *
  * A row with configs[b].temperature<=0 selects the lowest candidate rank attaining max(edge) and
  * writes its exact one-hot distribution. A positive-temperature row writes the FP32 softmax of
  * edge/temperature, then draws a candidate with counter key
@@ -45,7 +51,7 @@ namespace ninfer::ops {
  */
 void candidate_selector_path(const Tensor& candidate_ids, const Tensor& unary_scores,
                              const Tensor& projected_hidden, const Tensor& anchors,
-                             const Tensor& predecessor_codebook, const Tensor& successor_codebook,
+                             const Weight& predecessor_codebook, const Weight& successor_codebook,
                              const Tensor& base_positions, const SamplingConfig* configs,
                              Tensor& drafts, Tensor& proposal_q, WorkspaceArena& workspace,
                              cudaStream_t stream);

include/ninfer/ops/linear_swiglu.h

@@ -25,10 +25,11 @@ namespace ninfer::ops {
                                                                  std::int32_t max_tokens);
 
 /**
- * Policy-bearing capacity query. Q4/W8 admit A16Only. NVFP4 admits A16Only through T=16 and
- * AllowA4 for every positive T. Row-scaled FP8 admits A16Only and AllowA8 for every positive T.
- * A permissive policy covers whichever qualified route the private resolver selects across the
- * requested interval.
+ * Policy-bearing capacity query. Q4/W8 admit A16Only. NVFP4 admits A16Only at every positive T —
+ * fused through T=16, then a linear-then-silu_mul decomposition that materializes the gate/up
+ * projection — and AllowA4 for every positive T. Row-scaled FP8 admits A16Only and AllowA8 for
+ * every positive T. A permissive policy covers whichever qualified route the private resolver
+ * selects across the requested interval.
  */
 [[nodiscard]] std::size_t
 linear_swiglu_workspace_capacity_bytes(QType qtype, std::int32_t gate_up_rows,
@@ -75,8 +76,8 @@ void linear_swiglu(const Tensor& x, const Weight& gate_up_weight, Tensor& out, L
 
 /**
  * A16-only convenience form. Q4/W8 and row-scaled FP8 retain their complete positive-T domain.
- * NVFP4 is admitted only through T=16; larger NVFP4 extents require the policy-bearing AllowA4
- * form.
+ * NVFP4 keeps the complete positive-T domain: the fused small-T family through T=16, then the
+ * workspace-bearing linear-then-silu_mul decomposition.
  */
 void linear_swiglu(const Tensor& x, const Weight& gate_up_weight, Tensor& out, WorkspaceArena& ws,
                    cudaStream_t stream);

src/CMakeLists.txt

@@ -99,8 +99,11 @@ add_library(ninfer_ops STATIC
   ops/rmsnorm_rope/rmsnorm_rope.cpp
   ops/rmsnorm_rope/launch.cu
   ops/context_kv_materialize/context_kv_materialize.cpp
+  ops/context_kv_materialize/context_kv_key_post.cu
   ops/context_kv_materialize/materialize.cu
+  ops/context_kv_materialize/materialize_nvfp4.cu
   ops/candidate_selector/bf16/candidate_selector_path.cu
+  ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.cu
   ops/candidate_selector/bf16/candidate_selector_path_plan.cpp
   # Softmax Attention and KV-cache state transitions.
   ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp
@@ -142,6 +145,7 @@ add_library(ninfer_ops STATIC
   ops/attn_input_proj/fp8/fp8_attn_input_plan.cpp
   ops/attn_input_proj/nvfp4/nvfp4_attn_input_decode.cu
   ops/attn_input_proj/nvfp4/nvfp4_attn_input_small_t.cu
+  ops/attn_input_proj/nvfp4/nvfp4_dflash2_attn_input.cu
   ops/attn_input_proj/nvfp4/nvfp4_attn_input_w4a4.cu
   ops/attn_input_proj/nvfp4/nvfp4_attn_input_plan.cpp
   ops/attn_input_proj/q4_q5/q4_q5_attn_input_gemm_mma.cu
@@ -179,7 +183,9 @@ add_library(ninfer_ops STATIC
   ops/gdn_gating_proj/bf16/bf16_gdn_norm_gating_proj_27.cu
   ops/gdn_gating_proj/bf16/bf16_gdn_gating_proj_kernels.cu
   ops/gdn_gating_proj/bf16/bf16_gdn_gating_proj_plan.cpp
+  ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.cu
   ops/dynamic_grouped_conv/bf16/bf16_dynamic_grouped_conv_prepare_partial.cu
+  ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare.cu
   ops/dynamic_grouped_conv/bf16/bf16_dynamic_grouped_conv_prepare_reduce.cu
   ops/dynamic_grouped_conv/bf16/bf16_dynamic_grouped_conv_prepare_plan.cpp
   ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_materialized.cu
@@ -199,6 +205,7 @@ add_library(ninfer_ops STATIC
   ops/linear/nvfp4/nvfp4_format.cpp
   ops/linear/nvfp4/nvfp4_gemv.cu
   ops/linear/nvfp4/nvfp4_small_t.cu
+  ops/linear/nvfp4/nvfp4_small_t_dflash2.cu
   ops/linear/nvfp4/nvfp4_w4a4.cu
   ops/linear/nvfp4/nvfp4_dispatch.cpp
   ops/linear/q4/q4_small_t_mma.cu

src/artifact/binder.cpp

@@ -72,6 +72,18 @@ bool Binder::contains(std::string_view name) const noexcept {
     return reader_.find(name) != nullptr;
 }
 
+NumericFormat Binder::declared_format(std::string_view name) const {
+    const auto* object = reader_.find(name);
+    if (object == nullptr) {
+        throw ArtifactError("required artifact object is missing: " + std::string(name));
+    }
+    const auto* tensor = std::get_if<TensorDescriptor>(object);
+    if (tensor == nullptr) {
+        throw ArtifactError("required tensor is a resource: " + std::string(name));
+    }
+    return tensor->format;
+}
+
 const ObjectDescriptor& Binder::descriptor(ObjectHandle handle) const {
     if (handle.index >= reader_.objects().size()) {
         throw ArtifactError("artifact object handle is out of range");

src/artifact/binder.h

@@ -43,6 +43,10 @@ class Binder {
 
     ObjectHandle require_tensor(std::string_view name, NumericFormat format, StorageLayout layout,
                                 std::span<const std::uint64_t> shape);
+
+    // Declared tensor format lookup without binding or consuming the object; module binders use
+    // it to dispatch per-object weight encodings under one object-name contract.
+    [[nodiscard]] NumericFormat declared_format(std::string_view name) const;
     ObjectHandle require_resource(std::string_view name, ResourceEncoding encoding);
 
     [[nodiscard]] bool contains(std::string_view name) const noexcept;

src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_plan.h

@@ -26,6 +26,12 @@ void nvfp4_attn_input_w4a4_launch(const Tensor& x, const Weight& weight, Tensor&
                                   Tensor& k, Tensor& v, Nvfp4W4a4Workspace workspace,
                                   cudaStream_t stream);
 
+// Three-output DFlash2 route: the weight-only NVFP4 [6144,5120] parent writing q [4096,T],
+// k [1024,T], and v [1024,T] directly at every positive T (32-token chunks above the small-T
+// family). No transient workspace.
+void nvfp4_dflash2_attn_input(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k,
+                              Tensor& v, cudaStream_t stream);
+
 void nvfp4_attn_input_dispatch(const Tensor& x, const Weight& weight, Tensor& q, Tensor& gate,
                                Tensor& k, Tensor& v, LinearPolicy policy, WorkspaceArena* workspace,
                                cudaStream_t stream);

src/ops/attn_input_proj/nvfp4/nvfp4_dflash2_attn_input.cu

@@ -0,0 +1,105 @@
+#include "ops/attn_input_proj/nvfp4/nvfp4_attn_input_plan.h"
+
+#include "core/device.h"
+#include "ops/common/token_slices.h"
+#include "ops/linear/nvfp4/nvfp4_config.h"
+#include "ops/linear/nvfp4/nvfp4_gemv.cuh"
+#include "ops/linear/nvfp4/nvfp4_output.cuh"
+#include "ops/linear/nvfp4/nvfp4_small_t.cuh"
+
+#include <algorithm>
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+
+namespace ninfer::ops::detail {
+namespace {
+
+using Geometry = Nvfp4DFlash2QkvGeometry;
+using Output   = Nvfp4SplitOutput3<4096, 1024>;
+using Launch   = void (*)(const Tensor&, const Weight&, Tensor&, Tensor&, Tensor&, cudaStream_t);
+
+// The split-output epilogue owns the family's measured low-T warp mapping; see the four-output
+// route above for the crossover rationale.
+template <int ActiveTokens>
+struct Nvfp4DFlash2AttentionSmallTProductionSchedule {
+    static_assert(ActiveTokens >= kNvfp4FirstSmallT);
+    static_assert(ActiveTokens <= kNvfp4LastSmallT);
+    static constexpr int kWarpsPerCta       = ActiveTokens >= 17 ? 4 : (ActiveTokens >= 8 ? 16 : 8);
+    static constexpr int kValuesPerLane     = ActiveTokens >= 17 && ActiveTokens <= 20 ? 8 : 16;
+    static constexpr auto kActivationAccess = ActiveTokens <= 4
+                                                  ? Nvfp4SmallTActivationAccess::SharedPhase
+                                                  : Nvfp4SmallTActivationAccess::TokenPacked;
+    using Type =
+        Nvfp4SmallTSchedule<kWarpsPerCta, 1, 2, kValuesPerLane, ActiveTokens, 1, kActivationAccess,
+                            Nvfp4ScaleAccess::Direct, Nvfp4CodeCache::Default, 1,
+                            Nvfp4SmallTBlockOrder::RowsContiguous, 1>;
+};
+
+void launch_decode(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k, Tensor& v,
+                   cudaStream_t stream) {
+    using Schedule = typename Nvfp4LinearDecodeProductionSchedule<Geometry>::Type;
+
+    const Output output{static_cast<__nv_bfloat16*>(q.data), static_cast<__nv_bfloat16*>(k.data),
+                        static_cast<__nv_bfloat16*>(v.data)};
+    constexpr int kBlocks              = Geometry::kOutputRows / Schedule::kRowsPerCta;
+    const float inverse_weight_divisor = 1.0F / weight.weight_scale_divisor;
+    nvfp4_gemv_kernel<Geometry, Schedule, Nvfp4IdentityEpilogue, Output>
+        <<<dim3(kBlocks), dim3(Schedule::kThreads), 0, stream>>>(
+            static_cast<const __nv_bfloat16*>(x.data),
+            static_cast<const std::uint8_t*>(weight.qdata),
+            static_cast<const std::uint8_t*>(weight.scales), inverse_weight_divisor,
+            Nvfp4IdentityEpilogue{}, output);
+    CUDA_CHECK(cudaGetLastError());
+}
+
+template <int ActiveTokens>
+void launch_exact(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k, Tensor& v,
+                  cudaStream_t stream) {
+    using Schedule            = typename Nvfp4DFlash2AttentionSmallTProductionSchedule<ActiveTokens>::Type;
+    constexpr int kTokenTiles = (ActiveTokens + Schedule::kTokenTile - 1) / Schedule::kTokenTile;
+    constexpr int kBlocks     = (Geometry::kOutputRows / Schedule::kRowsPerCta) * kTokenTiles;
+
+    const Output output{static_cast<__nv_bfloat16*>(q.data), static_cast<__nv_bfloat16*>(k.data),
+                        static_cast<__nv_bfloat16*>(v.data)};
+    const float inverse_weight_divisor = 1.0F / weight.weight_scale_divisor;
+    nvfp4_small_t_kernel<Geometry, ActiveTokens, Schedule, Nvfp4IdentityEpilogue, Output>
+        <<<dim3(kBlocks), dim3(Schedule::kThreads), 0, stream>>>(
+            static_cast<const __nv_bfloat16*>(x.data),
+            static_cast<const std::uint8_t*>(weight.qdata),
+            static_cast<const std::uint8_t*>(weight.scales), inverse_weight_divisor,
+            Nvfp4IdentityEpilogue{}, output);
+    CUDA_CHECK(cudaGetLastError());
+}
+
+template <std::size_t... Offsets>
+constexpr auto make_launchers(std::index_sequence<Offsets...>) {
+    return std::array<Launch, sizeof...(Offsets)>{
+        &launch_exact<kNvfp4FirstSmallT + static_cast<int>(Offsets)>...};
+}
+
+constexpr auto kLaunchers =
+    make_launchers(std::make_index_sequence<kNvfp4LastSmallT - kNvfp4FirstSmallT + 1>{});
+
+} // namespace
+
+void nvfp4_dflash2_attn_input(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k,
+                              Tensor& v, cudaStream_t stream) {
+    constexpr std::int32_t kChunk = kNvfp4LastSmallT;
+    for (std::int32_t token_begin = 0; token_begin < x.ne[1]; token_begin += kChunk) {
+        const std::int32_t active = std::min(kChunk, x.ne[1] - token_begin);
+        const Tensor x_slice      = x.slice(1, token_begin, active);
+        Tensor q_slice            = q.slice(1, token_begin, active);
+        Tensor k_slice            = k.slice(1, token_begin, active);
+        Tensor v_slice            = v.slice(1, token_begin, active);
+        if (active == 1) {
+            launch_decode(x_slice, weight, q_slice, k_slice, v_slice, stream);
+        } else {
+            kLaunchers[active - kNvfp4FirstSmallT](x_slice, weight, q_slice, k_slice, v_slice,
+                                                   stream);
+        }
+    }
+}
+
+} // namespace ninfer::ops::detail