diff --git a/README.md b/README.md index 6b26b6e551..6c22e74063 100644 --- a/README.md +++ b/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. diff --git a/bench/ops/candidate_selector_bench.cu b/bench/ops/candidate_selector_bench.cu index 496eb01515..01f2c3c273 100644 --- a/bench/ops/candidate_selector_bench.cu +++ b/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(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; diff --git a/include/ninfer/ops/attn_input_proj.h b/include/ninfer/ops/attn_input_proj.h index ff5cb57812..d0ab5c0680 100644 --- a/include/ninfer/ops/attn_input_proj.h +++ b/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); diff --git a/include/ninfer/ops/candidate_selector.h b/include/ninfer/ops/candidate_selector.h index 013dd268fd..897926cbc3 100644 --- a/include/ninfer/ops/candidate_selector.h +++ b/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); diff --git a/include/ninfer/ops/linear_swiglu.h b/include/ninfer/ops/linear_swiglu.h index 0b977edd68..3635e3fe5e 100644 --- a/include/ninfer/ops/linear_swiglu.h +++ b/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); diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 9e991aa331..691f9470c2 100644 --- a/src/CMakeLists.txt +++ b/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 diff --git a/src/artifact/binder.cpp b/src/artifact/binder.cpp index 3c2e934a25..43f5971169 100644 --- a/src/artifact/binder.cpp +++ b/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(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"); diff --git a/src/artifact/binder.h b/src/artifact/binder.h index 4b77279bef..c09cbcba00 100644 --- a/src/artifact/binder.h +++ b/src/artifact/binder.h @@ -43,6 +43,10 @@ class Binder { ObjectHandle require_tensor(std::string_view name, NumericFormat format, StorageLayout layout, std::span 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; diff --git a/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_plan.h b/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_plan.h index 4bbd745b06..b47ba940f5 100644 --- a/src/ops/attn_input_proj/nvfp4/nvfp4_attn_input_plan.h +++ b/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); diff --git a/src/ops/attn_input_proj/nvfp4/nvfp4_dflash2_attn_input.cu b/src/ops/attn_input_proj/nvfp4/nvfp4_dflash2_attn_input.cu new file mode 100644 index 0000000000..ef49b058e9 --- /dev/null +++ b/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 +#include +#include +#include +#include + +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 +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; +}; + +void launch_decode(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k, Tensor& v, + cudaStream_t stream) { + using Schedule = typename Nvfp4LinearDecodeProductionSchedule::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 + <<>>( + static_cast(x.data), + static_cast(weight.qdata), + static_cast(weight.scales), inverse_weight_divisor, + Nvfp4IdentityEpilogue{}, output); + CUDA_CHECK(cudaGetLastError()); +} + +template +void launch_exact(const Tensor& x, const Weight& weight, Tensor& q, Tensor& k, Tensor& v, + cudaStream_t stream) { + using Schedule = typename Nvfp4DFlash2AttentionSmallTProductionSchedule::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 + <<>>( + static_cast(x.data), + static_cast(weight.qdata), + static_cast(weight.scales), inverse_weight_divisor, + Nvfp4IdentityEpilogue{}, output); + CUDA_CHECK(cudaGetLastError()); +} + +template +constexpr auto make_launchers(std::index_sequence) { + return std::array{ + &launch_exact(Offsets)>...}; +} + +constexpr auto kLaunchers = + make_launchers(std::make_index_sequence{}); + +} // 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 diff --git a/src/ops/candidate_selector/bf16/candidate_selector_path.cu b/src/ops/candidate_selector/bf16/candidate_selector_path.cu index 171da76dc1..cfee8afcc5 100644 --- a/src/ops/candidate_selector/bf16/candidate_selector_path.cu +++ b/src/ops/candidate_selector/bf16/candidate_selector_path.cu @@ -154,8 +154,8 @@ __global__ __launch_bounds__(32) void selector_lattice_walk_kernel(DeviceArgs a, void candidate_selector_path_launch(SelectorRoute route, const Tensor& candidate_ids, const Tensor& unary_scores, const Tensor& projected_hidden, - const Tensor& anchors, const Tensor& predecessor_codebook, - const Tensor& successor_codebook, const Tensor& base_positions, + const Tensor& anchors, const Weight& predecessor_codebook, + const Weight& successor_codebook, const Tensor& base_positions, const SamplingConfig* configs, Tensor& drafts, Tensor& proposal_q, const SelectorWorkspace& workspace, cudaStream_t stream) { @@ -163,8 +163,8 @@ void candidate_selector_path_launch(SelectorRoute route, const Tensor& candidate static_cast(unary_scores.data), static_cast(projected_hidden.data), static_cast(anchors.data), - static_cast(predecessor_codebook.data), - static_cast(successor_codebook.data), + static_cast(predecessor_codebook.qdata), + static_cast(successor_codebook.qdata), static_cast(base_positions.data), configs, static_cast(drafts.data), diff --git a/src/ops/candidate_selector/bf16/candidate_selector_path_kernels.h b/src/ops/candidate_selector/bf16/candidate_selector_path_kernels.h index af68c0f9ad..62fdba3517 100644 --- a/src/ops/candidate_selector/bf16/candidate_selector_path_kernels.h +++ b/src/ops/candidate_selector/bf16/candidate_selector_path_kernels.h @@ -4,8 +4,8 @@ namespace ninfer::ops::detail { void candidate_selector_path_launch(SelectorRoute route, const Tensor& candidate_ids, const Tensor& unary_scores, const Tensor& projected_hidden, - const Tensor& anchors, const Tensor& predecessor_codebook, - const Tensor& successor_codebook, const Tensor& base_positions, + const Tensor& anchors, const Weight& predecessor_codebook, + const Weight& successor_codebook, const Tensor& base_positions, const SamplingConfig* configs, Tensor& drafts, Tensor& proposal_q, const SelectorWorkspace& workspace, cudaStream_t stream); diff --git a/src/ops/candidate_selector/bf16/candidate_selector_path_plan.cpp b/src/ops/candidate_selector/bf16/candidate_selector_path_plan.cpp index 9d188175b3..5b4754d54f 100644 --- a/src/ops/candidate_selector/bf16/candidate_selector_path_plan.cpp +++ b/src/ops/candidate_selector/bf16/candidate_selector_path_plan.cpp @@ -19,8 +19,8 @@ const char* candidate_selector_path_route_name(int steps, int batch) { void candidate_selector_path_dispatch(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) { @@ -28,16 +28,22 @@ void candidate_selector_path_dispatch(const Tensor& candidate_ids, const Tensor& auto scope = workspace.scope(); const auto scratch = allocate_selector_workspace(workspace, route, candidate_ids.ne[1], candidate_ids.ne[2]); - const Tensor* live[]{ - &candidate_ids, &unary_scores, &projected_hidden, &anchors, &predecessor_codebook, - &successor_codebook, &base_positions, &drafts, &proposal_q}; + const Tensor tensors[]{ + candidate_ids, unary_scores, projected_hidden, anchors, + base_positions, drafts, proposal_q, + }; for (const auto* work : {&scratch.edges}) { if (!work->data) continue; const auto begin = reinterpret_cast(work->data), end = begin + work->bytes(); - for (const auto* tensor : live) { - const auto tb = reinterpret_cast(tensor->data); - if (begin < tb + tensor->bytes() && tb < end) + for (const auto& tensor : tensors) { + const auto tb = reinterpret_cast(tensor.data); + if (begin < tb + tensor.bytes() && tb < end) + throw std::invalid_argument("selector workspace overlaps operand"); + } + for (const auto* codebook : {&predecessor_codebook, &successor_codebook}) { + const auto cb = reinterpret_cast(codebook->qdata); + if (begin < cb + codebook->payload_bytes && cb < end) throw std::invalid_argument("selector workspace overlaps operand"); } const auto cb = reinterpret_cast(configs); diff --git a/src/ops/candidate_selector/bf16/candidate_selector_path_plan.h b/src/ops/candidate_selector/bf16/candidate_selector_path_plan.h index 92ca60caaa..ca9d8c0f12 100644 --- a/src/ops/candidate_selector/bf16/candidate_selector_path_plan.h +++ b/src/ops/candidate_selector/bf16/candidate_selector_path_plan.h @@ -27,8 +27,8 @@ SelectorWorkspace allocate_selector_workspace(Allocator& allocator, SelectorRout void candidate_selector_path_dispatch(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); diff --git a/src/ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.cu b/src/ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.cu new file mode 100644 index 0000000000..4ba394b204 --- /dev/null +++ b/src/ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.cu @@ -0,0 +1,265 @@ +// NVFP4 codebook route of candidate_selector_path: the [248320,256] codebooks are weight-only +// NVFP4 (packed E2M1 pairs with one E4M3 scale per 16-rank group and a payload divisor). The +// walk structure, draw, and lattice are shared with the BF16 family; only row access decodes. +// Successor rows stage raw codes and scales through cp_async and decode to BF16 after the wait; +// predecessor elements decode inline at their point of use. +#include "ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.h" + +#include "core/device.h" +#include "ops/common/memory.cuh" +#include "ops/common/warp.cuh" +#include "ops/kernel/sampling_device.cuh" +#include "ops/linear/nvfp4/nvfp4_codec.cuh" +#include +#include + +namespace ninfer::ops::detail { +namespace { + +constexpr int kCandidates = 16, kRank = 256; + +constexpr int kCodeBytesPerRow = kRank / 2; +constexpr int kScaleBytesPerRow = kRank / 16; + +// The codebook scale plane follows the registered K16M128x4 blocked arrangement shared by every +// NVFP4 payload: 512-byte tiles of 128 rows x 4 consecutive groups. +__device__ __forceinline__ int codebook_scale_offset(int token, int group) { + const int in_tile = (token % 32) * 16 + ((token % 128) / 32) * 4; + return (token / 128) * 4 * 512 + (group / 4) * 512 + in_tile + group % 4; +} + +struct DeviceArgs { + const std::int32_t* ids; + const float* unary; + const __nv_bfloat16* hidden; + const std::int32_t* anchors; + const std::uint8_t* predecessor_codes; + const std::uint8_t* predecessor_scales; + float predecessor_inverse_divisor; + const std::uint8_t* successor_codes; + const std::uint8_t* successor_scales; + float successor_inverse_divisor; + const std::int32_t* positions; + const SamplingConfig* configs; + std::int32_t* drafts; + float* q; + int steps; +}; + +struct alignas(16) SelectorShared { + std::uint8_t successor_staged[kCandidates * kCodeBytesPerRow]; + std::uint8_t successor_staged_scales[kCandidates * kScaleBytesPerRow]; + __nv_bfloat16 successors[kCandidates * kRank]; + float product[kRank]; + float edge[kCandidates]; + int predecessor, base_position; + float temperature; + unsigned long long seed; +}; + +__device__ __forceinline__ float decode_rank(const std::uint8_t* codes, + const std::uint8_t* scales, int token, int rank) { + const float2 pair = decode_nvfp4_e2m1x2(codes[token * kCodeBytesPerRow + (rank >> 1)]); + const float code = (rank & 1) != 0 ? pair.y : pair.x; + return code * decode_nvfp4_e4m3(scales[codebook_scale_offset(token, rank >> 4)]); +} + +// The probabilities written here are the same FP32 values consumed by the draw. +__device__ int draw_rank(float edge, float temperature, unsigned long long seed, int position, + float* q) { + const int lane = threadIdx.x & 31; + const float maximum = warp_max(edge); + if (temperature <= 0.0F) { + const unsigned winners = + __ballot_sync(kFullWarpMask, lane < kCandidates && edge == maximum); + const int selected = winners == 0 ? 0 : __ffs(winners) - 1; + if (lane < kCandidates) q[lane] = lane == selected ? 1.0F : 0.0F; + return selected; + } + const float weight = lane < kCandidates ? __expf((edge - maximum) / temperature) : 0.0F; + const float probability = weight / warp_sum(weight); + if (lane < kCandidates) q[lane] = probability; + float uniform = + lane == 0 ? sampling_uniform(seed, position, kSamplePurposeDFlash2Proposal, 0U) : 0.0F; + uniform = __shfl_sync(kFullWarpMask, uniform, 0); + float cumulative = probability; +#pragma unroll + for (int offset = 1; offset < kCandidates; offset *= 2) { + const float previous = __shfl_up_sync(kFullWarpMask, cumulative, offset); + if (lane >= offset) cumulative += previous; + } + const unsigned hits = __ballot_sync(kFullWarpMask, lane < kCandidates && uniform < cumulative); + return hits ? __ffs(hits) - 1 : kCandidates - 1; +} + +__device__ void score_row(const DeviceArgs& a, int column, SelectorShared& shared) { + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + int token = lane == 0 ? a.ids[column * kCandidates + warp] : 0; + token = __shfl_sync(kFullWarpMask, token, 0); + // Stage one packed successor row per warp: 128 code bytes (row-major) and the row's four + // 4-byte scale groups, one per 512-byte K16M128x4 tile, gathered into natural group order. + if (lane < kCodeBytesPerRow / 16) { + cp_async<16, Cache::cg>( + &shared.successor_staged[warp * kCodeBytesPerRow + lane * 16], + a.successor_codes + static_cast(token) * kCodeBytesPerRow + lane * 16); + } + if (lane < 4) { + const int in_tile = (token % 32) * 16 + ((token % 128) / 32) * 4; + cp_async<4>(&shared.successor_staged_scales[warp * kScaleBytesPerRow + lane * 4], + a.successor_scales + (static_cast(token / 128) * 4 + lane) * + 512 + + in_tile); + } + cp_commit(); + // Publish the preceding draw before reading its token. Successor prefetch is independent. + __syncthreads(); + if (tid < kRank) { + shared.product[tid] = + decode_rank(a.predecessor_codes, a.predecessor_scales, shared.predecessor, tid) * + a.predecessor_inverse_divisor * + __bfloat162float(a.hidden[static_cast(column) * kRank + tid]); + } + cp_wait<0>(); + __syncthreads(); + for (int item = tid; item < kCandidates * kRank; item += 512) { + const int candidate = item / kRank, rank = item % kRank; + const float2 pair = decode_nvfp4_e2m1x2( + shared.successor_staged[candidate * kCodeBytesPerRow + (rank >> 1)]); + const float code = (rank & 1) != 0 ? pair.y : pair.x; + const float scale = + decode_nvfp4_e4m3(shared.successor_staged_scales[candidate * kScaleBytesPerRow + + (rank >> 4)]); + shared.successors[item] = + __float2bfloat16_rn(code * scale * a.successor_inverse_divisor); + } + __syncthreads(); + { + const int c = warp; + float sum = 0; +#pragma unroll + for (int r = lane; r < kRank; r += 32) + sum = fmaf(shared.product[r], __bfloat162float(shared.successors[c * kRank + r]), sum); + sum = warp_reduce_sum(sum); + if (lane == 0) shared.edge[c] = a.unary[column * kCandidates + c] + sum; + } + __syncthreads(); +} + +__global__ __launch_bounds__(512, 1) void selector_walk_nvfp4_kernel(DeviceArgs a) { + __shared__ SelectorShared shared; + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31, batch = blockIdx.x; + if (tid == 0) { + shared.predecessor = a.anchors[batch]; + shared.base_position = a.positions[batch]; + shared.temperature = a.configs[batch].temperature; + shared.seed = a.configs[batch].seed; + } +#pragma unroll 1 + for (int step = 0; step < a.steps; ++step) { + const int column = batch * a.steps + step; + score_row(a, column, shared); + if (warp == 0) { + const float edge = lane < kCandidates ? shared.edge[lane] : -CUDART_INF_F; + const int selected = draw_rank(edge, shared.temperature, shared.seed, + shared.base_position + step, a.q + column * kCandidates); + if (lane == 0) { + shared.predecessor = a.ids[column * kCandidates + selected]; + a.drafts[column] = shared.predecessor; + } + } + } +} + +__global__ __launch_bounds__(512, 2) void selector_lattice_nvfp4_kernel(DeviceArgs a, float* edges) { + const int column = blockIdx.x, p = blockIdx.y; + const int step = column % a.steps, batch = column / a.steps; + if (step == 0 && p != 0) return; + __shared__ SelectorShared shared; + if (threadIdx.x == 0) + shared.predecessor = step == 0 ? a.anchors[batch] : a.ids[(column - 1) * kCandidates + p]; + score_row(a, column, shared); + if (threadIdx.x < kCandidates) + edges[(static_cast(column) * kCandidates + p) * kCandidates + threadIdx.x] = + shared.edge[threadIdx.x]; +} + +__global__ __launch_bounds__(32) void selector_lattice_walk_nvfp4_kernel(DeviceArgs a, + const float* edges) { + const int lane = threadIdx.x, batch = blockIdx.x; + const auto seed = a.configs[batch].seed; + const float temperature = a.configs[batch].temperature; + const int position = a.positions[batch]; + int predecessor_rank = 0; +#pragma unroll 1 + for (int step = 0; step < a.steps; ++step) { + const int column = batch * a.steps + step; + const float edge = + lane < kCandidates + ? edges[(static_cast(column) * kCandidates + predecessor_rank) * + kCandidates + + lane] + : -CUDART_INF_F; + int selected = + draw_rank(edge, temperature, seed, position + step, a.q + column * kCandidates); + predecessor_rank = selected; + if (lane == 0) a.drafts[column] = a.ids[column * kCandidates + selected]; + } +} + +} // namespace + +void candidate_selector_path_nvfp4_launch(SelectorRoute route, const Tensor& candidate_ids, + const Tensor& unary_scores, + const Tensor& projected_hidden, const Tensor& anchors, + const Weight& predecessor_codebook, + const Weight& successor_codebook, + const Tensor& base_positions, + const SamplingConfig* configs, Tensor& drafts, + Tensor& proposal_q, const SelectorWorkspace& workspace, + cudaStream_t stream) { + const DeviceArgs args{static_cast(candidate_ids.data), + static_cast(unary_scores.data), + static_cast(projected_hidden.data), + static_cast(anchors.data), + static_cast(predecessor_codebook.qdata), + static_cast(predecessor_codebook.scales), + 1.0F / predecessor_codebook.weight_scale_divisor, + static_cast(successor_codebook.qdata), + static_cast(successor_codebook.scales), + 1.0F / successor_codebook.weight_scale_divisor, + static_cast(base_positions.data), + configs, + static_cast(drafts.data), + static_cast(proposal_q.data), + candidate_ids.ne[1]}; + if (route == SelectorRoute::Direct) { + selector_walk_nvfp4_kernel<<>>(args); + } else { + auto* edges = static_cast(workspace.edges.data); + selector_lattice_nvfp4_kernel<<>>(args, edges); + CUDA_CHECK(cudaGetLastError()); + selector_lattice_walk_nvfp4_kernel<<>>(args, edges); + } + CUDA_CHECK(cudaGetLastError()); +} + +void candidate_selector_path_nvfp4_dispatch(const Tensor& candidate_ids, const Tensor& unary_scores, + const Tensor& projected_hidden, const Tensor& anchors, + 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) { + const auto route = candidate_selector_path_route(candidate_ids.ne[1], candidate_ids.ne[2]); + auto scope = workspace.scope(); + const auto scratch = + allocate_selector_workspace(workspace, route, candidate_ids.ne[1], candidate_ids.ne[2]); + candidate_selector_path_nvfp4_launch(route, candidate_ids, unary_scores, projected_hidden, + anchors, predecessor_codebook, successor_codebook, + base_positions, configs, drafts, proposal_q, scratch, + stream); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.h b/src/ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.h new file mode 100644 index 0000000000..86afbe4294 --- /dev/null +++ b/src/ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.h @@ -0,0 +1,18 @@ +#pragma once + +#include "ops/candidate_selector/bf16/candidate_selector_path_plan.h" + +namespace ninfer::ops::detail { + +// Weight-only NVFP4 codebook route of the selector walk; see the BF16 family for the route +// structure and the shared SelectorWorkspace contract. +void candidate_selector_path_nvfp4_dispatch(const Tensor& candidate_ids, const Tensor& unary_scores, + const Tensor& projected_hidden, const Tensor& anchors, + 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); + +} // namespace ninfer::ops::detail diff --git a/src/ops/context_kv_materialize/context_kv_common.cuh b/src/ops/context_kv_materialize/context_kv_common.cuh new file mode 100644 index 0000000000..6877d53be9 --- /dev/null +++ b/src/ops/context_kv_materialize/context_kv_common.cuh @@ -0,0 +1,108 @@ +#pragma once + +// Shared device-side pieces of context_kv_materialize: the packed-column mapping, the per-layer +// weight/cache view, the fused key head store, and the key post-processing kernel entry. Both +// the W8 and NVFP4 projection kernels consume them; the view's code/scale pointers are raw and +// interpreted per weight format. + +#include "core/tensor.h" +#include "ninfer/ops/context_kv_materialize.h" +#include "ops/common/dflash_rope.cuh" +#include "ops/common/warp.cuh" + +#include +#include + +#include +#include + +#include + +namespace ninfer::ops::detail { + +constexpr int kContextKVRows = 1024; +constexpr int kContextKVHeadDim = 128; + +__device__ __forceinline__ int context_column(int column, int width, int prefix) { + return width == prefix ? column : column / prefix * width + column % prefix; +} + +union alignas(16) ContextKVBf16x8 { + uint4 raw; + __nv_bfloat162 pair[4]; +}; + +__device__ __forceinline__ int swizzle_128(int row, int column) { + return (((column >> 3) ^ (row & 7)) << 3) | (column & 7); +} + +struct DeviceLayerView { + const std::uint8_t* key_codes; + const std::uint8_t* key_scales; + const std::uint8_t* value_codes; + const std::uint8_t* value_scales; + float key_inverse_divisor; + float value_inverse_divisor; + const __nv_bfloat16* key_norm; + __nv_bfloat16* cache_k; + __half* cache_v; + std::int32_t padded_capacity; +}; + +struct DeviceLayers { + DeviceLayerView layer[kContextKVMaterializeLayers]; +}; + +__device__ __forceinline__ void store_key_head(const float* input, DeviceLayerView layer, + const int* positions, const int* slots, int column, + int width, int head) { + const int lane = threadIdx.x & 31; + const int j = lane * 2; + float x0 = input[j], x1 = input[j + 1], y0 = input[j + 64], y1 = input[j + 65]; + float sum = warp_reduce_sum(x0 * x0 + x1 * x1 + y0 * y0 + y1 * y1); + const float inverse = rsqrtf(__shfl_sync(0xffffffffU, sum, 0) / 128.0f + 1.e-6f); + x0 *= inverse * __bfloat162float(layer.key_norm[j]); + x1 *= inverse * __bfloat162float(layer.key_norm[j + 1]); + y0 *= inverse * __bfloat162float(layer.key_norm[j + 64]); + y1 *= inverse * __bfloat162float(layer.key_norm[j + 65]); + float sin0, cos0, sin1, cos1; + dflash_rope_sincos(positions, column, j, &sin0, &cos0); + dflash_rope_sincos(positions, column, j + 1, &sin1, &cos1); + const auto dst = 128LL * ((positions[column] & 2047) + (long long)layer.padded_capacity * + (head + 8 * slots[column / width])); + auto* out = reinterpret_cast<__nv_bfloat162*>(layer.cache_k + dst); + out[lane] = __floats2bfloat162_rn(x0 * cos0 - y0 * sin0, x1 * cos1 - y1 * sin1); + out[lane + 32] = __floats2bfloat162_rn(y0 * cos0 + x0 * sin0, y1 * cos1 + x1 * sin1); +} + +inline DeviceLayers make_device_layers( + const std::array& layers) { + DeviceLayers result{}; + for (int index = 0; index < static_cast(kContextKVMaterializeLayers); ++index) { + const ContextKVMaterializeLayerView& source = layers[static_cast(index)]; + const auto inverse = [](const Weight& weight) { + return weight.qtype == QType::NVFP4 ? 1.0F / weight.weight_scale_divisor : 1.0F; + }; + result.layer[index] = { + static_cast(source.key_weight.qdata), + static_cast(source.key_weight.scales), + static_cast(source.value_weight.qdata), + static_cast(source.value_weight.scales), + inverse(source.key_weight), + inverse(source.value_weight), + static_cast(source.key_norm_weight.data), + static_cast<__nv_bfloat16*>(source.cache.k.data), + static_cast<__half*>(source.cache.v.data), + static_cast(source.cache.padded_capacity), + }; + } + return result; +} + +void context_kv_key_post_launch(const Tensor& key_scratch, const Tensor& positions, + const Tensor& counts, const Tensor& state_slots, + const DeviceLayers& layers, std::int32_t batch_size, + std::int32_t width, + ContextKVMaterializeExecutionEnvelope envelope, cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/context_kv_materialize/context_kv_key_post.cu b/src/ops/context_kv_materialize/context_kv_key_post.cu new file mode 100644 index 0000000000..14248aac72 --- /dev/null +++ b/src/ops/context_kv_materialize/context_kv_key_post.cu @@ -0,0 +1,48 @@ +// Format-neutral key post-processing of context_kv_materialize: RMSNorm + DFlash rope over the +// FP32 key scratch and the store into the cyclic cache. Shared by every weight format. +#include "ops/context_kv_materialize/context_kv_common.cuh" + +#include "core/device.h" + +namespace ninfer::ops::detail { +namespace { + +__global__ __launch_bounds__(256) void context_kv_key_post_kernel( + const float* __restrict__ key_scratch, const std::int32_t* __restrict__ positions, + const std::int32_t* __restrict__ counts, const std::int32_t* __restrict__ state_slots, + DeviceLayers layers, std::int32_t batch_size, std::int32_t width, std::int32_t min_count, + std::int32_t max_count) { + const int packed_column = static_cast(blockIdx.x); + const int physical_column = context_column(packed_column, width, max_count); + const int layer_index = static_cast(blockIdx.y); + const int batch = physical_column / width; + const int local = physical_column % width; + const int count = counts[batch]; + if (count < min_count || count > max_count || local >= count) return; + const auto layer = layers.layer[layer_index]; + const int head = threadIdx.x >> 5; + const float* input = key_scratch + + kContextKVRows * (packed_column + max_count * batch_size * layer_index) + + head * kContextKVHeadDim; + store_key_head(input, layer, positions, state_slots, physical_column, width, head); +} + +} // namespace + +void context_kv_key_post_launch(const Tensor& key_scratch, const Tensor& positions, + const Tensor& counts, const Tensor& state_slots, + const DeviceLayers& layers, std::int32_t batch_size, + std::int32_t width, + ContextKVMaterializeExecutionEnvelope envelope, + cudaStream_t stream) { + context_kv_key_post_kernel<<(kContextKVMaterializeLayers)), + 256, 0, stream>>>( + static_cast(key_scratch.data), static_cast(positions.data), + static_cast(counts.data), static_cast(state_slots.data), layers, + batch_size, width, static_cast(envelope.min_count), + static_cast(envelope.max_count)); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/context_kv_materialize/context_kv_materialize.cpp b/src/ops/context_kv_materialize/context_kv_materialize.cpp index ef371eaa9d..4576cad55a 100644 --- a/src/ops/context_kv_materialize/context_kv_materialize.cpp +++ b/src/ops/context_kv_materialize/context_kv_materialize.cpp @@ -4,6 +4,7 @@ #include "ops/context_kv_materialize/launch.h" #include +#include #include #include #include @@ -34,6 +35,28 @@ void require_tensor(const Tensor& tensor, DType dtype, std::int32_t n0, std::int } void require_weight(const Weight& weight, const char* name) { + if (weight.qtype == QType::NVFP4) { + // The NVFP4 key/value parents are 128-row-aligned slices of the draft module's packed + // query_key_value payload: the code and scale planes keep the registered arrangement but + // no longer sit at the canonical single-payload offsets, so the geometry is checked here + // instead of validate_nvfp4_weight. + constexpr std::uint64_t kCodeBytes = + static_cast(kKVSize) * static_cast(kHidden / 2); + constexpr std::uint64_t kScaleBytes = + static_cast(kKVSize) * static_cast(kHidden / 16); + if (weight.layout != QuantLayout::BlockScaleK16M128x4 || + weight.scale_dtype != DType::FP8_E4M3FN || weight.group != 16 || + weight.group_size != 16 || weight.ndim != 2 || weight.n != kKVSize || + weight.k != kHidden || weight.shape[0] != kKVSize || weight.shape[1] != kHidden || + weight.qhigh != nullptr || weight.high_plane_bytes != 0 || + weight.payload_bytes < kCodeBytes + kScaleBytes + sizeof(float) || + !aligned_to(weight.qdata, 16) || !aligned_to(weight.scales, 16) || + weight.scales < weight.qdata || + !std::isfinite(weight.weight_scale_divisor) || weight.weight_scale_divisor <= 0.0F) { + throw std::invalid_argument(std::string(kOp) + ": invalid " + name); + } + return; + } constexpr std::uint64_t kCodeBytes = static_cast(kKVSize) * static_cast(kHidden); constexpr std::uint64_t kScaleBytes = @@ -155,6 +178,11 @@ void context_kv_materialize( Tensor key_scratch; if (detail::context_kv_materialize_uses_scratch(route)) key_scratch = allocate_key_scratch(workspace, columns); + if (layers.front().key_weight.qtype == QType::NVFP4) { + detail::context_kv_materialize_nvfp4_launch(context, positions, counts, state_slots, + layers, envelope, route, key_scratch, stream); + return; + } detail::context_kv_materialize_launch(context, positions, counts, state_slots, layers, envelope, route, key_scratch, stream); } diff --git a/src/ops/context_kv_materialize/launch.h b/src/ops/context_kv_materialize/launch.h index 7194cd48ea..bb5a5b5270 100644 --- a/src/ops/context_kv_materialize/launch.h +++ b/src/ops/context_kv_materialize/launch.h @@ -37,4 +37,12 @@ void context_kv_materialize_launch( ContextKVMaterializeExecutionEnvelope envelope, ContextKVMaterializeRoute route, const Tensor& key_scratch, cudaStream_t stream); +// NVFP4 key/value parents: one MMA family serves every routed column count with the same +// envelope filtering, key scratch, and cache stores as the W8 launch. +void context_kv_materialize_nvfp4_launch( + const Tensor& context, const Tensor& positions, const Tensor& counts, const Tensor& state_slots, + const std::array& layers, + ContextKVMaterializeExecutionEnvelope envelope, ContextKVMaterializeRoute route, + const Tensor& key_scratch, cudaStream_t stream); + } // namespace ninfer::ops::detail diff --git a/src/ops/context_kv_materialize/materialize.cu b/src/ops/context_kv_materialize/materialize.cu index a6a79a2792..d4ed1fa989 100644 --- a/src/ops/context_kv_materialize/materialize.cu +++ b/src/ops/context_kv_materialize/materialize.cu @@ -3,8 +3,7 @@ #include "ops/common/memory.cuh" #include "ops/common/mma.cuh" #include "ops/linear/w8/w8_small_t_mma.cuh" -#include "ops/common/warp.cuh" -#include "ops/common/dflash_rope.cuh" +#include "ops/context_kv_materialize/context_kv_common.cuh" #include #include @@ -14,56 +13,6 @@ constexpr int kLayers = static_cast(kContextKVMaterializeLayers); constexpr int kRows = 1024; constexpr int kHeadDim = 128; -__device__ __forceinline__ int context_column(int column, int width, int prefix) { - return width == prefix ? column : column / prefix * width + column % prefix; -} - -struct DeviceLayerView { - const std::uint8_t* key_codes; - const std::uint8_t* key_scales; - const std::uint8_t* value_codes; - const std::uint8_t* value_scales; - const __nv_bfloat16* key_norm; - __nv_bfloat16* cache_k; - __half* cache_v; - std::int32_t padded_capacity; -}; - -struct DeviceLayers { - DeviceLayerView layer[kContextKVMaterializeLayers]; -}; - -__device__ __forceinline__ void store_key_head(const float* input, DeviceLayerView layer, - const int* positions, const int* slots, int column, - int width, int head) { - const int lane = threadIdx.x & 31; - const int j = lane * 2; - float x0 = input[j], x1 = input[j + 1], y0 = input[j + 64], y1 = input[j + 65]; - float sum = warp_reduce_sum(x0 * x0 + x1 * x1 + y0 * y0 + y1 * y1); - const float inverse = rsqrtf(__shfl_sync(0xffffffffU, sum, 0) / 128.0f + 1.e-6f); - x0 *= inverse * __bfloat162float(layer.key_norm[j]); - x1 *= inverse * __bfloat162float(layer.key_norm[j + 1]); - y0 *= inverse * __bfloat162float(layer.key_norm[j + 64]); - y1 *= inverse * __bfloat162float(layer.key_norm[j + 65]); - float sin0, cos0, sin1, cos1; - dflash_rope_sincos(positions, column, j, &sin0, &cos0); - dflash_rope_sincos(positions, column, j + 1, &sin1, &cos1); - const auto dst = 128LL * ((positions[column] & 2047) + (long long)layer.padded_capacity * - (head + 8 * slots[column / width])); - auto* out = reinterpret_cast<__nv_bfloat162*>(layer.cache_k + dst); - out[lane] = __floats2bfloat162_rn(x0 * cos0 - y0 * sin0, x1 * cos1 - y1 * sin1); - out[lane + 32] = __floats2bfloat162_rn(y0 * cos0 + x0 * sin0, y1 * cos1 + x1 * sin1); -} - -union alignas(16) Bf16x8 { - uint4 raw; - __nv_bfloat162 pair[4]; -}; - -__device__ __forceinline__ int swizzle_128(int row, int column) { - return (((column >> 3) ^ (row & 7)) << 3) | (column & 7); -} - template union alignas(16) MaterializeStorage { struct { @@ -154,7 +103,7 @@ __global__ __launch_bounds__(Rows / 16 * ColumnWarps * 32, 1) void context_kv_mm const int chunk = item - row * kChunksPerRow; const int col = chunk * 8; const uint2 packed = *reinterpret_cast(&mainloop.codes[row][col]); - Bf16x8 decoded; + ContextKVBf16x8 decoded; #pragma unroll for (int pair = 0; pair < 4; ++pair) { const unsigned word = (pair < 2 ? packed.x : packed.y) >> ((pair & 1) * 16); @@ -393,45 +342,6 @@ void launch_grouped(const Tensor& x, const Tensor& positions, const Tensor& coun CUDA_CHECK(cudaGetLastError()); } -__global__ __launch_bounds__(256) void context_kv_key_post_kernel( - const float* __restrict__ key_scratch, const std::int32_t* __restrict__ positions, - const std::int32_t* __restrict__ counts, const std::int32_t* __restrict__ state_slots, - DeviceLayers layers, std::int32_t batch_size, std::int32_t width, std::int32_t min_count, - std::int32_t max_count) { - const int packed_column = static_cast(blockIdx.x); - const int physical_column = context_column(packed_column, width, max_count); - const int layer_index = static_cast(blockIdx.y); - const int batch = physical_column / width; - const int local = physical_column % width; - const int count = counts[batch]; - if (count < min_count || count > max_count || local >= count) return; - const auto layer = layers.layer[layer_index]; - const int head = threadIdx.x >> 5; - const float* input = key_scratch + - kRows * (packed_column + max_count * batch_size * layer_index) + - head * kHeadDim; - store_key_head(input, layer, positions, state_slots, physical_column, width, head); -} - -DeviceLayers make_device_layers( - const std::array& layers) { - DeviceLayers result{}; - for (int index = 0; index < kLayers; ++index) { - const ContextKVMaterializeLayerView& source = layers[static_cast(index)]; - result.layer[index] = { - static_cast(source.key_weight.qdata), - static_cast(source.key_weight.scales), - static_cast(source.value_weight.qdata), - static_cast(source.value_weight.scales), - static_cast(source.key_norm_weight.data), - static_cast<__nv_bfloat16*>(source.cache.k.data), - static_cast<__half*>(source.cache.v.data), - static_cast(source.cache.padded_capacity), - }; - } - return result; -} - } // namespace void context_kv_materialize_launch( @@ -471,12 +381,8 @@ void context_kv_materialize_launch( key_scratch, stream); break; } - context_kv_key_post_kernel<<>>( - static_cast(key_scratch.data), static_cast(positions.data), - static_cast(counts.data), static_cast(state_slots.data), - device_layers, context.ne[2], context.ne[1], envelope.min_count, envelope.max_count); - CUDA_CHECK(cudaGetLastError()); + context_kv_key_post_launch(key_scratch, positions, counts, state_slots, device_layers, + context.ne[2], context.ne[1], envelope, stream); } } // namespace ninfer::ops::detail diff --git a/src/ops/context_kv_materialize/materialize_nvfp4.cu b/src/ops/context_kv_materialize/materialize_nvfp4.cu new file mode 100644 index 0000000000..fd8dd2fc7c --- /dev/null +++ b/src/ops/context_kv_materialize/materialize_nvfp4.cu @@ -0,0 +1,324 @@ +// NVFP4 projection kernels of context_kv_materialize: the weight-only NVFP4 [1024,5120] +// key/value parents (row slices of the draft module's query_key_value payload). One MMA family +// serves every column count: e2m1 codes stage raw and decode to exactly-representable BF16, the +// stored E4M3 scales apply in FP32 per 16-value group after each group's MMA accumulation, and +// the payload divisor folds into the captured scales. The activation staging, envelope filtering, +// key scratch, and cache stores are shared with the W8 family through context_kv_common.cuh. +#include "ops/context_kv_materialize/launch.h" +#include "core/device.h" +#include "ops/common/memory.cuh" +#include "ops/common/mma.cuh" +#include "ops/context_kv_materialize/context_kv_common.cuh" +#include "ops/linear/nvfp4/nvfp4_codec.cuh" +#include +#include + +namespace ninfer::ops::detail { +namespace { + +constexpr int kHidden = 5120, kRows = kContextKVRows, kHeadDim = kContextKVHeadDim; + +template +union alignas(16) MaterializeNvfp4Storage { + struct { + __nv_bfloat16 code_values[Rows][BlockK]; + __nv_bfloat16 activations[Columns][BlockK]; + std::uint8_t codes[Rows][BlockK / 2]; + std::uint8_t scales[Rows][BlockK / 16]; + } mainloop; + + float scores[Columns][Rows]; +}; + +template +__global__ __launch_bounds__(Rows / 16 * ColumnWarps * 32, 1) void context_kv_mma_nvfp4_kernel( + const __nv_bfloat16* hidden, const int* positions, const int* counts, const int* slots, + DeviceLayers layers, float* key_scratch, int width, int batch, int min_count, int max_count) { + constexpr int kBlockRows = Rows, kBlockK = BlockK; + constexpr int kBlockColumns = Columns; + constexpr int kColumnWarps = ColumnWarps; + constexpr int kWarps = Rows / 16 * kColumnWarps, kThreads = kWarps * 32; + constexpr int kWarpColumns = Columns / kColumnWarps, kTokenMmas = kWarpColumns / 8; + constexpr int kKTiles = kHidden / kBlockK; + constexpr int kGroupsPerTile = kBlockK / 16; + static_assert((Rows == 64 || Rows == 128) && Columns % (8 * ColumnWarps) == 0); + static_assert(kHidden % kBlockK == 0 && kThreads <= 1024); + static_assert((kBlockK % 16) == 0 && kBlockK <= 128); + const int tid = threadIdx.x, warp = tid >> 5, lane = tid & 31; + const int warp_row = warp / kColumnWarps, warp_col = warp % kColumnWarps; + const int gid = lane >> 2, lid = lane & 3; + const int a_matrix = lane >> 3, a_rowoff = (lane & 7) + ((a_matrix & 1) << 3); + const int a_coloff = (a_matrix >> 1) << 3, b_row = lane & 7, b_coloff = ((lane >> 3) & 1) << 3; + const int column_begin = blockIdx.y * Columns, + live_columns = min(Columns, max_count * batch - column_begin); + const int row_begin = blockIdx.x * Rows, layer_index = blockIdx.z >> 1; + const bool value = (blockIdx.z & 1) != 0; + const auto layer = layers.layer[layer_index]; + const auto* weight_codes = value ? layer.value_codes : layer.key_codes; + const auto* weight_scales = value ? layer.value_scales : layer.key_scales; + const float inverse_divisor = value ? layer.value_inverse_divisor : layer.key_inverse_divisor; + extern __shared__ __align__(16) unsigned char shared_bytes[]; + auto& storage = *reinterpret_cast*>(shared_bytes); + auto& mainloop = storage.mainloop; + { + float accumulators[kTokenMmas][4] = {}; + + const auto stage_activation = [&](int k_tile) { + const int k_begin = k_tile * kBlockK; + constexpr int kItems = kBlockColumns * (kBlockK / 8); + for (int item = tid; item < kItems; item += kThreads) { + const int column = item / (kBlockK / 8); + const int k8 = item - column * (kBlockK / 8); + auto* destination = &mainloop.activations[column][swizzle_128(column, k8 * 8)]; + if (column < live_columns) { + cp_async<16, Cache::ca>( + destination, + hidden + + static_cast(context_column(column_begin + column, width, + max_count)) * + kHidden + + k_begin + k8 * 8); + } else { + cp_async_zfill<16, Cache::ca>(destination, hidden + k_begin + k8 * 8, 0); + } + } + }; + + const auto stage_weight = [&](int k_tile) { + const int k_begin = k_tile * kBlockK; + constexpr int kChunks = kBlockRows * (kBlockK / 32); + for (int item = tid; item < kChunks; item += kThreads) { + const int local_row = item / (kBlockK / 32); + const int chunk = item - local_row * (kBlockK / 32); + cp_async<16, Cache::cg>( + &mainloop.codes[local_row][chunk * 16], + weight_codes + static_cast(row_begin + local_row) * (kHidden / 2) + + k_begin / 2 + chunk * 16); + } + // The scale plane is the registered K16M128x4 blocked arrangement: each k-tile's + // groups span BlockK/64 consecutive 512-byte tiles, four group bytes per row and tile. + for (int local_row = tid; local_row < kBlockRows; local_row += kThreads) { + const int row = row_begin + local_row; + const int in_tile = (row % 32) * 16 + ((row % 128) / 32) * 4; + const std::int64_t tile_base = + static_cast(row / 128) * (kHidden / 64) + k_begin / 64; +#pragma unroll + for (int half = 0; half < kBlockK / 64; ++half) { + cp_async<4>(&mainloop.scales[local_row][half * 4], + weight_scales + (tile_base + half) * 512 + in_tile); + } + } + }; + + // E2M1 magnitudes are exactly representable in BF16; decode pairs straight to registers + // and keep the stored E4M3 scale application in FP32. + const auto decode_e2m1_codes = [&]() { + constexpr int kChunksPerRow = kBlockK / 32; + for (int item = tid; item < kBlockRows * kChunksPerRow; item += kThreads) { + const int row = item / kChunksPerRow; + const int chunk = item - row * kChunksPerRow; + const int col = chunk * 32; + const uint4 packed = *reinterpret_cast(&mainloop.codes[row][col / 2]); + ContextKVBf16x8 decoded[4]; + const unsigned words[4] = {packed.x, packed.y, packed.z, packed.w}; +#pragma unroll + for (int word = 0; word < 4; ++word) { +#pragma unroll + for (int byte = 0; byte < 4; ++byte) { + const float2 values = decode_nvfp4_e2m1x2( + static_cast((words[word] >> (byte * 8)) & 0xffu)); + decoded[word].pair[byte] = + __floats2bfloat162_rn(values.x, values.y); + } + } +#pragma unroll + for (int vec = 0; vec < 4; ++vec) { + store_vec(&mainloop.code_values[row][swizzle_128(row, col + vec * 8)], + decoded[vec].raw); + } + } + }; + + stage_activation(0); + stage_weight(0); + cp_commit(); + +#pragma unroll 1 + for (int k_tile = 0; k_tile < kKTiles; ++k_tile) { + cp_wait<0>(); + __syncthreads(); + decode_e2m1_codes(); + __syncthreads(); + + // Capture the group scales (with the payload divisor folded in) before the next + // async weight stage reuses their shared plane. + float top_scales[kGroupsPerTile], bottom_scales[kGroupsPerTile]; +#pragma unroll + for (int g = 0; g < kGroupsPerTile; ++g) { + top_scales[g] = decode_nvfp4_e4m3(mainloop.scales[warp_row * 16 + gid][g]) * + inverse_divisor; + bottom_scales[g] = + decode_nvfp4_e4m3(mainloop.scales[warp_row * 16 + gid + 8][g]) * inverse_divisor; + } + __syncthreads(); + const int next = k_tile + 1; + if (next < kKTiles) { + stage_weight(next); + cp_commit(); + } + + const auto load_fragments = [&](int k_step, unsigned(&a)[4], + unsigned(&b)[kTokenMmas][2]) { + const int weight_row = warp_row * 16 + a_rowoff; + const int weight_col = k_step * 16 + a_coloff; + ldmatrix_x4( + a[0], a[1], a[2], a[3], + smem_addr( + &mainloop.code_values[weight_row][swizzle_128(weight_row, weight_col)])); +#pragma unroll + for (int token_mma = 0; token_mma < kTokenMmas; ++token_mma) { + const int activation_row = warp_col * kWarpColumns + token_mma * 8 + b_row; + const int activation_col = k_step * 16 + b_coloff; + ldmatrix_x2(b[token_mma][0], b[token_mma][1], + smem_addr(&mainloop.activations[activation_row][swizzle_128( + activation_row, activation_col)])); + } + }; + + unsigned a_fragments[4]; + unsigned b_fragments[kTokenMmas][2]; +#pragma unroll + for (int group = 0; group < kGroupsPerTile; ++group) { + float group_acc[kTokenMmas][4] = {}; + load_fragments(group, a_fragments, b_fragments); +#pragma unroll + for (int t = 0; t < kTokenMmas; ++t) + mma_bf16(group_acc[t][0], group_acc[t][1], group_acc[t][2], group_acc[t][3], + a_fragments[0], a_fragments[1], a_fragments[2], a_fragments[3], + b_fragments[t][0], b_fragments[t][1]); +#pragma unroll + for (int t = 0; t < kTokenMmas; ++t) { + accumulators[t][0] = + fmaf(group_acc[t][0], top_scales[group], accumulators[t][0]); + accumulators[t][1] = + fmaf(group_acc[t][1], top_scales[group], accumulators[t][1]); + accumulators[t][2] = + fmaf(group_acc[t][2], bottom_scales[group], accumulators[t][2]); + accumulators[t][3] = + fmaf(group_acc[t][3], bottom_scales[group], accumulators[t][3]); + } + } + + if (next < kKTiles) { + __syncthreads(); + stage_activation(next); + cp_commit(); + } + } + + __syncthreads(); + auto& scores = storage.scores; + const int local_row0 = warp_row * 16 + gid; + const int local_row1 = local_row0 + 8; +#pragma unroll + for (int token_mma = 0; token_mma < kTokenMmas; ++token_mma) { + const int column0 = warp_col * kWarpColumns + token_mma * 8 + 2 * lid; + if (column0 < Columns) { + scores[column0][local_row0] = accumulators[token_mma][0]; + scores[column0][local_row1] = accumulators[token_mma][2]; + } + if (column0 + 1 < Columns) { + scores[column0 + 1][local_row0] = accumulators[token_mma][1]; + scores[column0 + 1][local_row1] = accumulators[token_mma][3]; + } + } + } + __syncthreads(); + for (int local = warp; local < live_columns; local += kWarps) { + const int packed_column = column_begin + local; + const int column = context_column(packed_column, width, max_count); + const int request = column / width; + const int count = counts[request]; + if (count < min_count || count > max_count || column % width >= count) continue; + if constexpr (Rows == 128) { + if (!value) { + store_key_head(storage.scores[local], layer, positions, slots, column, width, + row_begin / 128); + continue; + } + } + for (int r = lane; r < Rows; r += 32) { + const int row = row_begin + r; + const float result = storage.scores[local][r]; + if (!value) { + key_scratch[row + 1024LL * (packed_column + max_count * batch * layer_index)] = + result; + } else { + const auto dst = row % 128 + 128LL * ((positions[column] & 2047) + + (long long)layer.padded_capacity * + (row / 128 + 8 * slots[request])); + layer.cache_v[dst] = __float2half_rn(__bfloat162float(__float2bfloat16_rn(result))); + } + } + } +} + +template +void launch_mma(const Tensor& x, const Tensor& positions, const Tensor& counts, const Tensor& slots, + DeviceLayers layers, ContextKVMaterializeExecutionEnvelope envelope, + const Tensor& scratch, cudaStream_t stream) { + constexpr int bytes = sizeof(MaterializeNvfp4Storage); + if constexpr (bytes > 48 * 1024) + CUDA_CHECK(cudaFuncSetAttribute( + context_kv_mma_nvfp4_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, bytes)); + context_kv_mma_nvfp4_kernel + <<>>( + static_cast(x.data), static_cast(positions.data), + static_cast(counts.data), static_cast(slots.data), layers, + static_cast(scratch.data), x.ne[1], x.ne[2], envelope.min_count, + envelope.max_count); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +void context_kv_materialize_nvfp4_launch( + const Tensor& context, const Tensor& positions, const Tensor& counts, const Tensor& state_slots, + const std::array& layers, + ContextKVMaterializeExecutionEnvelope envelope, ContextKVMaterializeRoute route, + const Tensor& key_scratch, cudaStream_t stream) { + const DeviceLayers device_layers = make_device_layers(layers); + using Route = ContextKVMaterializeRoute; + // Every column count serves from the MMA family; the small-column grouped schedules of the + // W8 table land on the 32-column MMA tile. + switch (route) { + case Route::KSplit16: + case Route::KSplit24: + case Route::Mma32: + launch_mma<64, 32, 128, 2>(context, positions, counts, state_slots, device_layers, + envelope, key_scratch, stream); + break; + case Route::Mma80: + launch_mma<64, 80, 128, 5>(context, positions, counts, state_slots, device_layers, + envelope, key_scratch, stream); + break; + case Route::Mma96: + launch_mma<64, 96, 128, 6>(context, positions, counts, state_slots, device_layers, + envelope, key_scratch, stream); + break; + case Route::Fused64: + launch_mma<128, 64, 128, 2>(context, positions, counts, state_slots, device_layers, + envelope, key_scratch, stream); + return; + case Route::Mma64: + launch_mma<64, 64, 64, 2>(context, positions, counts, state_slots, device_layers, + envelope, key_scratch, stream); + break; + } + context_kv_key_post_launch(key_scratch, positions, counts, state_slots, device_layers, + context.ne[2], context.ne[1], envelope, stream); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.cu b/src/ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.cu new file mode 100644 index 0000000000..d86d98de31 --- /dev/null +++ b/src/ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.cu @@ -0,0 +1,53 @@ +// Format-neutral finish step of linear_dynamic_grouped_conv_add: applies the two-tap dynamic +// convolution with the base kernel and finish delta to the materialized projection and +// accumulates into the residual. Shared by every projection weight format. +#include "ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.h" + +#include "core/device.h" + +#include + +namespace ninfer::ops::detail { +namespace { + +constexpr int kRows = 5120, kGroups = 320; + +__device__ __forceinline__ void finish_value(int row, int col, int width, float current, + float previous, const __nv_bfloat16* base, + const __nv_bfloat16* delta, __nv_bfloat16* residual) { + const int index = col * kRows + row, di = col * 2 * kGroups + row / 16; + float value = fmaf(__bfloat162float(base[2 * kRows + row]) + __bfloat162float(delta[di]), + current, __bfloat162float(residual[index])); + if (col % width != 0) + value = + fmaf(__bfloat162float(base[3 * kRows + row]) + __bfloat162float(delta[di + kGroups]), + previous, value); + residual[index] = __float2bfloat16_rn(value); +} + +__global__ void finish_kernel(const __nv_bfloat16* projected, const __nv_bfloat16* base, + const __nv_bfloat16* delta, __nv_bfloat16* residual, int width) { + const int row = blockIdx.x * blockDim.x + threadIdx.x, col = blockIdx.y; + if (row >= kRows) return; + const int index = col * kRows + row; + finish_value(row, col, width, __bfloat162float(projected[index]), + col % width ? __bfloat162float(projected[index - kRows]) : 0.0f, base, delta, + residual); +} + +} // namespace + +void dynamic_grouped_conv_add_finish_launch(const Tensor& projected, const Tensor& base_kernel, + const Tensor& finish_delta, Tensor& residual, + cudaStream_t stream) { + const int tokens = residual.ne[1] * residual.ne[2]; + const dim3 grid((kRows + 255) / 256, tokens); + finish_kernel<<>>( + static_cast(projected.data), + static_cast(base_kernel.data), + static_cast(finish_delta.data), + static_cast<__nv_bfloat16*>(residual.data), residual.ne[1]); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.h b/src/ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.h new file mode 100644 index 0000000000..e0545dc1f0 --- /dev/null +++ b/src/ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.h @@ -0,0 +1,16 @@ +#pragma once + +#include "core/tensor.h" + +#include + +namespace ninfer::ops::detail { + +// Finish step shared by every projection weight format of linear_dynamic_grouped_conv_add: +// projected is the materialized BF16 [5120, width*batch] projection, and the kernel folds the +// two-tap dynamic convolution into the residual in place. +void dynamic_grouped_conv_add_finish_launch(const Tensor& projected, const Tensor& base_kernel, + const Tensor& finish_delta, Tensor& residual, + cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare.cu b/src/ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare.cu new file mode 100644 index 0000000000..bbd1420aab --- /dev/null +++ b/src/ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare.cu @@ -0,0 +1,102 @@ +#include "ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare_plan.h" + +#include "core/device.h" +#include "ninfer/ops/linear.h" +#include "ninfer/ops/rmsnorm.h" +#include "ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.h" + +#include + +#include +#include + +namespace ninfer::ops::detail { +namespace { + +constexpr int kHidden = 5120, kGroups = 320, kCoefficientRows = 1280; + +// The BF16 route keeps the projected coefficients in FP32 through the conv application; this +// route reads them from the BF16 dynamic matrix the generic A16 linear materialized. The extra +// coefficient rounding belongs to the NVFP4 route's implementation profile. +template +__global__ __launch_bounds__(Capacity * 16, 4) void nvfp4_dynamic_grouped_conv_prepare_finish_kernel( + const __nv_bfloat16* base, const __nv_bfloat16* dynamic, __nv_bfloat16* prepared, + __nv_bfloat16* finish, int width, int batch_size) { + __shared__ float projected[4][Capacity]; + __shared__ float normalized[Capacity][16]; + const int tid = threadIdx.x, group = blockIdx.x, batch = blockIdx.y; + if (tid < 4 * Capacity && tid % Capacity < width) { + const int coefficient = tid / Capacity, position = tid % Capacity; + const int row = coefficient * kGroups + group; + projected[coefficient][position] = __bfloat162float( + dynamic[(static_cast(batch) * width + position) * kCoefficientRows + row]); + } + const int position = tid / 16, channel = tid % 16, hidden = group * 16 + channel; + const std::int64_t offset = static_cast(batch * width + position) * kHidden + hidden; + if (position < width) { normalized[position][channel] = __bfloat162float(prepared[offset]); } + __syncthreads(); + if (tid < 2 * Capacity && tid % Capacity < width) { + const int tap = tid / Capacity, pos = tid % Capacity; + finish[((batch * width + pos) * 2 + tap) * kGroups + group] = + __float2bfloat16_rn(projected[2 + tap][pos]); + } + if (position < width) { + float value = (__bfloat162float(base[hidden]) + projected[0][position]) * + normalized[position][channel]; + if (position > 0) + value = fmaf(__bfloat162float(base[kHidden + hidden]) + projected[1][position], + normalized[position - 1][channel], value); + prepared[offset] = __float2bfloat16_rn(value); + } +} + +void launch_finish(const Tensor& base, const Tensor& dynamic, Tensor& prepared, Tensor& finish, + cudaStream_t stream) { + const dim3 grid(kGroups, prepared.ne[2]); + if (prepared.ne[1] <= 8) { + nvfp4_dynamic_grouped_conv_prepare_finish_kernel<8> + <<>>( + static_cast(base.data), + static_cast(dynamic.data), + static_cast<__nv_bfloat16*>(prepared.data), + static_cast<__nv_bfloat16*>(finish.data), prepared.ne[1], prepared.ne[2]); + } else { + nvfp4_dynamic_grouped_conv_prepare_finish_kernel<16> + <<>>( + static_cast(base.data), + static_cast(dynamic.data), + static_cast<__nv_bfloat16*>(prepared.data), + static_cast<__nv_bfloat16*>(finish.data), prepared.ne[1], prepared.ne[2]); + } + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +void nvfp4_dynamic_grouped_conv_prepare_dispatch(const Tensor& residual, const Tensor& norm, + float eps, const Tensor& base, + const Weight& kernel_projection, Tensor& prepared, + Tensor& finish_delta, WorkspaceArena& workspace, + cudaStream_t stream) { + auto scope = workspace.scope(); + const int tokens = residual.ne[1] * residual.ne[2]; + Tensor dynamic = workspace.alloc(DType::BF16, {kCoefficientRows, tokens}); + rmsnorm(residual, norm, eps, false, prepared, stream); + linear(prepared.view({kHidden, tokens}), kernel_projection, dynamic, stream); + launch_finish(base, dynamic, prepared, finish_delta, stream); +} + +void nvfp4_linear_dynamic_grouped_conv_add_dispatch(const Tensor& x, const Weight& projection, + const Tensor& base_kernel, + const Tensor& finish_delta, Tensor& residual, + WorkspaceArena& workspace, + cudaStream_t stream) { + auto scope = workspace.scope(); + const int tokens = x.ne[1] * x.ne[2]; + Tensor projected = workspace.alloc(DType::BF16, {5120, tokens}); + linear(x.view({x.ne[0], tokens}), projection, projected, stream); + dynamic_grouped_conv_add_finish_launch(projected, base_kernel, finish_delta, residual, + stream); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare_plan.h b/src/ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare_plan.h new file mode 100644 index 0000000000..e6e7020a31 --- /dev/null +++ b/src/ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare_plan.h @@ -0,0 +1,33 @@ +#pragma once + +#include "core/arena.h" +#include "core/tensor.h" + +#include + +#include +#include + +namespace ninfer::ops::detail { + +// NVFP4 route of rmsnorm_dynamic_grouped_conv_prepare: the kernel_projection parent is the +// weight-only NVFP4 [1280,5120] matrix. The route materializes the dynamic coefficients in BF16 +// through the generic A16 linear and then applies the shared conv finish, so its transient +// footprint (1280 x W x B BF16) stays below the split-K capacity the BF16 route reports. +void nvfp4_dynamic_grouped_conv_prepare_dispatch(const Tensor& residual, const Tensor& norm, + float eps, const Tensor& base, + const Weight& kernel_projection, Tensor& prepared, + Tensor& finish_delta, WorkspaceArena& workspace, + cudaStream_t stream); + +// NVFP4 route of linear_dynamic_grouped_conv_add: the projection parent is the weight-only NVFP4 +// [5120,C] matrix (C in {4096,17408}). The route materializes the projection in BF16 through the +// generic A16 linear and then applies the shared format-neutral finish; its transient footprint +// matches the W8 route exactly. +void nvfp4_linear_dynamic_grouped_conv_add_dispatch(const Tensor& x, const Weight& projection, + const Tensor& base_kernel, + const Tensor& finish_delta, Tensor& residual, + WorkspaceArena& workspace, + cudaStream_t stream); + +} // namespace ninfer::ops::detail diff --git a/src/ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_materialized.cu b/src/ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_materialized.cu index 992ea7f899..08ef6842eb 100644 --- a/src/ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_materialized.cu +++ b/src/ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_materialized.cu @@ -1,5 +1,6 @@ #include "ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_kernels.h" #include "core/device.h" +#include "ops/dynamic_grouped_conv/dynamic_grouped_conv_add_finish.h" #include "ops/linear/w8/w8_config.h" #include "ops/linear/w8/w8_launch.h" #include "ops/linear/w8/w8_rowsplit_output.cuh" @@ -12,20 +13,7 @@ namespace ninfer::ops::detail { namespace { -constexpr int kRows = 5120, kGroups = 320; - -__device__ __forceinline__ void finish_value(int row, int col, int width, float current, - float previous, const __nv_bfloat16* base, - const __nv_bfloat16* delta, __nv_bfloat16* residual) { - const int index = col * kRows + row, di = col * 2 * kGroups + row / 16; - float value = fmaf(__bfloat162float(base[2 * kRows + row]) + __bfloat162float(delta[di]), - current, __bfloat162float(residual[index])); - if (col % width != 0) - value = - fmaf(__bfloat162float(base[3 * kRows + row]) + __bfloat162float(delta[di + kGroups]), - previous, value); - residual[index] = __float2bfloat16_rn(value); -} +constexpr int kRows = 5120; using Launch = W8Launch; @@ -70,16 +58,6 @@ constexpr auto make_launchers(std::index_sequence) { constexpr auto attention = make_launchers<4096>(std::make_index_sequence<11>{}); constexpr auto mlp = make_launchers<17408>(std::make_index_sequence<11>{}); -__global__ void finish_kernel(const __nv_bfloat16* projected, const __nv_bfloat16* base, - const __nv_bfloat16* delta, __nv_bfloat16* residual, int width) { - const int row = blockIdx.x * blockDim.x + threadIdx.x, col = blockIdx.y; - if (row >= kRows) return; - const int index = col * kRows + row; - finish_value(row, col, width, __bfloat162float(projected[index]), - col % width ? __bfloat162float(projected[index - kRows]) : 0.0f, base, delta, - residual); -} - void materialized(W8DynamicConvAddSchedule schedule, const Tensor& x, const Weight& weight, const Tensor& base, const Tensor& delta, Tensor& residual, Tensor& projected, cudaStream_t stream) { @@ -96,12 +74,7 @@ void materialized(W8DynamicConvAddSchedule schedule, const Tensor& x, const Weig launch_w8_mma_r64x32_c64_k128_a1(flat, weight, result, stream); break; } - const dim3 grid((kRows + 255) / 256, tokens); - finish_kernel<<>>(static_cast(projected.data), - static_cast(base.data), - static_cast(delta.data), - static_cast<__nv_bfloat16*>(residual.data), x.ne[1]); - CUDA_CHECK(cudaGetLastError()); + dynamic_grouped_conv_add_finish_launch(projected, base, delta, residual, stream); } } // namespace diff --git a/src/ops/linear/nvfp4/nvfp4_config.h b/src/ops/linear/nvfp4/nvfp4_config.h index 091145dc52..c9e9285ef7 100644 --- a/src/ops/linear/nvfp4/nvfp4_config.h +++ b/src/ops/linear/nvfp4/nvfp4_config.h @@ -110,6 +110,12 @@ using Nvfp4GdnInputGeometry = Nvfp4GemvGeometry<16384, 5120>; using Nvfp4MlpGateUpGeometry = Nvfp4GemvGeometry<34816, 5120>; using Nvfp4Residual6144Geometry = Nvfp4GemvGeometry<5120, 6144>; using Nvfp4Residual17408Geometry = Nvfp4GemvGeometry<5120, 17408>; +// DFlash2 drafter matrices: weight-only NVFP4 with no activation-quant divisor sites. +using Nvfp4DFlash2FeatureGeometry = Nvfp4GemvGeometry<5120, 25600>; +using Nvfp4DFlash2QkvGeometry = Nvfp4GemvGeometry<6144, 5120>; +using Nvfp4DFlash2AttnOutGeometry = Nvfp4GemvGeometry<5120, 4096>; +using Nvfp4DFlash2ConvProjGeometry = Nvfp4GemvGeometry<1280, 5120>; +using Nvfp4DFlash2SelectorGeometry = Nvfp4GemvGeometry<256, 5120>; using Nvfp4Activation5120Geometry = Nvfp4ActivationGeometry<5120>; using Nvfp4Activation6144Geometry = Nvfp4ActivationGeometry<6144>; @@ -121,6 +127,11 @@ enum class Nvfp4Problem : std::uint8_t { MlpGateUp, Residual6144, Residual17408, + DFlash2Feature, + DFlash2Qkv, + DFlash2AttnOut, + DFlash2ConvProj, + DFlash2Selector, }; inline constexpr bool is_nvfp4_linear_problem(std::int32_t output_rows, std::int32_t input_rows) { @@ -133,7 +144,17 @@ inline constexpr bool is_nvfp4_linear_problem(std::int32_t output_rows, std::int (output_rows == Nvfp4Residual6144Geometry::kOutputRows && input_rows == Nvfp4Residual6144Geometry::kInputRows) || (output_rows == Nvfp4Residual17408Geometry::kOutputRows && - input_rows == Nvfp4Residual17408Geometry::kInputRows); + input_rows == Nvfp4Residual17408Geometry::kInputRows) || + (output_rows == Nvfp4DFlash2FeatureGeometry::kOutputRows && + input_rows == Nvfp4DFlash2FeatureGeometry::kInputRows) || + (output_rows == Nvfp4DFlash2QkvGeometry::kOutputRows && + input_rows == Nvfp4DFlash2QkvGeometry::kInputRows) || + (output_rows == Nvfp4DFlash2AttnOutGeometry::kOutputRows && + input_rows == Nvfp4DFlash2AttnOutGeometry::kInputRows) || + (output_rows == Nvfp4DFlash2ConvProjGeometry::kOutputRows && + input_rows == Nvfp4DFlash2ConvProjGeometry::kInputRows) || + (output_rows == Nvfp4DFlash2SelectorGeometry::kOutputRows && + input_rows == Nvfp4DFlash2SelectorGeometry::kInputRows); } inline Nvfp4Problem resolve_nvfp4_problem(std::int32_t output_rows, std::int32_t input_rows) { @@ -157,6 +178,26 @@ inline Nvfp4Problem resolve_nvfp4_problem(std::int32_t output_rows, std::int32_t input_rows == Nvfp4Residual17408Geometry::kInputRows) { return Nvfp4Problem::Residual17408; } + if (output_rows == Nvfp4DFlash2FeatureGeometry::kOutputRows && + input_rows == Nvfp4DFlash2FeatureGeometry::kInputRows) { + return Nvfp4Problem::DFlash2Feature; + } + if (output_rows == Nvfp4DFlash2QkvGeometry::kOutputRows && + input_rows == Nvfp4DFlash2QkvGeometry::kInputRows) { + return Nvfp4Problem::DFlash2Qkv; + } + if (output_rows == Nvfp4DFlash2AttnOutGeometry::kOutputRows && + input_rows == Nvfp4DFlash2AttnOutGeometry::kInputRows) { + return Nvfp4Problem::DFlash2AttnOut; + } + if (output_rows == Nvfp4DFlash2ConvProjGeometry::kOutputRows && + input_rows == Nvfp4DFlash2ConvProjGeometry::kInputRows) { + return Nvfp4Problem::DFlash2ConvProj; + } + if (output_rows == Nvfp4DFlash2SelectorGeometry::kOutputRows && + input_rows == Nvfp4DFlash2SelectorGeometry::kInputRows) { + return Nvfp4Problem::DFlash2Selector; + } throw std::invalid_argument("unsupported NVFP4 problem"); } diff --git a/src/ops/linear/nvfp4/nvfp4_dispatch.cpp b/src/ops/linear/nvfp4/nvfp4_dispatch.cpp index 77114743d0..bf39e663f0 100644 --- a/src/ops/linear/nvfp4/nvfp4_dispatch.cpp +++ b/src/ops/linear/nvfp4/nvfp4_dispatch.cpp @@ -37,6 +37,13 @@ Nvfp4LinearRoute resolve_route(std::int32_t output_rows, std::int32_t input_rows case Nvfp4Problem::Residual6144: case Nvfp4Problem::Residual17408: return tokens >= 8 ? Nvfp4LinearRoute::W4A4 : Nvfp4LinearRoute::A16; + case Nvfp4Problem::DFlash2Feature: + case Nvfp4Problem::DFlash2Qkv: + case Nvfp4Problem::DFlash2AttnOut: + case Nvfp4Problem::DFlash2ConvProj: + case Nvfp4Problem::DFlash2Selector: + // Weight-only drafter matrices carry no activation-quant divisor sites. + throw std::invalid_argument("nvfp4 linear: DFlash2 problems admit only A16"); } throw std::logic_error("unreachable NVFP4 linear problem"); } diff --git a/src/ops/linear/nvfp4/nvfp4_gemv.cu b/src/ops/linear/nvfp4/nvfp4_gemv.cu index c74bd6dc6f..d81491850f 100644 --- a/src/ops/linear/nvfp4/nvfp4_gemv.cu +++ b/src/ops/linear/nvfp4/nvfp4_gemv.cu @@ -49,6 +49,21 @@ void launch_nvfp4_decode(const Tensor& x, const Weight& weight, Tensor& out, cud case Nvfp4Problem::Residual17408: launch_exact(x, weight, out, stream); return; + case Nvfp4Problem::DFlash2Feature: + launch_exact(x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2Qkv: + launch_exact(x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2AttnOut: + launch_exact(x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2ConvProj: + launch_exact(x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2Selector: + launch_exact(x, weight, out, stream); + return; } } diff --git a/src/ops/linear/nvfp4/nvfp4_launch.h b/src/ops/linear/nvfp4/nvfp4_launch.h index e73170c1b9..9a0d812014 100644 --- a/src/ops/linear/nvfp4/nvfp4_launch.h +++ b/src/ops/linear/nvfp4/nvfp4_launch.h @@ -8,5 +8,7 @@ namespace ninfer::ops::detail { void launch_nvfp4_decode(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t stream); void launch_nvfp4_small_t(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t stream); +void launch_nvfp4_small_t_dflash2(const Tensor& x, const Weight& weight, Tensor& out, + cudaStream_t stream); } // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_output.cuh b/src/ops/linear/nvfp4/nvfp4_output.cuh index 4c95c9e8c9..d1c7a621b8 100644 --- a/src/ops/linear/nvfp4/nvfp4_output.cuh +++ b/src/ops/linear/nvfp4/nvfp4_output.cuh @@ -30,4 +30,28 @@ struct Nvfp4ContiguousOutput { } }; +// Three-destination split for parents whose rows partition into query, key, and value segments. +// Row-group vectors never straddle a segment boundary when both row counts are multiples of eight. +template +struct Nvfp4SplitOutput3 { + __nv_bfloat16* query; + __nv_bfloat16* key; + __nv_bfloat16* value; + static_assert((QueryRows % 8) == 0 && (KvRows % 8) == 0); + + __device__ __forceinline__ void store(std::int32_t parent_row, std::int32_t token, + float result) const { + if (parent_row < QueryRows) { + query[static_cast(token) * QueryRows + parent_row] = + __float2bfloat16_rn(result); + } else if (parent_row < QueryRows + KvRows) { + key[static_cast(token) * KvRows + parent_row - QueryRows] = + __float2bfloat16_rn(result); + } else { + value[static_cast(token) * KvRows + parent_row - QueryRows - KvRows] = + __float2bfloat16_rn(result); + } + } +}; + } // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_small_t.cu b/src/ops/linear/nvfp4/nvfp4_small_t.cu index 8df412ae15..df654ae042 100644 --- a/src/ops/linear/nvfp4/nvfp4_small_t.cu +++ b/src/ops/linear/nvfp4/nvfp4_small_t.cu @@ -3,6 +3,7 @@ #include "core/device.h" #include "ops/linear/nvfp4/nvfp4_config.h" #include "ops/linear/nvfp4/nvfp4_small_t.cuh" +#include "ops/linear/nvfp4/nvfp4_small_t_launch.h" #include #include @@ -13,24 +14,6 @@ namespace { using Launch = void (*)(const Tensor&, const Weight&, Tensor&, cudaStream_t); -template -void launch_exact(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t stream) { - using Schedule = typename Nvfp4LinearSmallTProductionSchedule::Type; - constexpr int kTokenTiles = (ActiveTokens + Schedule::kTokenTile - 1) / Schedule::kTokenTile; - constexpr int kBlocks = (Geometry::kOutputRows / Schedule::kRowsPerCta) * kTokenTiles; - - const Nvfp4ContiguousOutput output{static_cast<__nv_bfloat16*>(out.data), - Geometry::kOutputRows}; - const float inverse_weight_divisor = 1.0F / weight.weight_scale_divisor; - nvfp4_small_t_kernel - <<>>( - static_cast(x.data), - static_cast(weight.qdata), - static_cast(weight.scales), inverse_weight_divisor, - Nvfp4IdentityEpilogue{}, output); - CUDA_CHECK(cudaGetLastError()); -} - template constexpr auto make_launchers(std::index_sequence) { return std::array{ @@ -64,6 +47,13 @@ void launch_nvfp4_small_t(const Tensor& x, const Weight& weight, Tensor& out, cu case Nvfp4Problem::Residual17408: launchers()[index](x, weight, out, stream); return; + case Nvfp4Problem::DFlash2Feature: + case Nvfp4Problem::DFlash2Qkv: + case Nvfp4Problem::DFlash2AttnOut: + case Nvfp4Problem::DFlash2ConvProj: + case Nvfp4Problem::DFlash2Selector: + launch_nvfp4_small_t_dflash2(x, weight, out, stream); + return; } } diff --git a/src/ops/linear/nvfp4/nvfp4_small_t_dflash2.cu b/src/ops/linear/nvfp4/nvfp4_small_t_dflash2.cu new file mode 100644 index 0000000000..dd04ad77b5 --- /dev/null +++ b/src/ops/linear/nvfp4/nvfp4_small_t_dflash2.cu @@ -0,0 +1,105 @@ +#include "ops/linear/nvfp4/nvfp4_small_t_launch.h" + +#include "ops/linear/nvfp4/nvfp4_config.h" + +#include +#include +#include +#include + +namespace ninfer::ops::detail { +namespace { + +using Launch = void (*)(const Tensor&, const Weight&, Tensor&, cudaStream_t); + +template +constexpr auto make_launchers(std::index_sequence) { + return std::array{ + &launch_exact(Offsets)>...}; +} + +template +const auto& launchers() { + static constexpr auto kLaunchers = make_launchers( + std::make_index_sequence{}); + return kLaunchers; +} + +} // namespace + +#define NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(TOKEN) \ + template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); + +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(2) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(3) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(4) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(5) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(6) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(7) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(8) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(9) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(10) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(11) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(12) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(13) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(14) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(15) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(16) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(17) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(18) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(19) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(20) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(21) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(22) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(23) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(24) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(25) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(26) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(27) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(28) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(29) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(30) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(31) +NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN(32) + +#undef NINFER_NVFP4_DFLASH2_INSTANTIATE_TOKEN + +void launch_nvfp4_small_t_dflash2(const Tensor& x, const Weight& weight, Tensor& out, + cudaStream_t stream) { + const std::size_t index = static_cast(x.ne[1] - kNvfp4FirstSmallT); + switch (resolve_nvfp4_problem(weight.n, weight.k)) { + case Nvfp4Problem::DFlash2Feature: + launchers()[index](x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2Qkv: + launchers()[index](x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2AttnOut: + launchers()[index](x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2ConvProj: + launchers()[index](x, weight, out, stream); + return; + case Nvfp4Problem::DFlash2Selector: + launchers()[index](x, weight, out, stream); + return; + case Nvfp4Problem::AttnInput: + case Nvfp4Problem::GdnInput: + case Nvfp4Problem::MlpGateUp: + case Nvfp4Problem::Residual6144: + case Nvfp4Problem::Residual17408: + break; + } + throw std::invalid_argument("nvfp4 small-T DFlash2: not a DFlash2 problem"); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_small_t_launch.h b/src/ops/linear/nvfp4/nvfp4_small_t_launch.h new file mode 100644 index 0000000000..3a9fa69ecf --- /dev/null +++ b/src/ops/linear/nvfp4/nvfp4_small_t_launch.h @@ -0,0 +1,93 @@ +#pragma once + +// ninfer::ops::detail - shared launcher template for the NVFP4 A16 small-T route (the build-speed +// TU split, mirroring the w8 small-T family). The DFlash2 drafter geometries expand in their own +// TU; the dispatcher references them through extern-template declarations and never re-expands +// the contraction. The six target-model geometries remain implicitly instantiated in the +// dispatcher TU they have always expanded in. + +#include "core/device.h" // CUDA_CHECK +#include "ops/linear/nvfp4/nvfp4_launch.h" +#include "ops/linear/nvfp4/nvfp4_config.h" +#include "ops/linear/nvfp4/nvfp4_small_t.cuh" + +#include +#include + +namespace ninfer::ops::detail { + +template +void launch_exact(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t stream) { + using Schedule = typename Nvfp4LinearSmallTProductionSchedule::Type; + constexpr int kTokenTiles = (ActiveTokens + Schedule::kTokenTile - 1) / Schedule::kTokenTile; + constexpr int kBlocks = (Geometry::kOutputRows / Schedule::kRowsPerCta) * kTokenTiles; + + const Nvfp4ContiguousOutput output{static_cast<__nv_bfloat16*>(out.data), + Geometry::kOutputRows}; + const float inverse_weight_divisor = 1.0F / weight.weight_scale_divisor; + nvfp4_small_t_kernel + <<>>( + static_cast(x.data), + static_cast(weight.qdata), + static_cast(weight.scales), inverse_weight_divisor, + Nvfp4IdentityEpilogue{}, output); + CUDA_CHECK(cudaGetLastError()); +} + +extern template void launch_exact( + const Tensor&, const Weight&, Tensor&, cudaStream_t); +extern template void launch_exact( + const Tensor&, const Weight&, Tensor&, cudaStream_t); +extern template void launch_exact( + const Tensor&, const Weight&, Tensor&, cudaStream_t); +extern template void launch_exact( + const Tensor&, const Weight&, Tensor&, cudaStream_t); +extern template void launch_exact( + const Tensor&, const Weight&, Tensor&, cudaStream_t); + +#define NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(TOKEN) \ + extern template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + extern template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + extern template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + extern template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); \ + extern template void launch_exact( \ + const Tensor&, const Weight&, Tensor&, cudaStream_t); + +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(3) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(4) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(5) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(6) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(7) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(8) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(9) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(10) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(11) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(12) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(13) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(14) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(15) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(16) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(17) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(18) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(19) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(20) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(21) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(22) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(23) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(24) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(25) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(26) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(27) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(28) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(29) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(30) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(31) +NINFER_NVFP4_DFLASH2_EXTERN_TOKEN(32) + +#undef NINFER_NVFP4_DFLASH2_EXTERN_TOKEN + +} // namespace ninfer::ops::detail diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp index 5912b5d660..fc9084cd13 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_plan.cpp @@ -16,6 +16,7 @@ namespace { enum class Nvfp4LinearSwiGluRoute { DecodeFusedA16, SmallTFusedA16, + LinearA16Post, FusedW4A4, LinearW4A4Post, TmaFusedW4A4, @@ -32,7 +33,7 @@ Nvfp4LinearSwiGluRoute resolve_route(LinearPolicy policy, std::int32_t tokens) { if (policy == LinearPolicy::A16Only) { if (tokens == 1) { return Nvfp4LinearSwiGluRoute::DecodeFusedA16; } if (tokens <= 16) { return Nvfp4LinearSwiGluRoute::SmallTFusedA16; } - throw std::invalid_argument("nvfp4 linear_swiglu A16 is registered only through T=16"); + return Nvfp4LinearSwiGluRoute::LinearA16Post; } if (tokens == 1) { return Nvfp4LinearSwiGluRoute::DecodeFusedA16; } if (tokens <= 4) { return Nvfp4LinearSwiGluRoute::SmallTFusedA16; } @@ -87,9 +88,19 @@ std::size_t nvfp4_linear_swiglu_workspace_capacity_bytes(LinearPolicy policy, } (void)resolve_route(policy, min_tokens); (void)resolve_route(policy, max_tokens); - if (policy == LinearPolicy::A16Only || max_tokens <= 4) { return 0; } + if (max_tokens <= 4) { return 0; } std::size_t maximum = 0; + if (policy == LinearPolicy::A16Only) { + // Beyond the fused A16 small-T family the route materializes the gate/up projection. + if (max_tokens > 16) { + WorkspaceLayoutBuilder layout; + layout.alloc(DType::BF16, {Nvfp4MlpGateUpGeometry::kOutputRows, max_tokens}, 256); + maximum = layout.peak_bytes(1); + } + return maximum; + } + if (min_tokens <= kFusedMaxTokens && max_tokens >= 5) { maximum = fused_workspace_bytes(std::min(max_tokens, kFusedMaxTokens)); } @@ -120,6 +131,16 @@ void nvfp4_linear_swiglu_dispatch(const Tensor& x, const Weight& weight, Tensor& case Nvfp4LinearSwiGluRoute::SmallTFusedA16: nvfp4_linear_swiglu_small_t_launch(x, weight, out, stream); return; + case Nvfp4LinearSwiGluRoute::LinearA16Post: { + auto scope = workspace.scope(); + Tensor projected = + workspace.alloc(DType::BF16, {Nvfp4MlpGateUpGeometry::kOutputRows, x.ne[1]}, 256); + linear(x, weight, projected, LinearPolicy::A16Only, workspace, stream); + constexpr std::int32_t kIntermediate = Nvfp4MlpGateUpGeometry::kOutputRows / 2; + silu_mul(projected.slice(0, 0, kIntermediate), + projected.slice(0, kIntermediate, kIntermediate), out, stream); + return; + } case Nvfp4LinearSwiGluRoute::FusedW4A4: nvfp4_linear_swiglu_w4a4_launch(x, weight, out, workspace, stream); return; diff --git a/src/ops/wrapper/attn_input_proj.cpp b/src/ops/wrapper/attn_input_proj.cpp index c7cdeb81cb..b32ae29777 100644 --- a/src/ops/wrapper/attn_input_proj.cpp +++ b/src/ops/wrapper/attn_input_proj.cpp @@ -267,6 +267,17 @@ void attn_input_proj(const Tensor& x, const Weight& query_key_value_weight, Tens require_matrix(q, kQRows, cols, "q"); require_matrix(k, kKvRows, cols, "k"); require_matrix(v, kKvRows, cols, "v"); + + if (query_key_value_weight.qtype == QType::NVFP4) { + if (hidden != 5120 || query_key_value_weight.n != kRows || + query_key_value_weight.k != hidden) { + throw std::invalid_argument("attn_input_proj: unsupported NVFP4 Q/K/V profile"); + } + (void)detail::validate_nvfp4_weight(query_key_value_weight, "attn_input_proj"); + detail::nvfp4_dflash2_attn_input(x, query_key_value_weight, q, k, v, stream); + return; + } + require_w8_rowsplit(query_key_value_weight, kRows, hidden, "query/key/value weight"); detail::w8_attn_input_dispatch(x, query_key_value_weight, q, k, v, stream); diff --git a/src/ops/wrapper/candidate_selector.cpp b/src/ops/wrapper/candidate_selector.cpp index 31abe1891d..a054c60bc7 100644 --- a/src/ops/wrapper/candidate_selector.cpp +++ b/src/ops/wrapper/candidate_selector.cpp @@ -1,6 +1,8 @@ #include "ninfer/ops/candidate_selector.h" #include "ops/candidate_selector/bf16/candidate_selector_path_plan.h" +#include "ops/candidate_selector/nvfp4/candidate_selector_path_nvfp4.h" +#include "ops/linear/nvfp4/nvfp4_format.h" #include #include @@ -44,9 +46,29 @@ bool overlaps(const Range& lhs, const Range& rhs) { return lhs_begin < rhs_begin + rhs.bytes && rhs_begin < lhs_begin + lhs.bytes; } +void require_codebook(const Weight& codebook, const char* label) { + if (codebook.qtype == QType::NVFP4) { + if (codebook.n != kCodebookRows || codebook.k != kRank) { + throw std::invalid_argument(std::string("candidate_selector_path: invalid ") + label); + } + (void)detail::validate_nvfp4_weight(codebook, "candidate_selector_path"); + return; + } + constexpr std::uint64_t kPayloadBytes = + static_cast(kCodebookRows) * kRank * sizeof(std::uint16_t); + if (codebook.qtype != QType::BF16_CTRL || codebook.layout != QuantLayout::Contiguous || + codebook.ndim != 2 || codebook.n != kCodebookRows || codebook.k != kRank || + codebook.shape[0] != kCodebookRows || codebook.shape[1] != kRank || + codebook.padded_shape[0] != kCodebookRows || codebook.padded_shape[1] != kRank || + codebook.payload_bytes < kPayloadBytes || codebook.qhigh != nullptr || + codebook.high_plane_bytes != 0 || !aligned_to(codebook.qdata, 16)) { + throw std::invalid_argument(std::string("candidate_selector_path: invalid ") + label); + } +} + void require_nonoverlap(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, const Tensor& drafts, const Tensor& proposal_q) { const std::array ranges{{ @@ -54,8 +76,9 @@ void require_nonoverlap(const Tensor& candidate_ids, const Tensor& unary_scores, {unary_scores.data, unary_scores.bytes(), "unary_scores"}, {projected_hidden.data, projected_hidden.bytes(), "projected_hidden"}, {anchors.data, anchors.bytes(), "anchors"}, - {predecessor_codebook.data, predecessor_codebook.bytes(), "predecessor_codebook"}, - {successor_codebook.data, successor_codebook.bytes(), "successor_codebook"}, + {predecessor_codebook.qdata, predecessor_codebook.payload_bytes, + "predecessor_codebook"}, + {successor_codebook.qdata, successor_codebook.payload_bytes, "successor_codebook"}, {base_positions.data, base_positions.bytes(), "base_positions"}, {configs, static_cast(candidate_ids.ne[2]) * sizeof(SamplingConfig), "configs"}, @@ -92,7 +115,7 @@ std::size_t candidate_selector_path_workspace_capacity_bytes(int min_steps, int 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) { @@ -107,10 +130,8 @@ void candidate_selector_path(const Tensor& candidate_ids, const Tensor& unary_sc require_tensor(unary_scores, DType::FP32, kCandidates, kSteps, batch_size, 1, "unary_scores"); require_tensor(projected_hidden, DType::BF16, kRank, kSteps, batch_size, 1, "projected_hidden"); require_tensor(anchors, DType::I32, batch_size, 1, 1, 1, "anchors"); - require_tensor(predecessor_codebook, DType::BF16, kRank, kCodebookRows, 1, 1, - "predecessor_codebook"); - require_tensor(successor_codebook, DType::BF16, kRank, kCodebookRows, 1, 1, - "successor_codebook"); + require_codebook(predecessor_codebook, "predecessor_codebook"); + require_codebook(successor_codebook, "successor_codebook"); require_tensor(base_positions, DType::I32, batch_size, 1, 1, 1, "base_positions"); require_tensor(drafts, DType::I32, kSteps, batch_size, 1, 1, "drafts"); require_tensor(proposal_q, DType::FP32, kCandidates, kSteps, batch_size, 1, "proposal_q"); @@ -120,9 +141,17 @@ void candidate_selector_path(const Tensor& candidate_ids, const Tensor& unary_sc require_nonoverlap(candidate_ids, unary_scores, projected_hidden, anchors, predecessor_codebook, successor_codebook, base_positions, configs, drafts, proposal_q); - detail::candidate_selector_path_dispatch( - candidate_ids, unary_scores, projected_hidden, anchors, predecessor_codebook, - successor_codebook, base_positions, configs, drafts, proposal_q, workspace, stream); + if (predecessor_codebook.qtype == QType::NVFP4) { + detail::candidate_selector_path_nvfp4_dispatch( + candidate_ids, unary_scores, projected_hidden, anchors, predecessor_codebook, + successor_codebook, base_positions, configs, drafts, proposal_q, workspace, stream); + return; + } + + detail::candidate_selector_path_dispatch(candidate_ids, unary_scores, projected_hidden, + anchors, predecessor_codebook, successor_codebook, + base_positions, configs, drafts, proposal_q, + workspace, stream); } } // namespace ninfer::ops diff --git a/src/ops/wrapper/dynamic_grouped_conv.cpp b/src/ops/wrapper/dynamic_grouped_conv.cpp index a460030643..76cd454d0d 100644 --- a/src/ops/wrapper/dynamic_grouped_conv.cpp +++ b/src/ops/wrapper/dynamic_grouped_conv.cpp @@ -1,7 +1,9 @@ #include "ninfer/ops/dynamic_grouped_conv.h" #include "ops/dynamic_grouped_conv/bf16/bf16_dynamic_grouped_conv_prepare_plan.h" +#include "ops/dynamic_grouped_conv/nvfp4/nvfp4_dynamic_grouped_conv_prepare_plan.h" #include "ops/dynamic_grouped_conv/w8/w8_dynamic_grouped_conv_add_plan.h" +#include "ops/linear/nvfp4/nvfp4_format.h" #include #include @@ -35,6 +37,13 @@ void require_tensor(const Tensor& tensor, DType dtype, std::int32_t d0, std::int } void require_kernel_projection_weight(const Weight& weight) { + if (weight.qtype == QType::NVFP4) { + if (weight.n != kCoefficientRows || weight.k != kHidden) { + throw std::invalid_argument( + "dynamic grouped conv prepare: invalid kernel_projection_weight"); + } + return; // the full NVFP4 payload contract is validated by the linear route + } constexpr std::uint64_t kPayloadBytes = static_cast(kCoefficientRows) * kHidden * sizeof(std::uint16_t); if (weight.qtype != QType::BF16_CTRL || weight.layout != QuantLayout::Contiguous || @@ -58,6 +67,12 @@ std::uint64_t required_w8_payload_bytes(std::int32_t input_rows) { } void require_finish_projection_weight(const Weight& weight, std::int32_t input_rows) { + if (weight.qtype == QType::NVFP4) { + if (weight.n != kHidden || weight.k != input_rows) { + throw std::invalid_argument("linear dynamic grouped conv add: invalid projection_weight"); + } + return; // the full NVFP4 payload contract is validated by the linear route + } const std::uint64_t payload_bytes = required_w8_payload_bytes(input_rows); if (weight.qtype != QType::W8G32_F16S || weight.layout != QuantLayout::RowSplit || weight.scale_dtype != DType::FP16 || weight.group_size != 32 || weight.group != 32 || @@ -94,10 +109,14 @@ bool overlaps(const Range& lhs, const Range& rhs) { void require_finish_nonoverlap(const Tensor& x, const Weight& projection_weight, const Tensor& base_kernel, const Tensor& finish_delta, const Tensor& residual, const WorkspaceArena& workspace) { + const bool nvfp4 = projection_weight.qtype == QType::NVFP4; const std::size_t code_bytes = - static_cast(kHidden) * static_cast(x.ne[0]); - const std::size_t scale_bytes = static_cast(kHidden) * - static_cast(x.ne[0] / 32) * sizeof(std::uint16_t); + nvfp4 ? static_cast(kHidden) * (static_cast(x.ne[0]) / 2) + : static_cast(kHidden) * static_cast(x.ne[0]); + const std::size_t scale_bytes = + nvfp4 ? static_cast(kHidden) * (static_cast(x.ne[0]) / 16) + : static_cast(kHidden) * + (static_cast(x.ne[0]) / 32) * sizeof(std::uint16_t); const std::array ranges{{ {x.data, x.bytes(), "x"}, {projection_weight.qdata, code_bytes, "projection codes"}, @@ -121,13 +140,15 @@ void require_nonoverlap(const Tensor& residual, const Tensor& norm_weight, const Tensor& base_kernel, const Weight& kernel_projection_weight, const Tensor& prepared, const Tensor& finish_delta, const WorkspaceArena& workspace) { - constexpr std::size_t kWeightBytes = - static_cast(kCoefficientRows) * kHidden * sizeof(std::uint16_t); + const std::size_t weight_bytes = + kernel_projection_weight.qtype == QType::NVFP4 + ? static_cast(kernel_projection_weight.payload_bytes) + : static_cast(kCoefficientRows) * kHidden * sizeof(std::uint16_t); const std::array ranges{{ {residual.data, residual.bytes(), "residual"}, {norm_weight.data, norm_weight.bytes(), "norm_weight"}, {base_kernel.data, base_kernel.bytes(), "base_kernel"}, - {kernel_projection_weight.qdata, kWeightBytes, "kernel_projection_weight"}, + {kernel_projection_weight.qdata, weight_bytes, "kernel_projection_weight"}, {prepared.data, prepared.bytes(), "prepared"}, {finish_delta.data, finish_delta.bytes(), "finish_delta"}, {workspace.base(), workspace.capacity(), "workspace"}, @@ -179,6 +200,14 @@ void rmsnorm_dynamic_grouped_conv_prepare(const Tensor& residual, const Tensor& require_nonoverlap(residual, norm_weight, base_kernel, kernel_projection_weight, prepared, finish_delta, workspace); + if (kernel_projection_weight.qtype == QType::NVFP4) { + (void)detail::validate_nvfp4_weight(kernel_projection_weight, kPrepareOp); + detail::nvfp4_dynamic_grouped_conv_prepare_dispatch(residual, norm_weight, eps, base_kernel, + kernel_projection_weight, prepared, + finish_delta, workspace, stream); + return; + } + detail::bf16_dynamic_grouped_conv_prepare_dispatch(residual, norm_weight, eps, base_kernel, kernel_projection_weight, prepared, finish_delta, workspace, stream); @@ -216,6 +245,14 @@ void linear_dynamic_grouped_conv_add(const Tensor& x, const Weight& projection_w require_finish_projection_weight(projection_weight, input_rows); require_finish_nonoverlap(x, projection_weight, base_kernel, finish_delta, residual, workspace); + if (projection_weight.qtype == QType::NVFP4) { + (void)detail::validate_nvfp4_weight(projection_weight, kAddOp); + detail::nvfp4_linear_dynamic_grouped_conv_add_dispatch(x, projection_weight, base_kernel, + finish_delta, residual, workspace, + stream); + return; + } + detail::w8_linear_dynamic_grouped_conv_add_dispatch(x, projection_weight, base_kernel, finish_delta, residual, workspace, stream); } 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 eb7fbb6d8b..30bef25a00 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 @@ -100,8 +100,8 @@ struct DFlash2LayerWeights { struct DFlash2CandidateSelectorWeights { Weight hidden_projection; - Tensor predecessor_codebook; - Tensor successor_codebook; + Weight predecessor_codebook; + Weight successor_codebook; }; struct DFlash2Weights { diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index a6afe01e63..4a2ae0f14f 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/src/targets/qwen3_6/impl/runtime/layouts_impl.h @@ -543,9 +543,16 @@ WorkspacePlan build_workspace_plan(const SequencePlanImpl& plan) { auto mlp = layout.scope(); prepare(); matrix(layout, DType::BF16, DFlashConfig::intermediate, tokens); - scratch(layout, ops::linear_swiglu_workspace_capacity_bytes( - QType::W8G32_F16S, 2 * DFlashConfig::intermediate, - DFlashConfig::hidden, tokens, tokens)); + // The profile admits both module weight formats; the W8 module needs no + // swiglu scratch while the NVFP4 A16 route materializes its gate/up + // projection beyond the fused small-T family, so the plan covers the max. + scratch(layout, + std::max(ops::linear_swiglu_workspace_capacity_bytes( + QType::W8G32_F16S, 2 * DFlashConfig::intermediate, + DFlashConfig::hidden, tokens, tokens), + ops::linear_swiglu_workspace_capacity_bytes( + QType::NVFP4, 2 * DFlashConfig::intermediate, + DFlashConfig::hidden, tokens, tokens))); scratch(layout, ops::linear_dynamic_grouped_conv_add_workspace_capacity_bytes( DFlashConfig::intermediate, width, width, batch, batch)); diff --git a/src/targets/qwen3_6_27b/impl/load/bindings.cpp b/src/targets/qwen3_6_27b/impl/load/bindings.cpp index f5659b86f6..d536c8e9e2 100644 --- a/src/targets/qwen3_6_27b/impl/load/bindings.cpp +++ b/src/targets/qwen3_6_27b/impl/load/bindings.cpp @@ -135,8 +135,30 @@ Weight materialized_weight(const artifact::MaterializedArtifact& materialized, } Weight row_view(const Weight& block, std::int32_t row_begin, std::int32_t row_count) { - if (row_begin < 0 || row_count <= 0 || row_begin + row_count > block.n || - block.layout != QuantLayout::RowSplit) { + if (row_begin < 0 || row_count <= 0 || row_begin + row_count > block.n) { + throw std::logic_error("invalid target row view"); + } + if (block.layout == QuantLayout::BlockScaleK16M128x4) { + // Row slicing at a 128-row boundary preserves the K16M128x4 tile arrangement exactly: + // the code plane is row-major and the sliced scale tiles stay contiguous and self-aligned. + if ((row_begin % 128) != 0 || (row_count % 128) != 0) { + throw std::logic_error("invalid NVFP4 target row view"); + } + const std::uint64_t code_row = static_cast(block.k) / 2; + const std::uint64_t scale_row = static_cast(block.k) / 16; + Weight out = block; + out.qdata = static_cast(block.qdata) + row_begin * code_row; + out.scales = static_cast(block.scales) + row_begin * scale_row; + out.payload = out.qdata; + out.payload_bytes = + block.payload_bytes - + (static_cast(row_begin) * (code_row + scale_row)); + out.n = row_count; + out.shape[0] = row_count; + out.padded_shape[0] = row_count; + return out; + } + if (block.layout != QuantLayout::RowSplit) { throw std::logic_error("invalid target row view"); } const std::uint64_t groups = static_cast(block.padded_shape[1] / block.group); @@ -461,15 +483,59 @@ void bind_qwen38_nvfp4full_text_layers(artifact::Binder& binder, BindingPlan& ou } } +// Weight-only NVFP4 DFlash2 matrix (the v2 module encoding): the payload divisor comes from +// the block-scale payload itself and the input divisor is fixed at 1.0F because the A16 drafter +// carries no activation-quant site. +WeightPlan bind_module_nvfp4_weight(artifact::Binder& binder, std::string_view name, + std::int32_t rows, std::int32_t columns, + artifact::TensorPlacement placement) { + const std::array shape = {static_cast(rows), + static_cast(columns)}; + const artifact::ObjectHandle parent = binder.require_tensor( + name, NumericFormat::NVFP4, artifact::StorageLayout::BlockScaleK16M128x4V1, shape); + if (placement == artifact::TensorPlacement::Device) { + binder.materialize_on_device(parent); + } else { + binder.validate_only(parent); + } + const artifact::BlockScaleGeometry geometry = + artifact::block_scale_geometry(NumericFormat::NVFP4, shape); + const std::uint32_t weight_bits = + read_u32_le(binder.payload(parent).data, geometry.weight_divisor_offset, name); + require_positive_finite(weight_bits, name); + return WeightPlan{.object = parent, + .format = NumericFormat::NVFP4, + .weight_scale_divisor_bits = weight_bits, + .input_scale_divisor_bits = 0x3F800000U};} + DFlash2Plan bind_dflash2(artifact::Binder& binder, artifact::TensorPlacement placement) { const auto bind_tensor = [&](std::string_view name, NumericFormat format, std::initializer_list shape) { return artifact::bind_tensor(binder, name, format, shape, placement); }; + // The module exists in either encoding under one object-name contract: the W8G32_F16S + // suffix or the weight-only NVFP4 v2 form. Matrices and codebooks dispatch per object on + // the declared format; norms and conv base kernels are BF16 in both. + const auto bind_matrix = [&](std::string_view name, std::int32_t rows, + std::int32_t columns) { + if (binder.declared_format(name) == NumericFormat::NVFP4) { + return bind_module_nvfp4_weight(binder, name, rows, columns, placement); + } + return bind_weight(binder, name, NumericFormat::W8G32_F16S, + {static_cast(rows), static_cast(columns)}, + placement); + }; + const auto bind_dense = [&](std::string_view name, std::int32_t rows, std::int32_t columns) { + if (binder.declared_format(name) == NumericFormat::NVFP4) { + return bind_module_nvfp4_weight(binder, name, rows, columns, placement); + } + return bind_weight(binder, name, NumericFormat::BF16, + {static_cast(rows), static_cast(columns)}, + placement); + }; DFlash2Plan out; - out.feature_projection = bind_weight(binder, "dflash2/feature_projection", - NumericFormat::W8G32_F16S, {5120, 25600}, placement); + out.feature_projection = bind_matrix("dflash2/feature_projection", 5120, 25600); out.context_norm = bind_tensor("dflash2/context_norm", NumericFormat::BF16, {5120}); for (std::size_t layer = 0; layer < out.layers.size(); ++layer) { DFlash2LayerPlan& target = out.layers[layer]; @@ -478,35 +544,27 @@ DFlash2Plan bind_dflash2(artifact::Binder& binder, artifact::TensorPlacement pla target.attention_conv.base_kernel = bind_tensor(prefix + "attention_conv/base_kernel", NumericFormat::BF16, {2, 2, 5120}); target.attention_conv.kernel_projection = - bind_weight(binder, prefix + "attention_conv/kernel_projection", NumericFormat::BF16, - {1280, 5120}, placement); - target.query_key_value = bind_weight(binder, prefix + "attention/query_key_value", - NumericFormat::W8G32_F16S, {6144, 5120}, placement); - target.query_norm = - bind_tensor(prefix + "attention/query_norm", NumericFormat::BF16, {128}); + bind_dense(prefix + "attention_conv/kernel_projection", 1280, 5120); + target.query_key_value = bind_matrix(prefix + "attention/query_key_value", 6144, 5120); + target.query_norm = bind_tensor(prefix + "attention/query_norm", NumericFormat::BF16, {128}); target.key_norm = bind_tensor(prefix + "attention/key_norm", NumericFormat::BF16, {128}); - target.attention_output = bind_weight(binder, prefix + "attention/output", - NumericFormat::W8G32_F16S, {5120, 4096}, placement); + target.attention_output = bind_matrix(prefix + "attention/output", 5120, 4096); target.post_attention_norm = bind_tensor(prefix + "post_attention_norm", NumericFormat::BF16, {5120}); target.mlp_conv.base_kernel = bind_tensor(prefix + "mlp_conv/base_kernel", NumericFormat::BF16, {2, 2, 5120}); target.mlp_conv.kernel_projection = - bind_weight(binder, prefix + "mlp_conv/kernel_projection", NumericFormat::BF16, - {1280, 5120}, placement); - target.gate_up = bind_weight(binder, prefix + "mlp/gate_up", NumericFormat::W8G32_F16S, - {34816, 5120}, placement); - target.down = bind_weight(binder, prefix + "mlp/down", NumericFormat::W8G32_F16S, - {5120, 17408}, placement); + bind_dense(prefix + "mlp_conv/kernel_projection", 1280, 5120); + target.gate_up = bind_matrix(prefix + "mlp/gate_up", 34816, 5120); + target.down = bind_matrix(prefix + "mlp/down", 5120, 17408); } out.final_norm = bind_tensor("dflash2/final_norm", NumericFormat::BF16, {5120}); out.candidate_selector.hidden_projection = - bind_weight(binder, "dflash2/candidate_selector/hidden_projection", NumericFormat::BF16, - {256, 5120}, placement); - out.candidate_selector.predecessor_codebook = bind_tensor( - "dflash2/candidate_selector/predecessor_codebook", NumericFormat::BF16, {248320, 256}); - out.candidate_selector.successor_codebook = bind_tensor( - "dflash2/candidate_selector/successor_codebook", NumericFormat::BF16, {248320, 256}); + bind_dense("dflash2/candidate_selector/hidden_projection", 256, 5120); + out.candidate_selector.predecessor_codebook = + bind_dense("dflash2/candidate_selector/predecessor_codebook", 248320, 256); + out.candidate_selector.successor_codebook = + bind_dense("dflash2/candidate_selector/successor_codebook", 248320, 256); return out; } @@ -609,7 +667,7 @@ ArtifactLoadPlan bind_artifact(artifact::Binder& binder, WeightsProfile weights_ binder, "vision/merger/fc2_bias", NumericFormat::BF16, {5120}, vision_placement); out.vision_merger_norm = qwen3_6::bind_vision_merger_norm(binder, vision_placement); - const bool has_dflash2 = binder.contains("dflash2/feature_projection"); + const bool has_dflash2 = binder.contains("dflash2/candidate_selector/hidden_projection"); if (features.dflash2() && !has_dflash2) { throw artifact::ArtifactError( "DFlash2 was selected but the artifact has no DFlash2 weight bundle"); @@ -618,7 +676,7 @@ ArtifactLoadPlan bind_artifact(artifact::Binder& binder, WeightsProfile weights_ const artifact::TensorPlacement placement = features.dflash2() ? artifact::TensorPlacement::Device : artifact::TensorPlacement::ValidateOnly; - out.dflash2 = bind_dflash2(binder, placement); + out.dflash2 = bind_dflash2(binder, placement); } load_plan.materialization = binder.finish(); @@ -761,12 +819,10 @@ LoadedModelData::LoadedModelData(BindingPlan plan, artifact::MaterializedArtifac artifact::materialized_tensor(backing, source.final_norm, NumericFormat::BF16, {5120}); dflash2.candidate_selector.hidden_projection = materialized_weight(backing, source.candidate_selector.hidden_projection, 256, 5120); - dflash2.candidate_selector.predecessor_codebook = - artifact::materialized_tensor(backing, source.candidate_selector.predecessor_codebook, - NumericFormat::BF16, {256, 248320}); - dflash2.candidate_selector.successor_codebook = - artifact::materialized_tensor(backing, source.candidate_selector.successor_codebook, - NumericFormat::BF16, {256, 248320}); + dflash2.candidate_selector.predecessor_codebook = materialized_weight( + backing, source.candidate_selector.predecessor_codebook, 248320, 256); + dflash2.candidate_selector.successor_codebook = materialized_weight( + backing, source.candidate_selector.successor_codebook, 248320, 256); } if (plan.features.vision) { diff --git a/src/targets/qwen3_6_27b/impl/load/bindings.h b/src/targets/qwen3_6_27b/impl/load/bindings.h index cb1750dfb5..d5384e88ec 100644 --- a/src/targets/qwen3_6_27b/impl/load/bindings.h +++ b/src/targets/qwen3_6_27b/impl/load/bindings.h @@ -125,8 +125,8 @@ struct DFlash2LayerPlan { struct DFlash2CandidateSelectorPlan { WeightPlan hidden_projection; - artifact::ObjectHandle predecessor_codebook; - artifact::ObjectHandle successor_codebook; + WeightPlan predecessor_codebook; + WeightPlan successor_codebook; }; struct DFlash2Plan { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index fc13920b65..a4cf031caa 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -269,6 +269,9 @@ ninfer_add_op_test(ninfer_linear_topk_test ninfer_add_op_test(ninfer_candidate_selector_test SOURCES ops/test_candidate_selector.cpp LIBRARIES ninfer_ops) +ninfer_add_op_test(ninfer_dflash2_nvfp4_routes_test + SOURCES ops/test_dflash2_nvfp4_routes.cpp + LIBRARIES ninfer_ops) # Stateful, fused, and exact-shape Ops remain central even when the current implementation domain # is one checkpoint/device. diff --git a/tests/ops/linear/test_nvfp4_a16.cpp b/tests/ops/linear/test_nvfp4_a16.cpp index bbb41dbde2..de5244ae3e 100644 --- a/tests/ops/linear/test_nvfp4_a16.cpp +++ b/tests/ops/linear/test_nvfp4_a16.cpp @@ -36,6 +36,22 @@ int run_nvfp4_a16() { {5120, 6144, 705U, Comparison::Sampled, true, new_problem_invocations}); failures += run_shape("NVFP4_A16", ActivationCompute::A16, make_nvfp4_weight, {5120, 17408, 707U, Comparison::Sampled, true, new_problem_invocations}); + // DFlash2 drafter problems (A16 weight-only). + constexpr std::array dflash2_invocations{ + Invocation{1, CallForm::Policy, ops::LinearPolicy::A16Only}, + Invocation{8, CallForm::Policy, ops::LinearPolicy::A16Only}, + Invocation{33, CallForm::Policy, ops::LinearPolicy::A16Only}, + }; + failures += run_shape("NVFP4_A16", ActivationCompute::A16, make_nvfp4_weight, + {5120, 25600, 711U, Comparison::Sampled, true, dflash2_invocations}); + failures += run_shape("NVFP4_A16", ActivationCompute::A16, make_nvfp4_weight, + {6144, 5120, 712U, Comparison::Sampled, true, dflash2_invocations}); + failures += run_shape("NVFP4_A16", ActivationCompute::A16, make_nvfp4_weight, + {5120, 4096, 713U, Comparison::Sampled, true, dflash2_invocations}); + failures += run_shape("NVFP4_A16", ActivationCompute::A16, make_nvfp4_weight, + {1280, 5120, 714U, Comparison::Sampled, true, dflash2_invocations}); + failures += run_shape("NVFP4_A16", ActivationCompute::A16, make_nvfp4_weight, + {256, 5120, 715U, Comparison::Sampled, true, dflash2_invocations}); return failures; } diff --git a/tests/ops/test_candidate_selector.cpp b/tests/ops/test_candidate_selector.cpp index e0680f786a..c01b0d3f93 100644 --- a/tests/ops/test_candidate_selector.cpp +++ b/tests/ops/test_candidate_selector.cpp @@ -438,8 +438,25 @@ int run(bool ties = false, bool dependent = false) { populate_codebook(predecessor_device, tokens, true); populate_codebook(successor_device, tokens, false); - Tensor predecessor(predecessor_device.p, DType::BF16, {kRank, kCodebookRows}); - Tensor successor(successor_device.p, DType::BF16, {kRank, kCodebookRows}); + const auto codebook_weight = [&](void* data) { + Weight weight{}; + weight.payload = data; + weight.payload_bytes = + static_cast(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; + }; + Weight predecessor = codebook_weight(predecessor_device.p); + Weight successor = codebook_weight(successor_device.p); int failures = 0; for (int batch_size = 1; batch_size <= kMaxBatch; ++batch_size) { diff --git a/tests/ops/test_context_kv_materialize.cpp b/tests/ops/test_context_kv_materialize.cpp index 46db694d9f..dc6286d0ab 100644 --- a/tests/ops/test_context_kv_materialize.cpp +++ b/tests/ops/test_context_kv_materialize.cpp @@ -367,6 +367,9 @@ int run_case(Fixture& fixture, const std::string& label, int width, int batch, counts_device.copy_from_host(next_counts.data(), batch * 4); slots_device.copy_from_host(next_slots.data(), batch * 4); positions_device.copy_from_host(next_positions.data(), columns * 4); + // The copies and reset run on the legacy default stream while the executable + // launches on a non-blocking stream: order them before the replay. + cuda_synchronize(); executable.launch(stream); cuda_synchronize(stream); failures += diff --git a/tests/ops/test_dflash2_nvfp4_routes.cpp b/tests/ops/test_dflash2_nvfp4_routes.cpp new file mode 100644 index 0000000000..7252a83392 --- /dev/null +++ b/tests/ops/test_dflash2_nvfp4_routes.cpp @@ -0,0 +1,539 @@ +// Oracle tests for the NVFP4 DFlash2 draft routes: the three-output attention input projection, +// the dynamic grouped conv prepare/finish pair, and the NVFP4 context_kv_materialize kernels. +// Every oracle evaluates the complete logical formula in FP64 from the represented public +// inputs, decoding the packed weights through quantized_weight::logical_weight_fp64. +#include + +#include "ninfer/ops/attn_input_proj.h" +#include "ninfer/ops/context_kv_materialize.h" +#include "ninfer/ops/dynamic_grouped_conv.h" + +#include "core/cyclic_kv_cache.h" +#include "core/decode_graph.h" +#include "ops/op_tester.h" +#include "ops/quantized_weight.h" + +#include +#include +#include +#include +#include + +using namespace ninfer; +using namespace ninfer::test; +using quantized_weight::PackedWeight; +using quantized_weight::PatternedWeightOptions; + +namespace { + +constexpr int kHidden = 5120; +constexpr double kRelative = 2.0e-2; + +DeviceBuffer payload_device(const std::vector& payload) { + DeviceBuffer buffer(payload.size()); + buffer.copy_from_host(payload.data(), payload.size()); + return buffer; +} + +float half_to_f32(std::uint16_t bits) { + const std::uint32_t sign = (bits >> 15) & 1U; + const std::uint32_t exponent = (bits >> 10) & 0x1fU; + const std::uint32_t fraction = bits & 0x3ffU; + float value = 0.0F; + if (exponent == 0) { + value = std::ldexp(static_cast(fraction), -24); + } else if (exponent != 31) { + value = std::ldexp(static_cast(fraction | 0x400U), static_cast(exponent) - 25); + } + return sign != 0U ? -value : value; +} + +float pattern_value(std::uint32_t seed, std::int32_t row, std::int32_t column, float scale) { + const std::uint32_t mixed = seed * 747796405U + static_cast(row) * 2891336453U + + static_cast(column) * 19349663U; + const int centered = static_cast((mixed >> 13U) % 509U) - 254; + return bf16_to_f32(f32_to_bf16(static_cast(centered) * scale)); +} + +bool sample_ok(double actual, double reference, double tolerance) { + return std::fabs(actual - reference) <= std::max(tolerance, std::fabs(reference) * kRelative); +} + +PackedWeight nvfp4_weight(std::int32_t rows, std::int32_t columns, std::uint32_t seed) { + PatternedWeightOptions options; + options.weight_scale_divisor = 512.0F; + options.input_scale_divisor = 1.0F; + return quantized_weight::make_patterned_weight(QType::NVFP4, rows, columns, seed, options); +} + +// --------------------------------------------------------------------------- +// Three-output attention input projection. +// --------------------------------------------------------------------------- +int verify_attn_input() { + const PackedWeight packed = nvfp4_weight(6144, kHidden, 501U); + DeviceBuffer weight_device = payload_device(packed.payload); + const Weight weight = packed.device_weight(weight_device.p); + int failures = 0; + for (const std::int32_t tokens : {1, 8, 33}) { + std::vector host(static_cast(kHidden) * tokens); + for (std::int32_t column = 0; column < tokens; ++column) { + for (std::int32_t row = 0; row < kHidden; ++row) { + host[static_cast(column) * kHidden + row] = + pattern_value(97U, row, column, 1.0F / 256.0F); + } + } + DeviceBuffer x_device = to_device_bf16(host); + DeviceBuffer q_device(static_cast(4096) * tokens * sizeof(std::uint16_t)); + DeviceBuffer k_device(static_cast(1024) * tokens * sizeof(std::uint16_t)); + DeviceBuffer v_device(static_cast(1024) * tokens * sizeof(std::uint16_t)); + Tensor x(x_device.p, DType::BF16, {kHidden, tokens}); + Tensor q(q_device.p, DType::BF16, {4096, tokens}); + Tensor k(k_device.p, DType::BF16, {1024, tokens}); + Tensor v(v_device.p, DType::BF16, {1024, tokens}); + ops::attn_input_proj(x, weight, q, k, v, nullptr); + cuda_synchronize(); + const auto q_got = from_device_bf16(q_device, static_cast(4096) * tokens); + const auto k_got = from_device_bf16(k_device, static_cast(1024) * tokens); + const auto v_got = from_device_bf16(v_device, static_cast(1024) * tokens); + const std::int32_t samples[] = {0, 123, 4095, 4096, 4600, 5120, 6143}; + for (const std::int32_t column : {0, tokens - 1}) { + for (const std::int32_t row : samples) { + double reference = 0.0; + for (std::int32_t input = 0; input < kHidden; ++input) { + reference += quantized_weight::logical_weight_fp64(packed, row, input) * + host[static_cast(column) * kHidden + input]; + } + const float actual = row < 4096 ? q_got[column * 4096 + row] + : row < 5120 ? k_got[column * 1024 + (row - 4096)] + : v_got[column * 1024 + (row - 5120)]; + if (!sample_ok(actual, reference, 2.0e-2)) { + std::cerr << "attn_input T=" << tokens << " row=" << row + << " column=" << column << ": actual=" << actual << " reference=" + << reference << '\n'; + ++failures; + } + } + } + } + return failures; +} + +// --------------------------------------------------------------------------- +// Dynamic grouped conv prepare + finish. +// --------------------------------------------------------------------------- +int verify_dynamic_conv(const PackedWeight& kernel_projection, const PackedWeight& projection, + std::int32_t width, std::int32_t batch, std::uint32_t seed) { + constexpr int kCoefficientRows = 1280; + const std::size_t columns = static_cast(width) * batch; + std::vector residual(columns * kHidden); + std::vector norm(kHidden); + std::vector base(kHidden * 4); + for (std::size_t column = 0; column < columns; ++column) { + for (std::int32_t row = 0; row < kHidden; ++row) { + residual[column * kHidden + row] = + pattern_value(seed, row, static_cast(column), 1.0F / 128.0F); + } + } + for (std::int32_t row = 0; row < kHidden; ++row) { + norm[row] = bf16_to_f32(f32_to_bf16(0.75F + static_cast(row % 17) * (1.0F / 128.0F))); + for (int block = 0; block < 4; ++block) { + base[static_cast(block) * kHidden + row] = + bf16_to_f32(f32_to_bf16(static_cast((row * 3 + block) % 11) * 0.03125F)); + } + } + + DeviceBuffer weight_device = payload_device(kernel_projection.payload); + DeviceBuffer projection_device = payload_device(projection.payload); + DeviceBuffer residual_device = to_device_bf16(residual); + DeviceBuffer norm_device = to_device_bf16(norm); + DeviceBuffer base_device = to_device_bf16(base); + GuardedDeviceBuffer prepared_device(columns * kHidden * sizeof(std::uint16_t)); + GuardedDeviceBuffer finish_device(columns * 2 * 320 * sizeof(std::uint16_t)); + const std::size_t prepare_bytes = ops::rmsnorm_dynamic_grouped_conv_prepare_workspace_capacity_bytes( + width, width, batch, batch); + GuardedDeviceBuffer scratch(std::max(prepare_bytes, 1)); + WorkspaceArena workspace(DeviceSpan{scratch.data(), scratch.bytes()}); + Tensor residual_tensor(residual_device.p, DType::BF16, {kHidden, width, batch}); + Tensor norm_tensor(norm_device.p, DType::BF16, {kHidden}); + Tensor base_tensor(base_device.p, DType::BF16, {kHidden, 2, 2}); + Tensor prepared_tensor(prepared_device.data(), DType::BF16, {kHidden, width, batch}); + Tensor finish_tensor(finish_device.data(), DType::BF16, {320, 2, width, batch}); + ops::rmsnorm_dynamic_grouped_conv_prepare(residual_tensor, norm_tensor, 1.0e-6F, base_tensor, + kernel_projection.device_weight(weight_device.p), + prepared_tensor, finish_tensor, workspace, nullptr); + cuda_synchronize(); + + // FP64 oracle: rmsnorm, projected coefficients, in-place conv application, finish delta. + std::vector normed(columns * kHidden); + for (std::size_t column = 0; column < columns; ++column) { + double sum = 0.0; + for (std::int32_t row = 0; row < kHidden; ++row) { + const double value = residual[column * kHidden + row]; + sum += value * value; + } + const double inverse = 1.0 / std::sqrt(sum / kHidden + 1.0e-6); + for (std::int32_t row = 0; row < kHidden; ++row) { + normed[column * kHidden + row] = + residual[column * kHidden + row] * inverse * norm[row]; + } + } + std::vector projected(columns * kCoefficientRows); + for (std::size_t column = 0; column < columns; ++column) { + for (std::int32_t row = 0; row < kCoefficientRows; ++row) { + double sum = 0.0; + for (std::int32_t input = 0; input < kHidden; ++input) { + sum += quantized_weight::logical_weight_fp64(kernel_projection, row, input) * + normed[column * kHidden + input]; + } + projected[column * kCoefficientRows + row] = sum; + } + } + int failures = 0; + const auto prepared_got = from_device_bf16(prepared_device.data(), columns * kHidden); + const auto finish_got = from_device_bf16(finish_device.data(), columns * 2 * 320); + for (std::size_t column = 0; column < columns; ++column) { + const std::int32_t position = static_cast(column % width); + for (std::int32_t row = 0; row < kHidden; row += 337) { + const std::int32_t group = row / 16; + double value = (base[row] + projected[column * kCoefficientRows + group]) * + normed[column * kHidden + row]; + if (position > 0) { + value += (base[kHidden + row] + + projected[column * kCoefficientRows + 320 + group]) * + normed[(column - 1) * kHidden + row]; + } + if (!sample_ok(prepared_got[column * kHidden + row], value, 4.0e-2)) { + std::cerr << "conv prepare W=" << width << " B=" << batch << " column=" << column + << " row=" << row << ": actual=" << prepared_got[column * kHidden + row] + << " reference=" << value << '\n'; + ++failures; + } + } + for (const std::int32_t group : {0, 137, 319}) { + for (int tap = 0; tap < 2; ++tap) { + const double reference = + projected[column * kCoefficientRows + (2 + tap) * 320 + group]; + const float actual = finish_got[(column * 2 + tap) * 320 + group]; + if (!sample_ok(actual, reference, 3.0e-3)) { + std::cerr << "conv finish delta W=" << width << " column=" << column + << " tap=" << tap << " group=" << group + << ": actual=" << actual << " reference=" << reference << '\n'; + ++failures; + } + } + } + } + if (failures != 0) { return failures; } + + // Finish: linear_dynamic_grouped_conv_add with the NVFP4 projection parent. + const std::int32_t input_rows = static_cast(projection.weight.k); + std::vector input(columns * input_rows); + for (std::size_t column = 0; column < columns; ++column) { + for (std::int32_t row = 0; row < input_rows; ++row) { + input[column * input_rows + row] = + pattern_value(seed + 31U, row, static_cast(column), 1.0F / 128.0F); + } + } + DeviceBuffer input_device = to_device_bf16(input); + GuardedDeviceBuffer residual_out_device(columns * kHidden * sizeof(std::uint16_t)); + residual_out_device.fill(0); + const std::size_t add_bytes = ops::linear_dynamic_grouped_conv_add_workspace_capacity_bytes( + input_rows, width, width, batch, batch); + GuardedDeviceBuffer add_scratch(std::max(add_bytes, 1)); + WorkspaceArena add_workspace(DeviceSpan{add_scratch.data(), add_scratch.bytes()}); + Tensor input_tensor(input_device.p, DType::BF16, {input_rows, width, batch}); + Tensor residual_out_tensor(residual_out_device.data(), DType::BF16, {kHidden, width, batch}); + ops::linear_dynamic_grouped_conv_add(input_tensor, + projection.device_weight(projection_device.p), + base_tensor, finish_tensor, residual_out_tensor, + add_workspace, nullptr); + cuda_synchronize(); + const auto out_got = from_device_bf16(residual_out_device.data(), columns * kHidden); + for (std::size_t column = 0; column < columns; ++column) { + const std::int32_t position = static_cast(column % width); + for (std::int32_t row = 0; row < kHidden; row += 337) { + const std::int32_t group = row / 16; + double z = 0.0; + for (std::int32_t input_row = 0; input_row < input_rows; input_row += 3) { + z += quantized_weight::logical_weight_fp64(projection, row, input_row) * + input[column * input_rows + input_row] * 3.0; + } + // The sampled-dot reconstruction above is exact only for all inputs; use the full dot. + z = 0.0; + for (std::int32_t input_row = 0; input_row < input_rows; ++input_row) { + z += quantized_weight::logical_weight_fp64(projection, row, input_row) * + input[column * input_rows + input_row]; + } + double value = (base[2 * kHidden + row] + + finish_got[(column * 2 + 0) * 320 + group]) * + z; + if (position > 0) { + double z_previous = 0.0; + for (std::int32_t input_row = 0; input_row < input_rows; ++input_row) { + z_previous += quantized_weight::logical_weight_fp64(projection, row, input_row) * + input[(column - 1) * input_rows + input_row]; + } + value += (base[3 * kHidden + row] + finish_got[(column * 2 + 1) * 320 + group]) * + z_previous; + } + if (!sample_ok(out_got[column * kHidden + row], value, 8.0e-2)) { + std::cerr << "conv add W=" << width << " B=" << batch << " column=" << column + << " row=" << row << ": actual=" << out_got[column * kHidden + row] + << " reference=" << value << '\n'; + ++failures; + } + } + } + return failures; +} + +// --------------------------------------------------------------------------- +// NVFP4 context_kv_materialize. +// --------------------------------------------------------------------------- +constexpr int kCachePadded = 2056; +constexpr int kCacheLanes = 2; +constexpr double kRopeTheta = 1.0e7; +constexpr std::uint16_t kCacheSentinel = 0x5a5aU; + +std::size_t cache_elements() { + return static_cast(128) * kCachePadded * 8 * kCacheLanes; +} + +double rope_angle(std::int32_t position, std::int32_t pair) { + return static_cast(position) * std::pow(kRopeTheta, -static_cast(pair) / 64.0); +} + +int verify_context_kv(std::int32_t width, std::int32_t batch, std::uint32_t seed, bool graph) { + constexpr int kLayers = static_cast(ops::kContextKVMaterializeLayers); + const std::size_t columns = static_cast(width) * batch; + std::array keys; + std::array values; + std::array key_device; + std::array value_device; + std::vector> norms(kLayers); + std::vector norm_device(kLayers); + for (int layer = 0; layer < kLayers; ++layer) { + keys[layer] = nvfp4_weight(1024, kHidden, 601U + 2U * layer); + values[layer] = nvfp4_weight(1024, kHidden, 602U + 2U * layer); + key_device[layer] = payload_device(keys[layer].payload); + value_device[layer] = payload_device(values[layer].payload); + norms[layer].resize(128); + for (int dim = 0; dim < 128; ++dim) { + norms[layer][dim] = + bf16_to_f32(f32_to_bf16(0.75F + static_cast((dim * 7 + layer) % 19) * + (1.0F / 128.0F))); + } + norm_device[layer] = to_device_bf16(norms[layer]); + } + std::vector context(columns * kHidden); + std::vector positions(columns); + for (std::size_t column = 0; column < columns; ++column) { + for (std::int32_t row = 0; row < kHidden; ++row) { + context[column * kHidden + row] = + pattern_value(seed, row, static_cast(column), 1.0F / 128.0F); + } + positions[column] = 17 + static_cast(column) * 13; + } + const std::vector counts(batch, width); + const std::vector slots = {0, 1}; + + DeviceBuffer context_device = to_device_bf16(context); + DeviceBuffer positions_device = to_device_i32(positions); + DeviceBuffer counts_device = to_device_i32(counts); + DeviceBuffer slots_device = to_device_i32(slots); + std::vector cache_k; + std::vector cache_v; + std::array views; + for (int layer = 0; layer < kLayers; ++layer) { + cache_k.emplace_back(cache_elements() * sizeof(std::uint16_t)); + cache_v.emplace_back(cache_elements() * sizeof(std::uint16_t)); + cache_k[layer].fill(0x5a); + cache_v[layer].fill(0x5a); + views[layer] = ops::ContextKVMaterializeLayerView{ + keys[layer].device_weight(key_device[layer].p), + values[layer].device_weight(value_device[layer].p), + Tensor(norm_device[layer].p, DType::BF16, {128}), + CyclicKVCacheLayerView{ + Tensor(cache_k[layer].data(), DType::BF16, {128, kCachePadded, 8, kCacheLanes}), + Tensor(cache_v[layer].data(), DType::FP16, {128, kCachePadded, 8, kCacheLanes}), + 2048U, static_cast(kCachePadded), 8U, 128U, + static_cast(kCacheLanes)}, + }; + } + const std::size_t capacity = ops::context_kv_materialize_workspace_capacity_bytes( + batch, width, width); + GuardedDeviceBuffer scratch(std::max(capacity, 1)); + WorkspaceArena workspace(DeviceSpan{scratch.data(), scratch.bytes()}); + Tensor context_tensor(context_device.p, DType::BF16, {kHidden, width, batch}); + Tensor positions_tensor(positions_device.p, DType::I32, {width, batch}); + Tensor counts_tensor(counts_device.p, DType::I32, {batch}); + Tensor slots_tensor(slots_device.p, DType::I32, {batch}); + if (!graph) { + ops::context_kv_materialize(context_tensor, positions_tensor, counts_tensor, slots_tensor, + views, {0U, static_cast(width)}, workspace, + nullptr); + cuda_synchronize(); + } else { + // Mirror the W8 test's replay contract: capture with the b%3 mixed-count pattern in the + // device buffers, then replay through the recorded executable with full counts, fresh + // slots, and shifted positions. + std::vector capture_counts(batch), capture_slots(batch, -1); + std::vector capture_positions(columns, -1); + for (std::int32_t b = 0; b < batch; ++b) { + capture_counts[b] = b == batch - 1 || b % 3 == 2 + ? width + : b % 3 == 1 ? std::max(1, width / 2) : 0; + if (capture_counts[b] != 0) { + capture_slots[b] = (b * 3 + 2) % 8; + for (std::int32_t i = 0; i < capture_counts[b]; ++i) { + capture_positions[static_cast(b) * width + i] = 2046 + b * 4096 + i; + } + } + } + counts_device.copy_from_host(capture_counts.data(), batch * 4); + slots_device.copy_from_host(capture_slots.data(), batch * 4); + positions_device.copy_from_host(capture_positions.data(), columns * 4); + for (auto& buffer : cache_k) { buffer.fill(0x5a); } + for (auto& buffer : cache_v) { buffer.fill(0x5a); } + cudaStream_t stream = nullptr; + cuda_check(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), + "context graph stream"); + DecodeGraphDefinition definition; + DecodeGraphExecutable executable; + cudaStream_t capture_stream = stream; + definition.capture(capture_stream, [&] { + ops::context_kv_materialize(context_tensor, positions_tensor, counts_tensor, + slots_tensor, views, + {0U, static_cast(width)}, workspace, + capture_stream); + }); + executable.instantiate(definition); + counts_device.copy_from_host(counts.data(), batch * 4); + slots_device.copy_from_host(slots.data(), batch * 4); + positions_device.copy_from_host(positions.data(), columns * 4); + for (auto& buffer : cache_k) { buffer.fill(0x5a); } + for (auto& buffer : cache_v) { buffer.fill(0x5a); } + cuda_synchronize(); + executable.launch(stream); + cuda_synchronize(stream); + cuda_check(cudaStreamDestroy(stream), "destroy context graph stream"); + } + + int failures = 0; + + for (int layer = 0; layer < kLayers; ++layer) { + const auto k_got = from_device_bf16(cache_k[layer].data(), cache_elements()); + const auto v_raw = from_device(cache_v[layer].data(), cache_elements()); + std::vector v_got(cache_elements()); + for (std::size_t index = 0; index < cache_elements(); ++index) { + v_got[index] = half_to_f32(v_raw[index]); + } + // Full FP64 oracle for sampled (column, head, dim) cells. + for (std::size_t column = 0; column < columns; column += std::max(1, columns / 3)) { + const std::int32_t position = positions[column]; + const std::size_t slot = static_cast(slots[column / width]); + for (const std::int32_t head : {0, 3, 7}) { + std::vector key_raw(128), key_normed(128); + for (int dim = 0; dim < 128; ++dim) { + const std::int32_t row = head * 128 + dim; + double sum = 0.0; + for (std::int32_t input = 0; input < kHidden; ++input) { + sum += quantized_weight::logical_weight_fp64( + keys[layer], row, input) * + context[column * kHidden + input]; + } + key_raw[dim] = sum; + } + double square = 0.0; + for (int dim = 0; dim < 128; ++dim) { square += key_raw[dim] * key_raw[dim]; } + const double inverse = 1.0 / std::sqrt(square / 128.0 + 1.0e-6); + for (int dim = 0; dim < 128; ++dim) { + key_normed[dim] = key_raw[dim] * inverse * norms[layer][dim]; + } + for (const std::int32_t pair : {0, 1, 31}) { + const double a_x = rope_angle(position, pair); + const double a_y = rope_angle(position, pair + 1); + const double kx = + key_normed[pair] * std::cos(a_x) - key_normed[pair + 64] * std::sin(a_x); + const double ky = + key_normed[pair + 64] * std::cos(a_x) + key_normed[pair] * std::sin(a_x); + const double kx1 = key_normed[pair + 1] * std::cos(a_y) - + key_normed[pair + 65] * std::sin(a_y); + const double ky1 = key_normed[pair + 65] * std::cos(a_y) + + key_normed[pair + 1] * std::sin(a_y); + const std::int64_t base = + (position & 2047) + + static_cast(kCachePadded) * (head + 8 * slot); + if (!sample_ok(k_got[base * 128 + pair], kx, 4.0e-2) || + !sample_ok(k_got[base * 128 + pair + 64], ky, 4.0e-2) || + !sample_ok(k_got[base * 128 + pair + 1], kx1, 4.0e-2) || + !sample_ok(k_got[base * 128 + pair + 65], ky1, 4.0e-2)) { + std::cerr << "context K layer=" << layer << " column=" << column + << " head=" << head << " pair=" << pair << ": actual k=" + << k_got[base * 128 + pair] << " reference=" << kx << '\n'; + ++failures; + } + } + for (const std::int32_t dim : {0, 64, 127}) { + const std::int32_t row = head * 128 + dim; + double value = 0.0; + for (std::int32_t input = 0; input < kHidden; ++input) { + value += quantized_weight::logical_weight_fp64(values[layer], row, input) * + context[column * kHidden + input]; + } + const std::int64_t base = + (position & 2047) + + static_cast(kCachePadded) * (head + 8 * slot); + if (!sample_ok(v_got[base * 128 + dim], value, 4.0e-2)) { + std::cerr << "context V layer=" << layer << " column=" << column + << " head=" << head << " dim=" << dim + << ": actual=" << v_got[base * 128 + dim] + << " reference=" << value << '\n'; + ++failures; + } + } + } + } + } + return failures; +} + +} // namespace + +int main() { + int device_count = 0; + if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count == 0) { + std::cout << "SKIP: no usable CUDA device\n"; + return 77; + } + try { + int failures = 0; + std::cerr << "section: attn_input\n"; + failures += verify_attn_input(); + std::cerr << "section: conv\n"; + const PackedWeight conv_projection_kernel = nvfp4_weight(1280, kHidden, 521U); + const PackedWeight attention_output = nvfp4_weight(kHidden, 4096, 523U); + const PackedWeight mlp_down = nvfp4_weight(kHidden, 17408, 525U); + std::cerr << "section: conv 2x1\n"; + failures += verify_dynamic_conv(conv_projection_kernel, attention_output, 2, 1, 701U); + std::cerr << "section: conv 8x2\n"; + failures += verify_dynamic_conv(conv_projection_kernel, attention_output, 8, 2, 703U); + std::cerr << "section: conv 4x3\n"; + failures += verify_dynamic_conv(conv_projection_kernel, mlp_down, 4, 3, 705U); + std::cerr << "section: context\n"; + failures += verify_context_kv(1, 1, 801U, false); + failures += verify_context_kv(16, 1, 803U, false); + failures += verify_context_kv(3, 2, 805U, false); + std::cerr << "section: context graph\n"; + failures += verify_context_kv(1, 1, 801U, true); + failures += verify_context_kv(8, 1, 807U, true); + failures += verify_context_kv(16, 1, 803U, true); + failures += verify_context_kv(3, 2, 805U, true); + failures += verify_context_kv(9, 2, 809U, true); + std::cout << (failures == 0 ? "OK" : "FAIL") << " dflash2 nvfp4 routes\n"; + return failures == 0 ? 0 : 1; + } catch (const std::exception& error) { + std::cerr << "dflash2 nvfp4 routes: " << error.what() << '\n'; + return 1; + } +}