From 9fe228a391f4d9aeabcd38738ba40cc4d1958957 Mon Sep 17 00:00:00 2001 From: Mykhailo Dementii Date: Wed, 9 Sep 2026 19:19:55 +0200 Subject: [PATCH 1/3] perf(ops): let causal attention finish with the sigmoid gate The three call sites in TextContext all issued causal_softmax_attention and then ops::sigmoid_mul over the whole output, which is one more graph node per full-attention layer and per MTP tail. 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, and the bytes are the bytes the caller's own sigmoid_mul produced: the reduce result is rounded to BF16 first, widened back and multiplied in FP32, because the standalone kernel reads what attention stored rather than the accumulator. The operator bench gains --gate off|standalone|fused. Over the 65 cells where the fold applies, kernel time is -8.33% median with 58 of 65 improved; over the 255 cells where the route does not fold, the median is exactly zero, which is the control that the standalone path is untouched. End to end on Qwen3.6-35B-A3B, ten mirrored passes with a zero-control arm: decode +0.311% median, all ten passes positive, fully separated from a control spanning -0.097 to +0.088%. Prefill is unchanged within that same control. The Op test now runs each case a second time with a gate and judges it against a third, ungated run, which separates "the Op repeats itself" from "the gate is exact" instead of conflating the two. Dropping the rounding before the multiply makes that check fail, so it is not vacuous. Co-Authored-By: Claude Opus 5 --- bench/ops/causal_softmax_attention_bench.cu | 51 +++++++++-- include/ninfer/ops/softmax_attention.h | 8 +- .../causal_cache/causal_softmax_attention.cpp | 54 ++++++++---- .../dense/causal_cache/launch.h | 15 ++-- .../dense/causal_cache/small_t.cu | 36 ++++---- .../dense/causal_cache/small_t.cuh | 28 ++++-- .../qwen3_6/impl/runtime/text_context_impl.h | 43 +++++---- tests/ops/softmax_attention/causal_cache.cpp | 88 +++++++++++++------ 8 files changed, 224 insertions(+), 99 deletions(-) diff --git a/bench/ops/causal_softmax_attention_bench.cu b/bench/ops/causal_softmax_attention_bench.cu index 5e6df581a5..69a45c1390 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}; @@ -112,7 +120,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 +243,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"); @@ -336,9 +354,9 @@ PagedKVLayerView make_cache_view(DeviceBuffer& k, DeviceBuffer& v, DeviceBuffer& .k_pages = Tensor( k.p, layout.key.data_dtype, {layout.key.data_leading_extent, kPagedKVPageSize, geometry.kv_heads, physical_pages}), - .v_pages = Tensor(v.p, layout.value.data_dtype, - {layout.value.data_leading_extent, kPagedKVPageSize, geometry.kv_heads, - physical_pages}), + .v_pages = Tensor(v.p, layout.value.data_dtype, + {layout.value.data_leading_extent, kPagedKVPageSize, geometry.kv_heads, + physical_pages}), .k_scale_pages = layout.key.has_scale() ? Tensor(k_scale.p, layout.key.scale_dtype, {layout.key.scale_leading_extent, kPagedKVPageSize, @@ -417,8 +435,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 +470,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 +481,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 +547,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 +569,7 @@ public: private: PagedKVStorageLayout storage_layout_; + GateMode gate_; std::int32_t batch_; bool masked_; std::int32_t visible_; @@ -561,6 +589,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 +599,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_; @@ -874,7 +904,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 +914,9 @@ int main(int argc, char** argv) { return 0; } + std::printf("gate=%s\n", options.gate == GateMode::Off ? "off" + : options.gate == GateMode::Standalone ? "standalone" + : "fused"); std::vector results; const std::vector context_profiles = options.row_contexts.empty() ? options.contexts : std::vector{0}; @@ -895,7 +928,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) || 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/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..0430a7d955 100644 --- a/src/ops/softmax_attention/dense/causal_cache/small_t.cu +++ b/src/ops/softmax_attention/dense/causal_cache/small_t.cu @@ -80,7 +80,7 @@ std::int32_t causal_small_t_launch_capacity(CausalAttentionExecutionEnvelope env const auto include = [&](std::uint32_t window) { if (window < envelope.min_visible_keys || window > envelope.max_visible_keys) { return; } const auto splits = causal_small_t_split_count(static_cast(window), - tokens, storage); + tokens, storage); capacity = capacity > splits ? capacity : splits; }; include(envelope.min_visible_keys); @@ -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( @@ -333,16 +334,16 @@ void causal_attention_small_t_launch_for(const Tensor& q, CacheInput input, cons constexpr int kDChunk = Geometry::QHeads == 24 ? 256 : 64; const auto launch_reduce = [&]() { const dim3 grid(Geometry::QHeads, div_up(kCausalHeadDim, kDChunk), - invocation.width * invocation.batch_size); + invocation.width * invocation.batch_size); causal_attention_small_t_reduce_output_kernel<<>>( + Offset><<>>( static_cast(partial_acc.data), static_cast(partial_m.data), static_cast(partial_l.data), static_cast(pos.data), invocation.valid_columns - ? static_cast(invocation.valid_columns->data) - : nullptr, + ? 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/src/targets/qwen3_6/impl/runtime/text_context_impl.h b/src/targets/qwen3_6/impl/runtime/text_context_impl.h index 8cdef782a1..63ece2d8e2 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_context_impl.h @@ -144,7 +144,7 @@ class ScopedValue { void DFlashFeatureSink::begin(const Tensor& value) { const bool prefill = features != nullptr && positions != nullptr && batch_features == nullptr; const bool batch = batch_features != nullptr && batch_lanes != nullptr && - batch_valid_columns != nullptr && batch_width > 0 && batch_size > 0; + batch_valid_columns != nullptr && batch_width > 0 && batch_size > 0; if ((!prefill && !batch) || layers.empty()) { throw std::logic_error("DFlash feature sink is incomplete"); } @@ -279,15 +279,15 @@ void TextContext::bind() { } const auto& source = *weights_.mtp; mtp_ = MtpW{&source, - &source.input_projection, - &source.embedding_norm, - &source.hidden_norm, - &source.input_norm, - &source.query_norm, - &source.key_norm, - &source.output, - &source.post_attention_norm, - &source.final_norm}; + &source.input_projection, + &source.embedding_norm, + &source.hidden_norm, + &source.input_norm, + &source.query_norm, + &source.key_norm, + &source.output, + &source.post_attention_norm, + &source.final_norm}; } for (int layer = 0; layer < kCfg.n_layers; ++layer) { @@ -396,17 +396,18 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po Tensor k_batch = kn.view({kCfg.head_dim, kCfg.n_kv, width, active_sequence_batch_}); Tensor v_batch = v.view({kCfg.head_dim, kCfg.n_kv, width, active_sequence_batch_}); Tensor a_batch = a.view({kCfg.head_dim, kCfg.n_q, width, active_sequence_batch_}); + Tensor gate_batch = gate.view({kCfg.head_dim, kCfg.n_q, 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_, *active_backend_kv_table_rows_, {kCfg.head_dim, kCfg.n_q, kCfg.n_kv}, kAttnScale, - 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, {kCfg.head_dim, kCfg.n_q, kCfg.n_kv}, kAttnScale, - 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_recipe::mtp_post_attention(work_, T); Tensor o = post.output; @@ -499,8 +500,8 @@ void TextContext::mtp_prefill_chunk(const Tensor& ids, const Tensor& hidden, if (final_chunk) { const std::size_t column_bytes = static_cast(kCfg.hidden) * dtype_size(DType::BF16); - const auto* x_src = static_cast(x.data) + - static_cast(T - 1) * column_bytes; + const auto* x_src = static_cast(x.data) + + static_cast(T - 1) * column_bytes; const auto* ah_src = static_cast(ah.data) + static_cast(T - 1) * column_bytes; CUDA_CHECK( @@ -528,7 +529,7 @@ void TextContext::mtp_prefill_chunk(const Tensor& ids, const Tensor& hidden, for (int axis = 0; axis < 3; ++axis) { const auto* src = static_cast(rope_positions.data) + static_cast(axis) * T + (T - 1); - auto* dst = static_cast(last_rope_position.data) + axis; + auto* dst = static_cast(last_rope_position.data) + axis; CUDA_CHECK( cudaMemcpyAsync(dst, src, sizeof(std::int32_t), cudaMemcpyDeviceToDevice, s)); } @@ -860,19 +861,23 @@ void TextContext::attn_mix(const FullLayerW& w, Tensor& x, int fidx, Phase ph) { Tensor k_batch = kn.view({kCfg.head_dim, kCfg.n_kv, width, active_sequence_batch_}); Tensor v_batch = v.view({kCfg.head_dim, kCfg.n_kv, width, active_sequence_batch_}); Tensor a_batch = a.view({kCfg.head_dim, kCfg.n_q, width, active_sequence_batch_}); + Tensor gate_batch = gate.view({kCfg.head_dim, kCfg.n_q, 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; every other route applies it inside the Op. One contract either + // way, and the result is the same bytes as the standalone multiply produced. ops::causal_softmax_attention(q_batch, k_batch, v_batch, position_batch, valid, kv_table_rows, {kCfg.head_dim, kCfg.n_q, kCfg.n_kv}, kAttnScale, batch_text_kv_->batch_layer_view(fidx), - *active_causal_attention_envelope_, work_, a_batch, s); + *active_causal_attention_envelope_, work_, a_batch, s, + &gate_batch); } else { ops::causal_softmax_attention(qn, kn, v, cache_positions, Tensor{}, kv_table_rows, {kCfg.head_dim, kCfg.n_q, kCfg.n_kv}, kAttnScale, batch_text_kv_->batch_layer_view(fidx), - *active_causal_attention_envelope_, work_, a, s); + *active_causal_attention_envelope_, work_, a, s, &gate); } - ops::sigmoid_mul(gate, a, s); Variant::attention_output_projection(a.view({kCfg.q_size, T}), *w.o_proj, x, ph, work_, s); } diff --git a/tests/ops/softmax_attention/causal_cache.cpp b/tests/ops/softmax_attention/causal_cache.cpp index 085a546fc8..99023fd78c 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" @@ -471,9 +472,9 @@ std::uint8_t encode_e2m1_rne_satfinite(float value) { const float lower_value = decode_e2m1(lower); const float lower_error = magnitude - lower_value; const float upper_error = upper_value - magnitude; - selected = lower_error < upper_error ? lower - : upper_error < lower_error ? upper - : ((lower & 1U) == 0U ? lower : upper); + selected = lower_error < upper_error ? lower + : upper_error < lower_error ? upper + : ((lower & 1U) == 0U ? lower : upper); break; } } @@ -698,8 +699,8 @@ HostCache make_cache(const Geometry& geometry, KvCacheStorage storage, std::int3 cache_index(geometry, logical_capacity, head, position, 0); const std::size_t k_scale = fp8_scale_index(geometry, logical_capacity, head, position); - const std::size_t v_code = logical_plane_index(kNvfp4CodeBytes, geometry, - logical_capacity, head, position, 0); + const std::size_t v_code = logical_plane_index(kNvfp4CodeBytes, geometry, + logical_capacity, head, position, 0); const std::size_t v_scale = logical_plane_index( kNvfp4QuantGroups, geometry, logical_capacity, head, position, 0); encode_fp8_rotated_row(logical_k, source, cache.k_fp8, source, cache.k_scale, @@ -1064,19 +1065,19 @@ class DeviceCache { if (storage_ == KvCacheStorage::BFloat16) { const auto k_physical = copy_from_guarded(k_, k_code_elements_); const auto v_physical = copy_from_guarded(v_, v_code_elements_); - cache.k_bf16 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); - cache.v_fp16 = gather_paged(v_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_bf16 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); + cache.v_fp16 = gather_paged(v_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); } else if (storage_ == KvCacheStorage::Int8Group64) { const auto k_physical = copy_from_guarded(k_, k_code_elements_); const auto v_physical = copy_from_guarded(v_, v_code_elements_); const auto ks_physical = copy_from_guarded(k_scale_, k_scale_elements_); const auto vs_physical = copy_from_guarded(v_scale_, v_scale_elements_); - cache.k_i8 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); - cache.v_i8 = gather_paged(v_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_i8 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); + cache.v_i8 = gather_paged(v_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); cache.k_scale = gather_paged(ks_physical, kQuantGroups, geometry_, logical_capacity_, block_table_host_); cache.v_scale = gather_paged(vs_physical, kQuantGroups, geometry_, @@ -1086,10 +1087,10 @@ class DeviceCache { const auto v_physical = copy_from_guarded(v_, v_code_elements_); const auto ks_physical = copy_from_guarded(k_scale_, k_scale_elements_); const auto vs_physical = copy_from_guarded(v_scale_, v_scale_elements_); - cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); - cache.v_fp8 = gather_paged(v_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); + cache.v_fp8 = gather_paged(v_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); cache.k_scale = gather_paged(ks_physical, kFp8QuantGroups, geometry_, logical_capacity_, block_table_host_); cache.v_scale = gather_paged(vs_physical, kFp8QuantGroups, geometry_, @@ -1099,8 +1100,8 @@ class DeviceCache { const auto v_physical = copy_from_guarded(v_, v_code_elements_); const auto ks_physical = copy_from_guarded(k_scale_, k_scale_elements_); const auto vs_physical = copy_from_guarded(v_scale_, v_scale_elements_); - cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); cache.v_nvfp4 = gather_paged(v_physical, kNvfp4CodeBytes, geometry_, logical_capacity_, block_table_host_); cache.k_scale = gather_paged(ks_physical, kFp8QuantGroups, geometry_, @@ -1722,13 +1723,13 @@ int run_a1_case(const Geometry& geometry, KvCacheStorage storage, const Attentio const std::int32_t total = test_case.base + test_case.tokens; const std::int32_t max_context = static_cast( std::max(static_cast(total + 3), test_case.envelope_max)); - const std::size_t q_elements = static_cast(kHeadDim) * - static_cast(geometry.q_heads) * - static_cast(test_case.tokens); + const std::size_t q_elements = static_cast(kHeadDim) * + static_cast(geometry.q_heads) * + static_cast(test_case.tokens); const std::size_t kv_elements = static_cast(kHeadDim) * static_cast(geometry.kv_heads) * static_cast(test_case.tokens); - std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); + std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); if (test_case.zero_q) std::fill(q.begin(), q.end(), 0.0f); std::vector k = make_bf16_values(kv_elements, test_case.seed + 1u, -0.25f, 0.25f); std::vector v = make_bf16_values(kv_elements, test_case.seed + 2u, -1.0f, 1.0f); @@ -1824,7 +1825,7 @@ int run_a3_case(const Geometry& geometry, KvCacheStorage storage, const Attentio const std::size_t q_elements = static_cast(kHeadDim) * static_cast(geometry.q_heads) * static_cast(test_case.tokens); - std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); + std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); if (test_case.zero_q) std::fill(q.begin(), q.end(), 0.0f); std::vector positions(static_cast(test_case.tokens)); for (std::int32_t token = 0; token < test_case.tokens; ++token) { @@ -2040,7 +2041,7 @@ int run_batch_case(const Geometry& geometry, KvCacheStorage storage, const std::string label = std::string("causal batch ") + geometry.name + " " + cache_name(storage) + " W=" + std::to_string(width) + " B=" + std::to_string(batch) + " phase=" + std::to_string(phase); - const auto output = copy_from_guarded(dout, q.size()); + const auto output = copy_from_guarded(dout, q.size()); failures += verify_attention(label, bf16_bits_to_double(output), reference, attention_criterion(storage)); failures += verify_invalid_columns_zero(label, output, geometry, width, valid); @@ -2076,6 +2077,43 @@ 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); + failures += verify_exact((label + " fused gate").c_str(), + copy_from_guarded(dout, q.size()), + copy_from_guarded(dexpected, output.size())); + failures += dgate.verify_guards((label + " fused gate input").c_str()); + failures += dout.verify_guards((label + " fused gate output").c_str()); + } } return failures; } From 2b7fb9e76dc0b808be4ae9186a33d479802791d6 Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Thu, 10 Sep 2026 00:23:58 +0200 Subject: [PATCH 2/3] test(attention): qualify the gated output against the FP64 oracle, and record the gate mode The gated checks compared production against production, so a defect shared by the attention result, the sigmoid or the BF16 boundary would have passed them. Judge the gated output against the same FP64 attention oracle the ungated output is judged by, multiplied by a host sigmoid, with the criterion widened by the one extra BF16 rounding the gated route performs. The bench now carries the gate mode each row was measured in -- the effective one, so a cached row reads standalone -- on the console line and in the CSV. Co-Authored-By: Claude Opus 5 --- bench/ops/causal_softmax_attention_bench.cu | 46 +++++++----- tests/ops/softmax_attention/causal_cache.cpp | 78 +++++++++++++------- 2 files changed, 79 insertions(+), 45 deletions(-) diff --git a/bench/ops/causal_softmax_attention_bench.cu b/bench/ops/causal_softmax_attention_bench.cu index 69a45c1390..91f7e619c5 100644 --- a/bench/ops/causal_softmax_attention_bench.cu +++ b/bench/ops/causal_softmax_attention_bench.cu @@ -108,8 +108,15 @@ 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"; +} + [[noreturn]] void usage(const char* message) { std::fprintf(stderr, "error: %s\n" @@ -354,9 +361,9 @@ PagedKVLayerView make_cache_view(DeviceBuffer& k, DeviceBuffer& v, DeviceBuffer& .k_pages = Tensor( k.p, layout.key.data_dtype, {layout.key.data_leading_extent, kPagedKVPageSize, geometry.kv_heads, physical_pages}), - .v_pages = Tensor(v.p, layout.value.data_dtype, - {layout.value.data_leading_extent, kPagedKVPageSize, geometry.kv_heads, - physical_pages}), + .v_pages = Tensor(v.p, layout.value.data_dtype, + {layout.value.data_leading_extent, kPagedKVPageSize, geometry.kv_heads, + physical_pages}), .k_scale_pages = layout.key.has_scale() ? Tensor(k_scale.p, layout.key.scale_dtype, {layout.key.scale_leading_extent, kPagedKVPageSize, @@ -736,15 +743,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 " @@ -761,7 +768,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," @@ -771,12 +778,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; @@ -914,9 +922,7 @@ int main(int argc, char** argv) { return 0; } - std::printf("gate=%s\n", options.gate == GateMode::Off ? "off" - : options.gate == GateMode::Standalone ? "standalone" - : "fused"); + 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}; @@ -998,6 +1004,10 @@ int main(int argc, char** argv) { result.workspace_peak = data.workspace_peak(); result.graph_calls = execution == Execution::Graph ? options.graph_calls : 1; + result.gate = entry == Entry::Cached && + options.gate == GateMode::Fused + ? GateMode::Standalone + : options.gate; result.timing.median_us /= result.graph_calls; result.timing.min_us /= result.graph_calls; result.timing.p95_us /= result.graph_calls; diff --git a/tests/ops/softmax_attention/causal_cache.cpp b/tests/ops/softmax_attention/causal_cache.cpp index 99023fd78c..5cb35d121e 100644 --- a/tests/ops/softmax_attention/causal_cache.cpp +++ b/tests/ops/softmax_attention/causal_cache.cpp @@ -472,9 +472,9 @@ std::uint8_t encode_e2m1_rne_satfinite(float value) { const float lower_value = decode_e2m1(lower); const float lower_error = magnitude - lower_value; const float upper_error = upper_value - magnitude; - selected = lower_error < upper_error ? lower - : upper_error < lower_error ? upper - : ((lower & 1U) == 0U ? lower : upper); + selected = lower_error < upper_error ? lower + : upper_error < lower_error ? upper + : ((lower & 1U) == 0U ? lower : upper); break; } } @@ -699,8 +699,8 @@ HostCache make_cache(const Geometry& geometry, KvCacheStorage storage, std::int3 cache_index(geometry, logical_capacity, head, position, 0); const std::size_t k_scale = fp8_scale_index(geometry, logical_capacity, head, position); - const std::size_t v_code = logical_plane_index(kNvfp4CodeBytes, geometry, - logical_capacity, head, position, 0); + const std::size_t v_code = logical_plane_index(kNvfp4CodeBytes, geometry, + logical_capacity, head, position, 0); const std::size_t v_scale = logical_plane_index( kNvfp4QuantGroups, geometry, logical_capacity, head, position, 0); encode_fp8_rotated_row(logical_k, source, cache.k_fp8, source, cache.k_scale, @@ -1065,19 +1065,19 @@ class DeviceCache { if (storage_ == KvCacheStorage::BFloat16) { const auto k_physical = copy_from_guarded(k_, k_code_elements_); const auto v_physical = copy_from_guarded(v_, v_code_elements_); - cache.k_bf16 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); - cache.v_fp16 = gather_paged(v_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_bf16 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); + cache.v_fp16 = gather_paged(v_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); } else if (storage_ == KvCacheStorage::Int8Group64) { const auto k_physical = copy_from_guarded(k_, k_code_elements_); const auto v_physical = copy_from_guarded(v_, v_code_elements_); const auto ks_physical = copy_from_guarded(k_scale_, k_scale_elements_); const auto vs_physical = copy_from_guarded(v_scale_, v_scale_elements_); - cache.k_i8 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); - cache.v_i8 = gather_paged(v_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_i8 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); + cache.v_i8 = gather_paged(v_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); cache.k_scale = gather_paged(ks_physical, kQuantGroups, geometry_, logical_capacity_, block_table_host_); cache.v_scale = gather_paged(vs_physical, kQuantGroups, geometry_, @@ -1087,10 +1087,10 @@ class DeviceCache { const auto v_physical = copy_from_guarded(v_, v_code_elements_); const auto ks_physical = copy_from_guarded(k_scale_, k_scale_elements_); const auto vs_physical = copy_from_guarded(v_scale_, v_scale_elements_); - cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); - cache.v_fp8 = gather_paged(v_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); + cache.v_fp8 = gather_paged(v_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); cache.k_scale = gather_paged(ks_physical, kFp8QuantGroups, geometry_, logical_capacity_, block_table_host_); cache.v_scale = gather_paged(vs_physical, kFp8QuantGroups, geometry_, @@ -1100,8 +1100,8 @@ class DeviceCache { const auto v_physical = copy_from_guarded(v_, v_code_elements_); const auto ks_physical = copy_from_guarded(k_scale_, k_scale_elements_); const auto vs_physical = copy_from_guarded(v_scale_, v_scale_elements_); - cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, - logical_capacity_, block_table_host_); + cache.k_fp8 = gather_paged(k_physical, kHeadDim, geometry_, + logical_capacity_, block_table_host_); cache.v_nvfp4 = gather_paged(v_physical, kNvfp4CodeBytes, geometry_, logical_capacity_, block_table_host_); cache.k_scale = gather_paged(ks_physical, kFp8QuantGroups, geometry_, @@ -1660,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); @@ -1723,13 +1736,13 @@ int run_a1_case(const Geometry& geometry, KvCacheStorage storage, const Attentio const std::int32_t total = test_case.base + test_case.tokens; const std::int32_t max_context = static_cast( std::max(static_cast(total + 3), test_case.envelope_max)); - const std::size_t q_elements = static_cast(kHeadDim) * - static_cast(geometry.q_heads) * - static_cast(test_case.tokens); + const std::size_t q_elements = static_cast(kHeadDim) * + static_cast(geometry.q_heads) * + static_cast(test_case.tokens); const std::size_t kv_elements = static_cast(kHeadDim) * static_cast(geometry.kv_heads) * static_cast(test_case.tokens); - std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); + std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); if (test_case.zero_q) std::fill(q.begin(), q.end(), 0.0f); std::vector k = make_bf16_values(kv_elements, test_case.seed + 1u, -0.25f, 0.25f); std::vector v = make_bf16_values(kv_elements, test_case.seed + 2u, -1.0f, 1.0f); @@ -1825,7 +1838,7 @@ int run_a3_case(const Geometry& geometry, KvCacheStorage storage, const Attentio const std::size_t q_elements = static_cast(kHeadDim) * static_cast(geometry.q_heads) * static_cast(test_case.tokens); - std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); + std::vector q = make_bf16_values(q_elements, test_case.seed, -0.25f, 0.25f); if (test_case.zero_q) std::fill(q.begin(), q.end(), 0.0f); std::vector positions(static_cast(test_case.tokens)); for (std::int32_t token = 0; token < test_case.tokens; ++token) { @@ -2041,7 +2054,7 @@ int run_batch_case(const Geometry& geometry, KvCacheStorage storage, const std::string label = std::string("causal batch ") + geometry.name + " " + cache_name(storage) + " W=" + std::to_string(width) + " B=" + std::to_string(batch) + " phase=" + std::to_string(phase); - const auto output = copy_from_guarded(dout, q.size()); + const auto output = copy_from_guarded(dout, q.size()); failures += verify_attention(label, bf16_bits_to_double(output), reference, attention_criterion(storage)); failures += verify_invalid_columns_zero(label, output, geometry, width, valid); @@ -2108,9 +2121,20 @@ int run_batch_case(const Geometry& geometry, KvCacheStorage storage, op_geometry(geometry), kAttentionScale, cache.view(), envelope, workspace, tout, device.stream, &tgate); cuda_synchronize(device.stream); - failures += verify_exact((label + " fused gate").c_str(), - copy_from_guarded(dout, q.size()), + 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()); } From d63848ece078d5614a49779b7e1f17008e54311f Mon Sep 17 00:00:00 2001 From: MichaelDementii <136074657+MichaelDementii@users.noreply.github.com> Date: Thu, 10 Sep 2026 00:55:04 +0200 Subject: [PATCH 3/3] bench(attention): name the gate mode in the profile header too The --profile path returns before the console line that named the mode, so profiles captured with --gate off, standalone and fused carried no way to tell them apart. The PROFILE header now names the effective mode, through the same helper the CSV and the result line use. Co-Authored-By: Claude Opus 5 --- bench/ops/causal_softmax_attention_bench.cu | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/bench/ops/causal_softmax_attention_bench.cu b/bench/ops/causal_softmax_attention_bench.cu index 91f7e619c5..6c7e1e4ad8 100644 --- a/bench/ops/causal_softmax_attention_bench.cu +++ b/bench/ops/causal_softmax_attention_bench.cu @@ -117,6 +117,13 @@ 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" @@ -819,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); @@ -1004,10 +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 = entry == Entry::Cached && - options.gate == GateMode::Fused - ? GateMode::Standalone - : options.gate; + 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;