superme1one · Nov 5, 2026
diff --git a/csrc/layers/attention/backends/causal_prefill_policy.hpp 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 <cstddef>
+#include <string_view>
+
+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
@@ -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 <algorithm>
+#include <base/causal_softmax_infinilm.h>
 #include <base/flash_attn_varlen_func.h>
 #include <base/gemm.h>
 #include <base/reshape_and_cache_flash.h>
+#include <cstdlib>
 #include <cstring>
 #include <limits>
 #include <stdexcept>
@@ -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<infinilm::runtime::Tensor> 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<int64_t>(seq_len);
-            const auto columns = static_cast<int64_t>(total_seq_len);
-            infinilm::execution::submit<infini::ops::Fill>(
-                runtime, causal_policy, infini::ops::Fill::DefaultConfig(causal_device),
-                mask, -std::numeric_limits<double>::infinity(), mask);
-            infinilm::execution::submit<infini::ops::Triu>(
-                runtime, causal_policy, infini::ops::Triu::DefaultConfig(causal_device),
-                mask, columns - rows + 1, mask);
-            infinilm::execution::submit<infini::ops::Tril>(
-                runtime, causal_policy, infini::ops::Tril::DefaultConfig(causal_device),
-                attn_weight_softmax, columns - rows, attn_weight_softmax);
-            infinilm::execution::submit<infini::ops::Add>(
-                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<infini::ops::CausalSoftmaxInfinilm>(
+                runtime, causal_policy,
+                infini::ops::CausalSoftmaxInfinilm::DefaultConfig(causal_device),
+                attn_weight_softmax, attn_weight_softmax);
+        } else {
+            std::optional<infinilm::runtime::Tensor> 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<int64_t>(seq_len);
+                const auto columns = static_cast<int64_t>(total_seq_len);
+                infinilm::execution::submit<infini::ops::Fill>(
+                    runtime, causal_policy, infini::ops::Fill::DefaultConfig(causal_device),
+                    mask, -std::numeric_limits<double>::infinity(), mask);
+                infinilm::execution::submit<infini::ops::Triu>(
+                    runtime, causal_policy, infini::ops::Triu::DefaultConfig(causal_device),
+                    mask, columns - rows + 1, mask);
+                infinilm::execution::submit<infini::ops::Tril>(
+                    runtime, causal_policy, infini::ops::Tril::DefaultConfig(causal_device),
+                    attn_weight_softmax, columns - rows, attn_weight_softmax);
+                infinilm::execution::submit<infini::ops::Add>(
+                    runtime, causal_policy, infini::ops::Add::DefaultConfig(causal_device),
+                    attn_weight_softmax, mask, attn_weight_softmax);
+            }
+            infinilm::execution::submit<infini::ops::Softmax>(
+                runtime, causal_policy, infini::ops::Softmax::DefaultConfig(causal_device),
+                attn_weight_softmax, int64_t{-1}, std::optional<infini::ops::DataType>{},
+                attn_weight_softmax);
         }
-        infinilm::execution::submit<infini::ops::Softmax>(
-            runtime, causal_policy, infini::ops::Softmax::DefaultConfig(causal_device),
-            attn_weight_softmax, int64_t{-1}, std::optional<infini::ops::DataType>{},
-            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
@@ -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
@@ -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
@@ -0,0 +1,22 @@
+#include <cassert>
+#include <initializer_list>
+
+#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));
+}