From 2961dac61c4589cf19c891df02fcff81d1942292 Mon Sep 17 00:00:00 2001 From: MichaelDementii Date: Wed, 16 Sep 2026 20:35:14 +0200 Subject: [PATCH] perf(attention): fold the sigmoid gate into the causal reduce epilogue Every full-attention call site in TextContext issued causal_softmax_attention and then ops::sigmoid_mul over the whole output - a second pass, and a second node in the decode graph, for work the reduce epilogue is already holding in registers. The benchmark reports the node count itself, so this is measured rather than asserted: on the append entry with BF16 or INT8 storage the Op's graph is three nodes in 92 of the 128 cells swept, five in 24 and seven in 12, and handing the gate to the Op removes exactly one of them, never two. On qwen3.6-35b-a3b that is ten of the forty text layers, plus the MTP tail once per draft step - thirteen calls in an mtp3 decode round. causal_softmax_attention now takes an optional gate. Where the route reaches the shared BF16/INT8 reducer the multiply is folded into its store; every other route - FP8, NVFP4, K8V4, the prompt kernel, chunked small-T - applies the standalone elementwise kernel inside the Op. One contract either way, so no caller has to know which route it landed on. The bytes are the bytes the caller's own sigmoid_mul produced, and that is a rounding question rather than a formality. The standalone kernel reads what attention stored, so it sees the reduce result already rounded to BF16. The fused epilogue therefore rounds first, widens back, multiplies in FP32 and rounds to nearest; keeping the FP32 accumulator would be more accurate and would move tokens. One case is not a multiplication by one: a masked column stores an exact zero, and 0 * sigmoid(g) is zero for every finite gate but not for a NaN one, which the standalone multiply would propagate, so the fused path does the multiply there too. causal_softmax_attention_cached is untouched and keeps the standalone multiply; extending it is a separate decision. All three forms live in one binary, so the comparison is one process per pass: --gate standalone is what the model issued, --gate fused hands the gate to the Op, and --gate standalone again is the identical-baseline null. Six passes, the arm order rotated through six permutations, 480 cells over both entries, both geometries, all five storages, B in {1,4} and W in {1,4,8,16}, graph execution, cold cache, 20 warmups and 300 samples. A pass is dropped when the two identical binaries disagree by more than 1 %; 396 of 2880 were, and over the rest the null's p95 is 0.50 % with a worst of 1.00 %. Before that rule it is 4.05 %, which is what the drop rate is buying. Cells are classified by the bench's own graph_nodes column, not by an assumption about which routes fold. band cells median best worst outside the graph is one node shorter 76 -6.19% -15.13% +0.00% 60 / 76 the node count is unchanged, control 404 +0.00% -4.42% +5.65% 19 / 404 storage geometry cells median median saving bf16 d256-h16-kv2 16 -4.43% 1.94us bf16 d256-h24-kv4 20 -4.60% 1.99us int8 d256-h16-kv2 16 -7.50% 2.00us int8 d256-h24-kv4 24 -7.93% 2.02us Which cells fold is not a rule worth stating: 96 cells are on the append entry with BF16 or INT8 storage at W <= 8 and 76 of them fold, the twenty exceptions being every d256-h16-kv2 cell at W=8, both storages, and four d256-h24-kv4 BF16 cells. The table reports the node column rather than predicting it. The saving is a node and it does not grow with the cell: its median is 2.000 us, 48 of the 76 cells lie between 1.76 and 2.05 with the largest mode exactly at 2.048, and across an eleven-fold range of cell cost - 11.49 to 123.97 us - the correlation between saving and cost is r = -0.09. It is not constant either: fourteen cells save under 0.25 us and five save nothing measurable, and this sweep does not say why. Eager timing, where a launch is not amortised by a graph, puts 74 of those 76 cells - two lose every pass to the null rule - at a median of -14.73 % and 4.05 us, every one faster and none saving zero. Eager is the noisier instrument here: 617 of 2880 passes dropped, 21.4 %. The cells the product actually issues are narrower than the band. The artifact is d256-h16-kv2: query_norm is [256], the output projection is 2048x4096 so sixteen query heads, and the fused QKGV projection is 9216 rows, leaving two KV heads. Decode at mtp3 issues W=4 for the target block and W=1 per draft step, at B=1, BF16, append. All eight such cells fold, at a median of -2.95 % and a median saving of 1.03 us on cells costing 15.58 to 40.16 us. W=8, which DFlash7 would issue on this geometry, does not fold at all. Output is unchanged: NINFER_OP_REPORT_STATS=1 at %.17g over the whole test suite gives 11900 cases shared between the arms, all byte-identical, none differing in any field. The 335 cases the candidate adds are the gated oracle checks the test gains. Each case in causal_cache.cpp now runs a second time with a gate and is judged against a third, ungated run multiplied by a host sigmoid - an independent path from the inputs rather than a second run of the Op. Strength control: keeping the FP32 accumulator instead of rounding before the multiply, which is strictly more accurate and not what the standalone multiply produces, fails 65 exact-parity checks; restoring it passes. No end-to-end number. On this machine the identical-baseline arm reads +0.55 % median on decode over a -0.44 to +0.89 % band, against a candidate reading +1.10 %, so the stand does not resolve the change. What the Op measurement implies, as a prediction rather than a measurement: thirteen calls per round at the 1.03 us the production-shaped cells save is 13.4 us against a round the same run reports as 4891 us, or 0.27 %. ctest 120/120 on both arms. --- bench/ops/causal_softmax_attention_bench.cu | 88 ++++++++++++++----- include/ninfer/ops/softmax_attention.h | 8 +- src/models/qwen3_5/execution/text.cpp | 21 +++-- .../causal_cache/causal_softmax_attention.cpp | 54 ++++++++---- .../dense/causal_cache/launch.h | 15 ++-- .../dense/causal_cache/small_t.cu | 26 +++--- .../dense/causal_cache/small_t.cuh | 28 ++++-- tests/ops/softmax_attention/causal_cache.cpp | 62 +++++++++++++ 8 files changed, 234 insertions(+), 68 deletions(-) diff --git a/bench/ops/causal_softmax_attention_bench.cu b/bench/ops/causal_softmax_attention_bench.cu index 5e6df581a5..6c7e1e4ad8 100644 --- a/bench/ops/causal_softmax_attention_bench.cu +++ b/bench/ops/causal_softmax_attention_bench.cu @@ -10,6 +10,7 @@ #include "core/device.h" #include "core/paged_kv_cache.h" #include "core/paged_kv_storage.h" +#include "ninfer/ops/sigmoid_mul.h" #include "ninfer_bench_common.h" #include @@ -45,6 +46,12 @@ enum class Execution : std::uint8_t { Eager, Graph, Both }; enum class CacheMode : std::uint8_t { Cold, Warm, Both }; enum class CacheState : std::uint8_t { Cold, Warm }; enum class PageMapping : std::uint8_t { Identity, Fragmented }; +// off: no gate at all. standalone: the Op runs ungated and the standalone sigmoid_mul follows, +// which is exactly what a full-attention layer used to issue. fused: the gate is handed to the Op. +// standalone and fused compute the same bytes, so their medians compare directly; off is the +// control that neither form changes. The cached entry takes no gate parameter, so it applies the +// standalone multiply in both gated modes and is a second control inside the same table. +enum class GateMode : std::uint8_t { Off, Standalone, Fused }; struct Geometry { const char* name; @@ -62,6 +69,7 @@ struct Options { Execution execution = Execution::Graph; CacheMode cache = CacheMode::Cold; PageMapping mapping = PageMapping::Identity; + GateMode gate = GateMode::Off; std::vector batches{1}; std::vector tokens{1, 2, 4, 6, 8, 12, 16, 1024}; std::vector contexts{0, 128, 2048, 8192}; @@ -100,8 +108,22 @@ struct Result { bench::ColdTiming timing; std::size_t graph_nodes = 0, workspace_peak = 0; int graph_calls = 1; + // The mode this row was actually measured in, not the one asked for on the command line: the + // cached entry takes no gate parameter, so it is standalone whenever a gate is requested. + GateMode gate = GateMode::Off; }; +const char* gate_name(GateMode gate) noexcept { + return gate == GateMode::Off ? "off" : gate == GateMode::Standalone ? "standalone" : "fused"; +} + +// The cached entry takes no gate parameter, so a requested fused gate is a standalone multiply +// there. Every line that names a gate names this one. +GateMode effective_gate(GateMode requested, Entry entry) noexcept { + return entry == Entry::Cached && requested == GateMode::Fused ? GateMode::Standalone + : requested; +} + [[noreturn]] void usage(const char* message) { std::fprintf(stderr, "error: %s\n" @@ -112,7 +134,7 @@ struct Result { "[--context L,...] [--row-contexts L0,...] [--valid-columns V0,...] " "[--table-rows R0,...] " "[--execution eager|graph|both] [--cache cold|warm|both] " - "[--mapping identity|fragmented] " + "[--mapping identity|fragmented] [--gate off|standalone|fused] " "[--warmup N] [--repeat N] [--graph-calls N] [--profile] [--csv-out PATH]\n", message); std::exit(2); @@ -235,6 +257,16 @@ Options parse_options(int argc, char** argv) { options.mapping = PageMapping::Fragmented; else usage("--mapping expects identity or fragmented"); + } else if (argument == "--gate") { + const std::string_view value(next("--gate requires a value")); + if (value == "off") + options.gate = GateMode::Off; + else if (value == "standalone") + options.gate = GateMode::Standalone; + else if (value == "fused") + options.gate = GateMode::Fused; + else + usage("--gate expects off, standalone, or fused"); } else if (argument == "--graph-calls") { options.graph_calls = parse_i32(next("--graph-calls requires a value"), 1, 128, "--graph-calls"); @@ -417,8 +449,9 @@ class Case { public: Case(Geometry geometry, KvCacheStorage storage, std::int32_t tokens, std::span contexts, std::span valid_columns, - std::span table_rows, PageMapping mapping) - : storage_layout_(paged_kv_storage_layout(storage, kHeadDim)), + std::span table_rows, PageMapping mapping, + GateMode gate = GateMode::Off) + : storage_layout_(paged_kv_storage_layout(storage, kHeadDim)), gate_(gate), batch_(static_cast(contexts.size())), masked_(std::any_of(valid_columns.begin(), valid_columns.end(), [tokens](std::int32_t valid) { return valid != tokens; })), @@ -451,6 +484,8 @@ public: block_table_(static_cast(logical_pages_) * batch_ * sizeof(std::int32_t)), output_(bench::make_zeros(static_cast(kHeadDim) * geometry.query_heads * tokens * batch_ * 2)), + gate_buffer_(bench::make_bf16(static_cast(kHeadDim) * geometry.query_heads * + tokens * batch_)), workspace_bytes_(workspace_capacity(geometry, storage, tokens, batch_, visible_)), workspace_(std::max(workspace_bytes_, 1)), q_tensor_(q_.p, DType::BF16, {kHeadDim, geometry.query_heads, tokens, batch_}), @@ -460,6 +495,8 @@ public: valid_columns_tensor_(valid_columns_.p, DType::I32, {batch_}), table_rows_tensor_(table_rows_.p, DType::I32, {batch_}), output_tensor_(output_.p, DType::BF16, {kHeadDim, geometry.query_heads, tokens, batch_}), + gate_tensor_(gate_buffer_.p, DType::BF16, + {kHeadDim, geometry.query_heads, tokens, batch_}), cache_view_(make_cache_view(cache_k_, cache_v_, cache_k_scale_, cache_v_scale_, block_table_, geometry, storage, padded_, physical_pages_)), batch_cache_view_(make_batch_cache_view(cache_k_, cache_v_, cache_k_scale_, @@ -524,12 +561,16 @@ public: ops::causal_softmax_attention( q_tensor_, k_tensor_, v_tensor_, positions_tensor_, validity, table_rows_tensor_, {kHeadDim, q_tensor_.ne[1], k_tensor_.ne[1]}, kScale, batch_cache_view_, envelope_, - workspace_, output_tensor_, stream); + workspace_, output_tensor_, stream, + gate_ == GateMode::Fused ? &gate_tensor_ : nullptr); } else { ops::causal_softmax_attention_cached( q_tensor_, positions_tensor_, {kHeadDim, q_tensor_.ne[1], cache_view_.num_kv_heads}, kScale, cache_view_, envelope_, workspace_, output_tensor_, stream); } + if (gate_ == GateMode::Standalone || (gate_ == GateMode::Fused && entry == Entry::Cached)) { + ops::sigmoid_mul(gate_tensor_, output_tensor_, stream); + } } [[nodiscard]] std::size_t workspace_bytes() const noexcept { return workspace_bytes_; } @@ -542,6 +583,7 @@ public: private: PagedKVStorageLayout storage_layout_; + GateMode gate_; std::int32_t batch_; bool masked_; std::int32_t visible_; @@ -561,6 +603,7 @@ private: DeviceBuffer cache_v_scale_; DeviceBuffer block_table_; DeviceBuffer output_; + DeviceBuffer gate_buffer_; std::size_t workspace_bytes_; WorkspaceArena workspace_; Tensor q_tensor_; @@ -570,6 +613,7 @@ private: Tensor valid_columns_tensor_; Tensor table_rows_tensor_; Tensor output_tensor_; + Tensor gate_tensor_; PagedKVLayerView cache_view_; PagedKVBatchLayerView batch_cache_view_; ops::CausalAttentionExecutionEnvelope envelope_; @@ -706,15 +750,15 @@ void report(const Result& result) { const double pv_tflops = result.pv_flops / seconds / 1.0e12; std::printf( "entry=%-6s geometry=%-14s kv=%-6s mapping=%-10s execution=%-5s cache=%-4s " - "B=%d W=%d contexts=%s valid=%s rows=%s " + "gate=%-10s B=%d W=%d contexts=%s valid=%s rows=%s " "workspace=%9zu peak=%9zu nodes=%zu calls=%d median=%10.3f us min=%10.3f us p95=%10.3f us " "logical_payload=%8.1f GB/s physical_payload=%8.1f GB/s math=%7.2f TFLOP/s\n", entry_name(result.entry), result.geometry.name, storage_name(result.storage), mapping_name(result.mapping), execution_name(result.execution), cache_name(result.cache), - result.batch, result.tokens, result.row_contexts.c_str(), result.valid_columns.c_str(), - result.table_rows.c_str(), result.workspace_bytes, result.workspace_peak, - result.graph_nodes, result.graph_calls, result.timing.median_us, result.timing.min_us, - result.timing.p95_us, logical_gbps, physical_gbps, tflops); + gate_name(result.gate), result.batch, result.tokens, result.row_contexts.c_str(), + result.valid_columns.c_str(), result.table_rows.c_str(), result.workspace_bytes, + result.workspace_peak, result.graph_nodes, result.graph_calls, result.timing.median_us, + result.timing.min_us, result.timing.p95_us, logical_gbps, physical_gbps, tflops); std::printf(" vectors K=%.0f V=%.0f bytes cache logical=%.0f physical=%.0f bytes " "qk_flops=%.0f pv_flops=%.0f qk_full_op=%7.2f TFLOP/s " "pv_full_op=%7.2f TFLOP/s unique_kv=%.0f bytes " @@ -731,7 +775,7 @@ void write_csv(const Options& options, const std::vector& results) { std::ofstream output(path); if (!output) { throw std::runtime_error("failed to open CSV output"); } output - << "entry,geometry,kv_dtype,mapping,execution,cache,B,W,row_contexts,valid_columns," + << "entry,geometry,kv_dtype,mapping,execution,cache,gate,B,W,row_contexts,valid_columns," "table_rows,workspace_bytes,logical_bytes,physical_bytes,logical_cache_bytes," "key_vector_bytes,value_vector_bytes,physical_cache_bytes,qk_flops,pv_flops," "qk_full_op_tflops,pv_full_op_tflops," @@ -741,12 +785,13 @@ void write_csv(const Options& options, const std::vector& results) { output << entry_name(result.entry) << ',' << result.geometry.name << ',' << storage_name(result.storage) << ',' << mapping_name(result.mapping) << ',' << execution_name(result.execution) << ',' << cache_name(result.cache) << ',' - << result.batch << ',' << result.tokens << ',' << result.row_contexts << ',' - << result.valid_columns << ',' << result.table_rows << ',' << result.workspace_bytes - << ',' << result.logical_bytes << ',' << result.physical_bytes << ',' - << result.logical_cache_bytes << ',' << result.key_vector_bytes << ',' - << result.value_vector_bytes << ',' << result.physical_cache_bytes << ',' - << result.qk_flops << ',' << result.pv_flops << ','; + << gate_name(result.gate) << ',' << result.batch << ',' << result.tokens << ',' + << result.row_contexts << ',' << result.valid_columns << ',' << result.table_rows + << ',' << result.workspace_bytes << ',' << result.logical_bytes << ',' + << result.physical_bytes << ',' << result.logical_cache_bytes << ',' + << result.key_vector_bytes << ',' << result.value_vector_bytes << ',' + << result.physical_cache_bytes << ',' << result.qk_flops << ',' << result.pv_flops + << ','; const double seconds = result.timing.median_us * 1.0e-6; output << result.qk_flops / seconds / 1.0e12 << ',' << result.pv_flops / seconds / 1.0e12 << ',' << result.unique_kv_bytes << ',' << result.unique_kv_bytes / seconds / 1.0e9; @@ -781,9 +826,10 @@ void profile(Case& data, Entry entry, const Geometry& geometry, KvCacheStorage s } std::printf( "PROFILE entry=%s geometry=%s kv=%s mapping=%s dispatch=public execution=%s cache=%s " - "B=%d W=%d contexts=%.*s valid=%.*s rows=%.*s graph_calls=%d\n", + "gate=%s B=%d W=%d contexts=%.*s valid=%.*s rows=%.*s graph_calls=%d\n", entry_name(entry), geometry.name, storage_name(storage), mapping_name(options.mapping), - execution_name(execution), cache_name(cache), batch, width, + execution_name(execution), cache_name(cache), + gate_name(effective_gate(options.gate, entry)), batch, width, static_cast(contexts.size()), contexts.data(), static_cast(valid_columns.size()), valid_columns.data(), static_cast(table_rows.size()), table_rows.data(), options.graph_calls); @@ -874,7 +920,7 @@ int main(int argc, char** argv) { options.row_contexts.empty() ? options.contexts.front() : 0; const RowProfile rows = make_row_profile(options, batch, width, context); Case data(geometry, storage, width, rows.contexts, rows.valid_columns, rows.table_rows, - options.mapping); + options.mapping, options.gate); const std::string context_name = profile_name(rows.contexts); const std::string valid_name = profile_name(rows.valid_columns); const std::string table_name = profile_name(rows.table_rows); @@ -884,6 +930,7 @@ int main(int argc, char** argv) { return 0; } + std::printf("gate=%s\n", gate_name(options.gate)); std::vector results; const std::vector context_profiles = options.row_contexts.empty() ? options.contexts : std::vector{0}; @@ -895,7 +942,7 @@ int main(int argc, char** argv) { const RowProfile rows = make_row_profile(options, batch, tokens, context); Case data(geometry, storage, tokens, rows.contexts, rows.valid_columns, - rows.table_rows, options.mapping); + rows.table_rows, options.mapping, options.gate); for (const Entry entry : {Entry::Append, Entry::Cached}) { if ((options.entry == Entry::Append && entry != Entry::Append) || (options.entry == Entry::Cached && entry != Entry::Cached) || @@ -965,6 +1012,7 @@ int main(int argc, char** argv) { result.workspace_peak = data.workspace_peak(); result.graph_calls = execution == Execution::Graph ? options.graph_calls : 1; + result.gate = effective_gate(options.gate, entry); result.timing.median_us /= result.graph_calls; result.timing.min_us /= result.graph_calls; result.timing.p95_us /= result.graph_calls; diff --git a/include/ninfer/ops/softmax_attention.h b/include/ninfer/ops/softmax_attention.h index 77e9931b6c..114038e221 100644 --- a/include/ninfer/ops/softmax_attention.h +++ b/include/ninfer/ops/softmax_attention.h @@ -132,13 +132,19 @@ void packed_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v, * Inputs, output, every cache plane/table, and live workspace suballocations are pairwise * non-overlapping. The Op overwrites every addressed cache row but owns no cache allocation, * frontier, request identity, or commit authority. + * + * An optional gate asks the Op to finish with out *= sigmoid(gate). When present it is a contiguous + * BF16 tensor shaped exactly like out and disjoint from every other operand. Where the route allows + * it the multiply is folded into the reduce epilogue; every other route applies the standalone + * elementwise kernel inside the Op. Either way the result is bit-identical to calling sigmoid_mul + * on the ungated output, so no caller has to know which route it landed on. */ void causal_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& valid_columns, const Tensor& kv_table_rows, AttentionHeadGeometry geometry, float scale, PagedKVBatchLayerView cache, CausalAttentionExecutionEnvelope envelope, WorkspaceArena& workspace, - Tensor& out, cudaStream_t stream); + Tensor& out, cudaStream_t stream, const Tensor* gate = nullptr); /** * Read-only single-sequence causal attention over an already populated cache. diff --git a/src/models/qwen3_5/execution/text.cpp b/src/models/qwen3_5/execution/text.cpp index d9037515b3..b6e63c3622 100644 --- a/src/models/qwen3_5/execution/text.cpp +++ b/src/models/qwen3_5/execution/text.cpp @@ -376,6 +376,9 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po Tensor a_batch = a.view({dimension(config_.attention->head_dim), dimension(config_.attention->num_attention_heads), width, active_sequence_batch_}); + Tensor gate_batch = gate.view({dimension(config_.attention->head_dim), + dimension(config_.attention->num_attention_heads), width, + active_sequence_batch_}); Tensor position_batch = positions.view({width, active_sequence_batch_}); ops::causal_softmax_attention( q_batch, k_batch, v_batch, position_batch, *active_valid_columns_, @@ -384,7 +387,7 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po dimension(config_.attention->num_attention_heads), dimension(config_.attention->num_key_value_heads)}, static_cast(1.0 / std::sqrt(static_cast(config_.attention->head_dim))), - batch_mtp_kv_->batch_layer_view(0), envelope, work_, a_batch, s); + batch_mtp_kv_->batch_layer_view(0), envelope, work_, a_batch, s, &gate_batch); } else { ops::causal_softmax_attention( qn, kn, v, positions, Tensor{}, io_.backend_kv_table_row, @@ -392,9 +395,8 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po dimension(config_.attention->num_attention_heads), dimension(config_.attention->num_key_value_heads)}, static_cast(1.0 / std::sqrt(static_cast(config_.attention->head_dim))), - batch_mtp_kv_->batch_layer_view(0), envelope, work_, a, s); + batch_mtp_kv_->batch_layer_view(0), envelope, work_, a, s, &gate); } - ops::sigmoid_mul(gate, a, s); const auto post = workspace::mtp_post_attention(work_, config_, T); Tensor o = post.output; @@ -899,8 +901,14 @@ void TextContext::attn_mix(const BlockParameters& w, Tensor& x, int fidx, Phase Tensor a_batch = a.view({dimension(config_.attention->head_dim), dimension(config_.attention->num_attention_heads), width, active_sequence_batch_}); + Tensor gate_batch = gate.view({dimension(config_.attention->head_dim), + dimension(config_.attention->num_attention_heads), width, + active_sequence_batch_}); Tensor position_batch = cache_positions.view({width, active_sequence_batch_}); const Tensor valid = active_valid_columns_ != nullptr ? *active_valid_columns_ : Tensor{}; + // The gate rides along with the attention call: where the route can, the reduce epilogue + // applies it at the store, and every other route applies it inside the Op. One contract + // either way, and the same bytes the standalone multiply produced. ops::causal_softmax_attention( q_batch, k_batch, v_batch, position_batch, valid, kv_table_rows, {dimension(config_.attention->head_dim), @@ -908,7 +916,7 @@ void TextContext::attn_mix(const BlockParameters& w, Tensor& x, int fidx, Phase dimension(config_.attention->num_key_value_heads)}, static_cast(1.0 / std::sqrt(static_cast(config_.attention->head_dim))), batch_text_kv_->batch_layer_view(fidx), *active_causal_attention_envelope_, work_, - a_batch, s); + a_batch, s, &gate_batch); } else { ops::causal_softmax_attention( qn, kn, v, cache_positions, Tensor{}, kv_table_rows, @@ -916,10 +924,9 @@ void TextContext::attn_mix(const BlockParameters& w, Tensor& x, int fidx, Phase dimension(config_.attention->num_attention_heads), dimension(config_.attention->num_key_value_heads)}, static_cast(1.0 / std::sqrt(static_cast(config_.attention->head_dim))), - batch_text_kv_->batch_layer_view(fidx), *active_causal_attention_envelope_, work_, a, - s); + batch_text_kv_->batch_layer_view(fidx), *active_causal_attention_envelope_, work_, a, s, + &gate); } - ops::sigmoid_mul(gate, a, s); ops::linear_add(a.view({dimension(config_.attention->query_width()), T}), p.output.weight, x, p.output.policy, work_, s); diff --git a/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp b/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp index fd69b1d07c..365fd9f30e 100644 --- a/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp +++ b/src/ops/softmax_attention/dense/causal_cache/causal_softmax_attention.cpp @@ -1,6 +1,8 @@ // ninfer::ops - causal cached Softmax Attention validation and finite route dispatch. #include "ninfer/ops/softmax_attention.h" +#include "ninfer/ops/sigmoid_mul.h" + #include "core/layout.h" #include "core/paged_kv_storage.h" #include "ops/softmax_attention/dense/causal_cache/launch.h" @@ -286,11 +288,10 @@ void for_each_small_t_chunk(const Tensor& q, const Tensor& positions, WorkspaceA KvCacheStorage cache_storage, CausalAttentionExecutionEnvelope envelope, Tensor& out, Launch&& launch) { for (std::int32_t begin = 0; begin < q.ne[2]; - begin += - causal_attention_chunk_tokens(q.ne[1], q.ne[2], 1, cache_storage, envelope)) { - const std::int32_t count = std::min( - causal_attention_chunk_tokens(q.ne[1], q.ne[2], 1, cache_storage, envelope), - q.ne[2] - begin); + begin += causal_attention_chunk_tokens(q.ne[1], q.ne[2], 1, cache_storage, envelope)) { + const std::int32_t count = + std::min(causal_attention_chunk_tokens(q.ne[1], q.ne[2], 1, cache_storage, envelope), + q.ne[2] - begin); auto chunk_scope = workspace.scope(); const std::int32_t splits = detail::causal_attention_split_capacity(q.ne[1], count, cache_storage, envelope); @@ -308,11 +309,11 @@ void launch_chunked_small_t(const Tensor& q, const Tensor& k, const Tensor& v, CausalAttentionExecutionEnvelope envelope, WorkspaceArena& workspace, Tensor& out, cudaStream_t stream) { for (std::int32_t begin = 0; begin < q.ne[2]; - begin += causal_attention_chunk_tokens(q.ne[1], q.ne[2], q.ne[3], cache.storage, - envelope)) { - const std::int32_t count = std::min(causal_attention_chunk_tokens( - q.ne[1], q.ne[2], q.ne[3], cache.storage, envelope), - q.ne[2] - begin); + begin += + causal_attention_chunk_tokens(q.ne[1], q.ne[2], q.ne[3], cache.storage, envelope)) { + const std::int32_t count = std::min( + causal_attention_chunk_tokens(q.ne[1], q.ne[2], q.ne[3], cache.storage, envelope), + q.ne[2] - begin); auto chunk_scope = workspace.scope(); const std::int32_t splits = detail::causal_attention_split_capacity( q.ne[1], count, cache.storage, envelope, q.ne[3]); @@ -425,12 +426,12 @@ std::size_t causal_softmax_attention_workspace_capacity_bytes( if (route == detail::CausalAttentionRoute::SmallT) { return chunk_capacity(width); } std::size_t maximum = 0; for (std::int32_t begin = 0; begin < width; - begin += causal_attention_chunk_tokens(q_heads, width, batch_size, - cache_storage, envelope)) { + begin += + causal_attention_chunk_tokens(q_heads, width, batch_size, cache_storage, envelope)) { maximum = std::max( maximum, - chunk_capacity(std::min(causal_attention_chunk_tokens( - q_heads, width, batch_size, cache_storage, envelope), + chunk_capacity(std::min(causal_attention_chunk_tokens(q_heads, width, batch_size, + cache_storage, envelope), width - begin))); } return maximum; @@ -451,7 +452,7 @@ void causal_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& kv_table_rows, AttentionHeadGeometry geometry, float scale, PagedKVBatchLayerView cache, CausalAttentionExecutionEnvelope envelope, WorkspaceArena& workspace, - Tensor& out, cudaStream_t stream) { + Tensor& out, cudaStream_t stream, const Tensor* gate) { constexpr const char* op = "causal_softmax_attention"; validate_batched_attention_tensors(q, positions, valid_columns, kv_table_rows, out, cache, geometry, envelope, scale, op); @@ -465,6 +466,14 @@ void causal_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v, require_shape(v, kHeadDim, kv_heads, width, batch, op, "v"); require_contiguous_nonnull(k, op, "k"); require_contiguous_nonnull(v, op, "v"); + if (gate != nullptr) { + require_shape(*gate, kHeadDim, static_cast(q.ne[1]), width, batch, op, + "gate"); + require_contiguous_nonnull(*gate, op, "gate"); + if (gate->dtype != DType::BF16) { + throw std::invalid_argument("causal_softmax_attention: gate must be BF16"); + } + } auto scope = workspace.scope(); const detail::CausalAttentionRoute route = @@ -472,6 +481,7 @@ void causal_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v, if (route == detail::CausalAttentionRoute::ChunkedSmallT) { launch_chunked_small_t(q, k, v, positions, valid_columns, kv_table_rows, scale, cache, envelope, workspace, out, stream); + if (gate != nullptr) { sigmoid_mul(*gate, out, stream); } return; } if (route == detail::CausalAttentionRoute::SmallT) { @@ -479,13 +489,21 @@ void causal_softmax_attention(const Tensor& q, const Tensor& k, const Tensor& v, detail::causal_attention_split_capacity(q.ne[1], width, cache.storage, envelope, batch); SmallTWorkspace partial = allocate_small_t_workspace(workspace, q.ne[1], width, splits, batch); - detail::causal_attention_small_t_launch(q, k, v, positions, valid_columns, kv_table_rows, - scale, cache, envelope, 0, width, partial.acc, - partial.m, partial.l, out, stream); + // Only the shared BF16/INT8 reducer carries a gate. The FP8, NVFP4 and K8V4 storages reach + // their own reduce kernels, so they take the standalone multiply like the routes that + // cannot fold it at all; every caller still sees one contract. + const bool fusable = cache.storage == KvCacheStorage::BFloat16 || + cache.storage == KvCacheStorage::Int8Group64; + detail::causal_attention_small_t_launch( + q, k, v, positions, valid_columns, kv_table_rows, scale, cache, envelope, 0, width, + partial.acc, partial.m, partial.l, out, stream, + (gate != nullptr && fusable) ? gate->data : nullptr); + if (gate != nullptr && !fusable) { sigmoid_mul(*gate, out, stream); } return; } detail::causal_attention_prompt_launch(q, k, v, positions, valid_columns, kv_table_rows, scale, cache, out, stream); + if (gate != nullptr) { sigmoid_mul(*gate, out, stream); } } void causal_softmax_attention_cached(const Tensor& q, const Tensor& positions, diff --git a/src/ops/softmax_attention/dense/causal_cache/launch.h b/src/ops/softmax_attention/dense/causal_cache/launch.h index 71670c1798..83126e52d9 100644 --- a/src/ops/softmax_attention/dense/causal_cache/launch.h +++ b/src/ops/softmax_attention/dense/causal_cache/launch.h @@ -34,11 +34,16 @@ CausalAttentionRoute causal_attention_resolve_route(std::int32_t q_heads, std::i const char* causal_attention_route_name(CausalAttentionRoute route); -void causal_attention_small_t_launch( - const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& positions, - const Tensor& valid_columns, const Tensor& table_rows, float scale, PagedKVBatchLayerView cache, - CausalAttentionExecutionEnvelope envelope, std::int32_t column_begin, std::int32_t width, - Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, cudaStream_t stream); +// A non-null gate is applied by the reduce epilogue at the store. Only the shared BF16/INT8 +// reducer accepts one; the caller keeps the standalone multiply for every other storage. +void causal_attention_small_t_launch(const Tensor& q, const Tensor& k, const Tensor& v, + const Tensor& positions, const Tensor& valid_columns, + const Tensor& table_rows, float scale, + PagedKVBatchLayerView cache, + CausalAttentionExecutionEnvelope envelope, + std::int32_t column_begin, std::int32_t width, + Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, + Tensor& out, cudaStream_t stream, const void* gate = nullptr); void causal_attention_cached_small_t_launch(const Tensor& q, const Tensor& positions, float scale, const PagedKVLayerView& cache, diff --git a/src/ops/softmax_attention/dense/causal_cache/small_t.cu b/src/ops/softmax_attention/dense/causal_cache/small_t.cu index 711eea9a4c..8a6616d198 100644 --- a/src/ops/softmax_attention/dense/causal_cache/small_t.cu +++ b/src/ops/softmax_attention/dense/causal_cache/small_t.cu @@ -258,7 +258,8 @@ void causal_attention_small_t_launch_for(const Tensor& q, CacheInput input, cons const CausalSmallTInvocation& invocation, CausalAttentionExecutionEnvelope envelope, Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, - Tensor& out, cudaStream_t stream) { + Tensor& out, cudaStream_t stream, + const void* gate = nullptr) { const auto logical_capacity = static_cast(envelope.max_visible_keys); const auto implementation_window = static_cast(envelope.max_visible_keys); const auto splits = causal_attention_split_capacity( @@ -342,7 +343,7 @@ void causal_attention_small_t_launch_for(const Tensor& q, CacheInput input, cons ? static_cast(invocation.valid_columns->data) : nullptr, invocation.width, invocation.full_width, invocation.column_begin, invocation.batch_size, - splits, static_cast<__nv_bfloat16*>(out.data)); + splits, static_cast<__nv_bfloat16*>(out.data), static_cast(gate)); }; const auto launch_profile = [&]() { if (invocation.column_begin == 0) @@ -368,11 +369,14 @@ void causal_attention_small_t_launch_for(const Tensor& q, CacheInput input, cons CUDA_CHECK(cudaGetLastError()); } -void causal_attention_small_t_launch( - const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& pos, - const Tensor& valid_columns, const Tensor& table_rows, float scale, PagedKVBatchLayerView cache, - CausalAttentionExecutionEnvelope envelope, std::int32_t column_begin, std::int32_t width, - Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, cudaStream_t stream) { +void causal_attention_small_t_launch(const Tensor& q, const Tensor& k, const Tensor& v, + const Tensor& pos, const Tensor& valid_columns, + const Tensor& table_rows, float scale, + PagedKVBatchLayerView cache, + CausalAttentionExecutionEnvelope envelope, + std::int32_t column_begin, std::int32_t width, + Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, + Tensor& out, cudaStream_t stream, const void* gate) { if (cache.storage == KvCacheStorage::Fp8KeyNvfp4Value) { causal_attention_small_t_k8v4_launch(q, k, v, pos, valid_columns, table_rows, scale, cache, envelope, column_begin, width, partial_acc, partial_m, @@ -402,14 +406,14 @@ void causal_attention_small_t_launch( .batch_size = q.ne[3], }; if (q.ne[1] == CausalD256H24Kv4::QHeads) { - causal_attention_small_t_launch_for(q, input, pos, scale, cache, - invocation, envelope, partial_acc, - partial_m, partial_l, out, stream); + causal_attention_small_t_launch_for( + q, input, pos, scale, cache, invocation, envelope, partial_acc, partial_m, partial_l, + out, stream, gate); return; } causal_attention_small_t_launch_for(q, input, pos, scale, cache, invocation, envelope, partial_acc, partial_m, - partial_l, out, stream); + partial_l, out, stream, gate); } void causal_attention_cached_small_t_launch(const Tensor& q, const Tensor& pos, float scale, diff --git a/src/ops/softmax_attention/dense/causal_cache/small_t.cuh b/src/ops/softmax_attention/dense/causal_cache/small_t.cuh index 4cb8127178..dde2021ccf 100644 --- a/src/ops/softmax_attention/dense/causal_cache/small_t.cuh +++ b/src/ops/softmax_attention/dense/causal_cache/small_t.cuh @@ -198,7 +198,7 @@ __launch_bounds__(256) __global__ void causal_attention_small_t_reduce_output_ke const float* partial_acc, const float* partial_m, const float* partial_l, const std::int32_t* positions, const std::int32_t* valid_columns, std::int32_t tokens, std::int32_t full_width, std::int32_t column_begin, std::int32_t batch_size, - std::int32_t split_count, __nv_bfloat16* out) { + std::int32_t split_count, __nv_bfloat16* out, const __nv_bfloat16* __restrict__ gate) { static_assert(DChunk > 0 && DChunk <= kCausalHeadDim); const int q_head = static_cast(blockIdx.x); @@ -225,9 +225,15 @@ __launch_bounds__(256) __global__ void causal_attention_small_t_reduce_output_ke if constexpr (Masked) { const int absolute_column = token + (Offset ? column_begin : 0); if (absolute_column >= valid_columns[batch]) { - if (tid < DChunk && d_start + tid < kCausalHeadDim) - out[causal_q_index(q_head, d_start + tid, output_column)] = - __float2bfloat16(0.0f); + if (tid < DChunk && d_start + tid < kCausalHeadDim) { + const auto index = causal_q_index(q_head, d_start + tid, output_column); + // 0 * sigmoid(g) is 0 for every finite g, but not for a NaN gate, and the + // standalone multiply would propagate that NaN. Do the multiply either way. + out[index] = + gate == nullptr + ? __float2bfloat16(0.0f) + : __float2bfloat16_rn(0.0f * sigmoid(__bfloat162float(gate[index]))); + } return; } } @@ -261,8 +267,18 @@ __launch_bounds__(256) __global__ void causal_attention_small_t_reduce_output_ke weights[split]; } - const float value = (head_l > 0.0f) ? numerator / head_l : 0.0f; - out[causal_q_index(q_head, d, output_column)] = __float2bfloat16(value); + const float value = (head_l > 0.0f) ? numerator / head_l : 0.0f; + const auto out_index = causal_q_index(q_head, d, output_column); + if (gate == nullptr) { + out[out_index] = __float2bfloat16(value); + return; + } + // Bit-exact with the standalone multiply: that kernel reads what attention stored, so it sees + // the reduce result already rounded to BF16. Round here first, widen back, multiply in FP32, + // round to nearest. Keeping the FP32 value would be more accurate and would move tokens. + const __nv_bfloat16 reduced = __float2bfloat16(value); + out[out_index] = + __float2bfloat16_rn(__bfloat162float(reduced) * sigmoid(__bfloat162float(gate[out_index]))); } } // namespace ninfer::ops diff --git a/tests/ops/softmax_attention/causal_cache.cpp b/tests/ops/softmax_attention/causal_cache.cpp index 085a546fc8..5cb35d121e 100644 --- a/tests/ops/softmax_attention/causal_cache.cpp +++ b/tests/ops/softmax_attention/causal_cache.cpp @@ -1,6 +1,7 @@ #include "core/arena.h" #include "core/paged_kv_cache.h" #include "ninfer/ops/kv_cache_append.h" +#include "ninfer/ops/sigmoid_mul.h" #include "ninfer/ops/softmax_attention.h" #include "ops/op_tester.h" #include "ops/softmax_attention/oracle.h" @@ -1659,6 +1660,19 @@ ReductionCriterion attention_criterion(KvCacheStorage storage) { throw std::logic_error("unregistered causal-attention test storage"); } +// The gated route rounds to BF16 twice -- once on the attention result, once on the product with +// the gate -- where the ungated route rounds once. The second rounding moves an element by at +// most 2^-9 of itself, and the gate cannot enlarge an element, so 2^-9 is what both the +// relative-L2 bound and the bound taken relative to the largest reference have to grow by. The +// absolute floor is unchanged: it is there for elements near zero, which the extra rounding +// cannot move by more than it already covers. +ReductionCriterion gated_attention_criterion(KvCacheStorage storage) { + ReductionCriterion criterion = attention_criterion(storage); + criterion.relative_l2 += 0x1p-9; + criterion.gross_relative_to_max_reference += 0x1p-9; + return criterion; +} + int verify_attention(const std::string& label, const std::vector& actual, const std::vector& reference, const ReductionCriterion& criterion) { return verify_reduction(label.c_str(), actual, reference, criterion); @@ -2076,6 +2090,54 @@ int run_batch_case(const Geometry& geometry, KvCacheStorage storage, std::cerr << label << ": workspace mismatch\n"; ++failures; } + + // Handing the Op a gate must produce exactly what applying sigmoid_mul afterwards + // produces -- on the route that folds the multiply into the reduce epilogue and on the + // routes that fall back to the standalone kernel alike. Re-running the Op is safe: + // appending the same k/v to the same rows again leaves the cache byte-identical, which + // the standalone-parity check above has just established. + { + const auto gate_bits = + to_bf16_bits(make_bf16_values(q.size(), test_case.seed + 97u, -3.0F, 3.0F)); + GuardedDeviceBuffer dgate(gate_bits.size() * sizeof(std::uint16_t)); + GuardedDeviceBuffer dexpected(output.size() * sizeof(std::uint16_t)); + dgate.copy_from_host(gate_bits.data(), gate_bits.size() * sizeof(std::uint16_t)); + dexpected.copy_from_host(output.data(), output.size() * sizeof(std::uint16_t)); + Tensor tgate(dgate.data(), DType::BF16, {kHeadDim, geometry.q_heads, width, batch}); + Tensor texpected(dexpected.data(), DType::BF16, + {kHeadDim, geometry.q_heads, width, batch}); + // The reference is a second ungated run of the Op, not the first run's output: that + // separates "the gate is exact" from "the Op repeats itself", and only the first is + // what this check is about. + ops::causal_softmax_attention(tq, tk, tv, tp, masked ? tvalid : Tensor{}, tlanes, + op_geometry(geometry), kAttentionScale, cache.view(), + envelope, workspace, texpected, device.stream); + cuda_synchronize(device.stream); + failures += + verify_exact((label + " ungated repeat").c_str(), + copy_from_guarded(dexpected, output.size()), output); + ops::sigmoid_mul(tgate, texpected, device.stream); + ops::causal_softmax_attention(tq, tk, tv, tp, masked ? tvalid : Tensor{}, tlanes, + op_geometry(geometry), kAttentionScale, cache.view(), + envelope, workspace, tout, device.stream, &tgate); + cuda_synchronize(device.stream); + const auto gated_output = copy_from_guarded(dout, q.size()); + failures += verify_exact((label + " fused gate").c_str(), gated_output, + copy_from_guarded(dexpected, output.size())); + // Both checks above compare production against production, so an error shared by the + // attention result, the sigmoid or the BF16 boundary would pass them. This one does + // not: the gated output is qualified against the same FP64 attention oracle the + // ungated output is judged by, multiplied by a host sigmoid of the gate. + std::vector gated_reference(reference.size()); + for (std::size_t i = 0; i < gated_reference.size(); ++i) + gated_reference[i] = + reference[i] / (1.0 + std::exp(-double(bf16_to_f32(gate_bits[i])))); + failures += verify_attention(label + " fused gate against the oracle", + bf16_bits_to_double(gated_output), gated_reference, + gated_attention_criterion(storage)); + failures += dgate.verify_guards((label + " fused gate input").c_str()); + failures += dout.verify_guards((label + " fused gate output").c_str()); + } } return failures; }