From 5d2509fe88ff3d4d69374bef189e63e5d324ceba Mon Sep 17 00:00:00 2001 From: superme1one <2459380227@qq.com> Date: Thu, 5 Nov 2026 16:35:30 +0800 Subject: [PATCH] perf: add opt-in Iluvatar fused causal Prefill --- .../backends/causal_prefill_policy.hpp | 20 ++++ .../layers/attention/backends/static_attn.cpp | 66 ++++++----- .../iluvatar-four-optimizations-results.json | 60 ++++++++++ .../iluvatar-four-optimizations.md | 103 ++++++++++++++++++ test/runtime/causal_prefill_policy_test.cpp | 22 ++++ 5 files changed, 246 insertions(+), 25 deletions(-) create mode 100644 csrc/layers/attention/backends/causal_prefill_policy.hpp create mode 100644 docs/performance/iluvatar-four-optimizations-results.json create mode 100644 docs/performance/iluvatar-four-optimizations.md create mode 100644 test/runtime/causal_prefill_policy_test.cpp diff --git a/csrc/layers/attention/backends/causal_prefill_policy.hpp b/csrc/layers/attention/backends/causal_prefill_policy.hpp new file mode 100644 index 000000000..4cc80dd4b --- /dev/null +++ b/csrc/layers/attention/backends/causal_prefill_policy.hpp @@ -0,0 +1,20 @@ +#ifndef INFINILM_CAUSAL_PREFILL_POLICY_HPP_ +#define INFINILM_CAUSAL_PREFILL_POLICY_HPP_ + +#include +#include + +namespace infinilm::layers::attention::backends::detail { + +inline bool UseFusedCausalPrefill(const char *flag, bool iluvatar, bool bf16, + bool contiguous, std::size_t rank, + std::size_t heads, std::size_t query, + std::size_t key) { + return flag != nullptr && std::string_view(flag) == "1" && iluvatar && bf16 + && contiguous && rank == 3 && heads > 0 && query > 1 && query <= key + && key <= 4096 && heads <= (std::size_t{65535} * 65535) / query; +} + +} // namespace infinilm::layers::attention::backends::detail + +#endif diff --git a/csrc/layers/attention/backends/static_attn.cpp b/csrc/layers/attention/backends/static_attn.cpp index 31e5bb758..231bf986f 100644 --- a/csrc/layers/attention/backends/static_attn.cpp +++ b/csrc/layers/attention/backends/static_attn.cpp @@ -4,11 +4,14 @@ #include "../../../utils.hpp" #include "../attention_checks.hpp" #include "attention_layer.hpp" +#include "causal_prefill_policy.hpp" #include "runtime/runtime.hpp" #include +#include #include #include #include +#include #include #include #include @@ -268,34 +271,47 @@ infinilm::runtime::Tensor StaticAttentionImpl::forward(const AttentionLayer &lay if (seq_len > total_seq_len) { throw std::invalid_argument("causal_softmax has unsupported shape or dtype"); } - std::optional causal_mask; - if (seq_len != 1) { - causal_mask = infinilm::runtime::Tensor::empty( - runtime, {seq_len, total_seq_len}, attn_weight_softmax->dtype()); - } const auto causal_device = runtime.device().type(); constexpr auto causal_policy = infinilm::execution::RecordingPolicy::DeviceOnly; - if (causal_mask) { - const auto &mask = *causal_mask; - const auto rows = static_cast(seq_len); - const auto columns = static_cast(total_seq_len); - infinilm::execution::submit( - runtime, causal_policy, infini::ops::Fill::DefaultConfig(causal_device), - mask, -std::numeric_limits::infinity(), mask); - infinilm::execution::submit( - runtime, causal_policy, infini::ops::Triu::DefaultConfig(causal_device), - mask, columns - rows + 1, mask); - infinilm::execution::submit( - runtime, causal_policy, infini::ops::Tril::DefaultConfig(causal_device), - attn_weight_softmax, columns - rows, attn_weight_softmax); - infinilm::execution::submit( - runtime, causal_policy, infini::ops::Add::DefaultConfig(causal_device), - attn_weight_softmax, mask, attn_weight_softmax); + const bool fused = detail::UseFusedCausalPrefill( + std::getenv("INFINILM_ILUVATAR_FUSED_CAUSAL_PREFILL"), + causal_device == infinilm::runtime::Device::Type::kIluvatar, + attn_weight_softmax->dtype() == infinilm::runtime::DataType::kBFloat16, + attn_weight_softmax->is_contiguous(), attn_weight_softmax->ndim(), + batch_size * num_heads_, seq_len, total_seq_len); + if (fused) { + infinilm::execution::submit( + runtime, causal_policy, + infini::ops::CausalSoftmaxInfinilm::DefaultConfig(causal_device), + attn_weight_softmax, attn_weight_softmax); + } else { + std::optional causal_mask; + if (seq_len != 1) { + causal_mask = infinilm::runtime::Tensor::empty( + runtime, {seq_len, total_seq_len}, attn_weight_softmax->dtype()); + } + if (causal_mask) { + const auto &mask = *causal_mask; + const auto rows = static_cast(seq_len); + const auto columns = static_cast(total_seq_len); + infinilm::execution::submit( + runtime, causal_policy, infini::ops::Fill::DefaultConfig(causal_device), + mask, -std::numeric_limits::infinity(), mask); + infinilm::execution::submit( + runtime, causal_policy, infini::ops::Triu::DefaultConfig(causal_device), + mask, columns - rows + 1, mask); + infinilm::execution::submit( + runtime, causal_policy, infini::ops::Tril::DefaultConfig(causal_device), + attn_weight_softmax, columns - rows, attn_weight_softmax); + infinilm::execution::submit( + runtime, causal_policy, infini::ops::Add::DefaultConfig(causal_device), + attn_weight_softmax, mask, attn_weight_softmax); + } + infinilm::execution::submit( + runtime, causal_policy, infini::ops::Softmax::DefaultConfig(causal_device), + attn_weight_softmax, int64_t{-1}, std::optional{}, + attn_weight_softmax); } - infinilm::execution::submit( - runtime, causal_policy, infini::ops::Softmax::DefaultConfig(causal_device), - attn_weight_softmax, int64_t{-1}, std::optional{}, - attn_weight_softmax); INFINILM_TRACE_OP(CausalSoftmax, attn_weight_softmax); } diff --git a/docs/performance/iluvatar-four-optimizations-results.json b/docs/performance/iluvatar-four-optimizations-results.json new file mode 100644 index 000000000..d9af46bb9 --- /dev/null +++ b/docs/performance/iluvatar-four-optimizations-results.json @@ -0,0 +1,60 @@ +{ + "official_commit": "6e7f9537173cd33bb30539ebf8d92ab6d869660d", + "same_card_per_question": true, + "new_independent_questions": 100, + "model_processes": 50, + "formal_requests": 200, + "warmup_requests": 100, + "roles": { + "official": { + "correct": 82, + "total": 100, + "accuracy_pct": 82, + "completed": 100, + "parseable": 100, + "truncated": 0, + "errors": 0, + "formal_seconds": 749.203817545, + "requests_per_second": 0.13347502729988908, + "tokens": 13894, + "tokens_per_second": 18.54502029304659, + "median_ttft_ms": 2521.8428675, + "median_tpot_ms": 69.50871872384035, + "median_latency_ms": 11942.6921875 + }, + "combined": { + "correct": 84, + "total": 100, + "accuracy_pct": 84, + "completed": 100, + "parseable": 100, + "truncated": 0, + "errors": 0, + "formal_seconds": 505.326616325, + "requests_per_second": 0.19789181248209806, + "tokens": 14213, + "tokens_per_second": 28.126363308080595, + "median_ttft_ms": 1802.3383135, + "median_tpot_ms": 44.95782571872432, + "median_latency_ms": 8004.2276055 + } + }, + "changes": { + "accuracy_change_percentage_points": 2, + "time_reduction_pct": 32.55151609066005, + "request_throughput_change_pct": 48.26130137248717, + "token_throughput_change_pct": 51.66531426566574 + }, + "measurement_scope": "Historical four-item private combination; cumulative per-version same-card formal wave makespans", + "token_different_questions": 34, + "regression_ids": [ + 1097 + ], + "improvement_ids": [ + 672, + 1259, + 1203 + ], + "current_pr_revision_gpu_tested": false, + "original_summary_sha256": "6dfa8dbd43c100ec45b27886edb5294c7a8eb15e03578a9388862add06567ea9" +} diff --git a/docs/performance/iluvatar-four-optimizations.md b/docs/performance/iluvatar-four-optimizations.md new file mode 100644 index 000000000..7e3cc59cd --- /dev/null +++ b/docs/performance/iluvatar-four-optimizations.md @@ -0,0 +1,103 @@ +# Iluvatar inference optimizations and Prefill integration + +This opt-in change submits the fused causal Softmax through the normal runtime +execution path. Enable `INFINILM_ILUVATAR_FUSED_CAUSAL_PREFILL=1` only with the +companion InfiniOps FP32 causal provider. Unset it to retain the existing +Fill/Triu/Tril/Add + Softmax path. The environment flag is exact (`1`), with +eligibility restricted to contiguous 3D Iluvatar BF16, 2 <= Q <= K <= 4096. +The default remains unchanged, including Decode Attention. + +The operator submission preserves owner checks, stream/recording policy, +normal dispatch and runtime profiling. It removes the experimental direct +kernel entry point and counter. The compatibility operator already exists; +the companion PR specializes its FP32 accumulation order without adding APIs. +An older InfiniOps build exposes the same operator name but does not have this +FP32 provider: install the companion change before enabling the flag. + +## Four retained optimizations + +| Item | Bottleneck addressed | Narrow implementation | +| --- | --- | --- | +| Gate/Up GEMV | BLAS overhead and weight traffic at B1/B2 | BF16 K4096/N32768, shared weight loads for two rows | +| Down GEMV | Small-row Down projection | BF16 K16384/N4096, same four-accumulator scheme | +| Prefill causal Softmax | Mask allocations and separate masking launches | Register causal mask, FP32 max/exp/sum, BF16 exponent store/divide | +| Decode QK reduction | Repeated CTA barriers across eight token scores | Same FP32 trees, 48→6 QK barriers, CTA64/tile8/split4 | + +Gate/Up and Down share one InfiniOps PR. Prefill has an InfiniOps provider PR +and this InfiniLM integration PR. Decode has a separate InfiniOps PR. These +changes are independent of the earlier R5 modernization submissions. + +## Historical same-card application evidence + +The four-item combination was compared with official InfiniLM +`6e7f9537173cd33bb30539ebf8d92ab6d869660d` on TG-V200/CoreX using 100 new +fixed GSM8K targets, eight demonstrations from a disjoint test-corpus subset, +greedy decoding, max_tokens=512, EOS=[2,73440], and strict last-`####` numeric +scoring. The model, prompts, weights, answer references and scoring were held +fixed. This is an internal paired evaluation, not a standard benchmark split. + +Each target's two versions ran on the same physical card. The first 20 targets +ran serially on one card; the remaining 80 were sharded over two cards. Each +model had two fixed warmup requests before four formal requests. Fifty model +processes completed normally (scope return, maps, atexit and OS exit0), with +200 formal and 100 warmup requests. Repeated warmups add no independent quality +samples. Two physical cache rows and two active requests were used throughout. + +| Metric | Official | Four-item combination | +| --- | ---: | ---: | +| Correct / 100 | 82 | 84 | +| Completed / parseable | 100 / 100 | 100 / 100 | +| Errors / truncated | 0 / 0 | 0 / 0 | +| Sum of formal paired-wave makespans, s | 749.203818 | 505.326616 | +| Requests/s over that sum | 0.133475 | 0.197892 | +| Published tokens | 13894 | 14213 | +| Published tokens/s | 18.545020 | 28.126363 | +| Median first TokenOutput, ms | 2521.843 | 1802.338 | +| Median TPOT, ms | 69.509 | 44.958 | +| Median request completion latency, ms | 11942.692 | 8004.228 | + +Time fell **32.5515%**, request throughput rose **48.2613%**, and token +throughput rose **51.6653%**. The time denominator is the sum of each version's +same-card formal wave makespans, including queue and text processing, excluding +model load, warmup and final shutdown. It is not the parallel run's host wall +time. Output lengths changed; token throughput is not substituted for completed +request throughput. These are combined gains, not sums of component gains. + +### Token differences and task regressions + +34/100 targets had different token sequences: 26 remained correct, four +remained wrong, one regressed, and three improved. Accuracy increased by two +percentage points while preserving this individual regression: + +| Target ID | Reference | Official prediction | Combined prediction | Change | +| --- | ---: | ---: | ---: | --- | +| 1097 | 480 | 480 | 448 | Correct → wrong | +| 672 | 2 | 9 | 2 | Wrong → correct | +| 1259 | 7200 | 720 | 7200 | Wrong → correct | +| 1203 | 240 | 420 | 240 | Wrong → correct | + +99 request latencies decreased. Target 672 took 73.5054% longer as its output +grew from 57 to 238 tokens and its answer became correct. Token differences and +correctness are recorded separately. Finite sample accuracy is not a claim of +universal per-task non-regression. + +## Current submission and remaining validation + +This PR is rebased on main `cc6d52ae04f4c9ff17229bc7fc13a9ad105aca6d`, +which adds automatic operator-tuning initialization after the measured baseline. +The companion InfiniOps PRs are based on master +`1a68b261d19f0ed6f76bbabeab7a25f479f05e07`. The original GPU results apply +to the recorded private builds; this revised normal-dispatch integration still +needs a current-base Iluvatar smoke build and the same two-task application +regression. It is submitted as a draft with the flag off by default. + +Focused host check: + +```sh +c++ -std=c++17 -I. test/runtime/causal_prefill_policy_test.cpp -o /tmp/causal-prefill-policy +/tmp/causal-prefill-policy +``` + +The companion provider includes focused device tests for masked values, +rectangular masks and in-place/out-of-place output. No dataset expansion is +needed to validate this integration. diff --git a/test/runtime/causal_prefill_policy_test.cpp b/test/runtime/causal_prefill_policy_test.cpp new file mode 100644 index 000000000..0b69e55c2 --- /dev/null +++ b/test/runtime/causal_prefill_policy_test.cpp @@ -0,0 +1,22 @@ +#include +#include + +#include "csrc/layers/attention/backends/causal_prefill_policy.hpp" + +int main() { + using infinilm::layers::attention::backends::detail::UseFusedCausalPrefill; + for (const char *flag : {"", "0", "true", "01", "1 "}) { + assert(!UseFusedCausalPrefill(flag, true, true, true, 3, 32, 2910, 2910)); + } + assert(!UseFusedCausalPrefill(nullptr, true, true, true, 3, 32, 2910, 2910)); + assert(UseFusedCausalPrefill("1", true, true, true, 3, 32, 2910, 2910)); + assert(UseFusedCausalPrefill("1", true, true, true, 3, 64, 521, 4096)); + assert(!UseFusedCausalPrefill("1", false, true, true, 3, 32, 2910, 2910)); + assert(!UseFusedCausalPrefill("1", true, false, true, 3, 32, 2910, 2910)); + assert(!UseFusedCausalPrefill("1", true, true, false, 3, 32, 2910, 2910)); + assert(!UseFusedCausalPrefill("1", true, true, true, 2, 32, 2910, 2910)); + assert(!UseFusedCausalPrefill("1", true, true, true, 3, 0, 2910, 2910)); + assert(!UseFusedCausalPrefill("1", true, true, true, 3, 32, 1, 4096)); + assert(!UseFusedCausalPrefill("1", true, true, true, 3, 32, 2910, 2909)); + assert(!UseFusedCausalPrefill("1", true, true, true, 3, 32, 2910, 4097)); +}