Enable CUDA-graph capture for causal fixed-capacity decode via on-device KV length

crates/onnx-runtime-ep-cuda/src/kernels/standard_attention.rs

@@ -170,6 +170,7 @@ const ATTN_SPLIT_MIN_CHUNK: u64 = 128;
 fn attention_split_config(
     is_decode: bool,
     dev_length_eligible: bool,
+    is_causal: bool,
     want_qk: bool,
     cap: u64,
     total_rows: u64,
@@ -194,7 +195,22 @@ fn attention_split_config(
     // Split only on the capture-safe fixed-capacity decode route, and never when
     // the raw QK scores are an observable output (the score scratch is reused as
     // per-key storage inside the split kernel).
-    if !enabled || !is_decode || !dev_length_eligible || want_qk || cap == 0 {
+    //
+    // Also never for the causal form. The multi-split path reorders the fp32
+    // key reduction (chunk-local online softmax merged by `attention_combine`)
+    // versus the monolithic `attention_row` serial ascending reduction. That is
+    // a legitimately different — but not bit-identical — fp32 result, and at a
+    // near-tie between the top-1 and top-2 logits it can flip the greedy argmax.
+    // The monolithic reduction is the reference: it matches both the parent CUDA
+    // build and an independent CPU dense-prefill oracle token-for-token, whereas
+    // the split reduction diverges at such a tie (measured on granite-1b: the
+    // decode stream flips at the first near-tie once causal decode is routed
+    // through the split kernel). The non-causal fixed-capacity decode has always
+    // taken the split path, so its byte-identity baseline already includes it;
+    // gating the causal form out keeps the newly capture-enabled causal decode
+    // bit-for-bit equal to the pre-capture monolithic path while still capturing
+    // (the split is a throughput optimization, not a capture requirement).
+    if !enabled || !is_decode || !dev_length_eligible || is_causal || want_qk || cap == 0 {
         return None;
     }
     attention_split_geometry(chunk_override, cap, total_rows)
@@ -506,7 +522,13 @@ extern "C" __global__ void attention_row(
   }
 
   // Stage 3: attention mask + causal frontier + padding frontier.
-  const long long offset = offsets[b];
+  // When the valid length is read on-device (`dev_len`), the causal frontier is
+  // derived from it too (`offset = total_seq - q_seq`, i.e. the on-device past
+  // length) rather than the host `offsets[b]`, which under a frozen fixed-
+  // capacity binding reports the padded capacity. Query row `i` (absolute
+  // position `past + i`) attends keys `[0, past + i]`.
+  const long long offset =
+      (dev_len != nullptr) ? (long long)total_seq - (long long)q_seq : offsets[b];
   const long long pad_limit = pad_limits[b];
   const long long causal_limit = (long long)i + offset;
   for (unsigned long long j = tid; j < total_seq; j += nthreads) {
@@ -658,7 +680,11 @@ extern "C" __global__ void attention_split(
   const unsigned long long qoff = q_is_3d
       ? (b * q_seq + i) * (q_heads * head_size) + qh * head_size
       : ((b * q_heads + qh) * q_seq + i) * head_size;
-  const long long offset = offsets[b];
+  // See `attention_row`: with an on-device valid length the causal frontier is
+  // derived on-device (`total_seq - q_seq`) instead of the host `offsets[b]`,
+  // which reports padded capacity under a frozen fixed-capacity binding.
+  const long long offset =
+      (dev_len != nullptr) ? (long long)total_seq - (long long)q_seq : offsets[b];
   const long long pad_limit = pad_limits[b];
   const long long causal_limit = (long long)i + offset;
   const bool single = (total_seq <= chunk);
@@ -1239,18 +1265,19 @@ fn std_attention_carve(
 }
 
 /// Static route signal for staged dense KV growth. On the session's
-/// mask-driven, non-causal fixed-capacity decode contract, the mask key width
-/// and past-cache sequence extent are the same physical capacity, so execution
-/// takes the fixed-stride append path (`stage_key/value == false`). Other
-/// in-op-cache routes may grow a dense aliased cache and reserve only the
-/// present outputs that actually exist. Actual pointer aliasing is checked by
-/// `run`; a non-aliased dispatch can safely consume less than this conservative
-/// dense-route reservation.
+/// mask-driven fixed-capacity decode contract, the mask key width and past-cache
+/// sequence extent are the same physical capacity, so execution takes the
+/// fixed-stride append path (`stage_key/value == false`) in both the causal and
+/// non-causal form. Other in-op-cache routes may grow a dense aliased cache and
+/// reserve only the present outputs that actually exist. Actual pointer aliasing
+/// is checked by `run`; a non-aliased dispatch can safely consume less than this
+/// conservative dense-route reservation.
 fn std_attention_staging_route(
     inputs: &[TensorMetadata<'_>],
     is_causal: bool,
     output_count: usize,
 ) -> (bool, bool) {
+    let _ = is_causal;
     let has_past_key = inputs.get(4).is_some_and(|input| input.present);
     let has_past_value = inputs.get(5).is_some_and(|input| input.present);
     if !has_past_key || !has_past_value {
@@ -1266,8 +1293,13 @@ fn std_attention_staging_route(
         .get(3)
         .filter(|mask| mask.present)
         .and_then(|mask| mask.shape.last().copied());
-    let fixed_capacity_append = !is_causal
-        && past_key_capacity.is_some()
+    // Fixed-capacity append when the past caches and the mask all share one
+    // physical key width (the frozen decode contract). The `is_causal` attribute
+    // is irrelevant: for a growing eager cache the past extent (valid length) and
+    // the mask width (valid length + current) differ, so this stays false there
+    // regardless of causality; it only holds once the KV+mask are frozen to a
+    // common capacity, where execution appends in place without staging.
+    let fixed_capacity_append = past_key_capacity.is_some()
         && past_key_capacity == past_value_capacity
         && past_key_capacity == mask_key_capacity;
     if fixed_capacity_append {
@@ -2169,13 +2201,19 @@ impl StandardAttentionKernel {
             // (whose extent is frozen when a CUDA graph is captured). Scanning
             // the LAST query row makes it correct for BOTH phases — prefill
             // returns prompt_len, decode returns total_seq (row 0 would report 1
-            // under a causal prefill mask). Eligible only for the fixed-capacity,
-            // mask-masked (non-causal) fixed-slot-append path; every other path
-            // passes a null pointer and keeps the host-derived length, so
-            // eager/dense/GQA are bit-for-bit unchanged.
+            // under a causal prefill mask). Eligible for the fixed-capacity,
+            // masked fixed-slot-append path in BOTH causal and non-causal form:
+            // the standard additive causal-mask builder (`Where(And(mask, causal),
+            // 0, -inf)`) is frozen to physical capacity alongside the KV, so its
+            // last-row frontier is the true valid length regardless of the op's
+            // `is_causal` attribute (the causal frontier and the padding frontier
+            // coincide at the last query row). When `dev_len` is set the causal
+            // offset and score-loop extent are both derived from it (see
+            // `attention_row`), so no host length leaks into a captured causal
+            // decode. Every other path passes a null pointer and keeps the
+            // host-derived length, so eager/dense/GQA are bit-for-bit unchanged.
             let dev_length_eligible = has_past_key
                 && has_past_value
-                && !self.is_causal
                 && mask.kind != 0
                 && capacity_key
                 && capacity_value
@@ -2200,6 +2238,7 @@ impl StandardAttentionKernel {
             let split_cfg = attention_split_config(
                 q_seq == 1,
                 dev_length_eligible,
+                self.is_causal,
                 want_qk,
                 key_cap as u64,
                 total_rows_u,

crates/onnx-runtime-session/src/executor/geometry.rs

@@ -418,19 +418,19 @@ pub(super) fn kernel_input_uses_physical_capacity(node: &Node, input_index: usiz
     // attention mask (input 3) instead of the growing cache extent. This mirrors
     // the GQA treatment and is what lets the decode step bind the KV cache at a
     // fixed capacity so whole-step CUDA-graph capture stays shape-static. Gated
-    // to the mask-driven, non-causal form (a present mask input and no
-    // `is_causal` attribute): that path derives length from the mask frontier,
-    // so the cache extent is pure capacity. Causal-attribute or mask-less
-    // Attention still reads the cache extent as the valid length.
+    // only on a present mask input (input 3): the standard additive causal-mask
+    // builder (`Where(And(attention_mask, causal), 0, -inf)`) is frozen to
+    // physical capacity alongside the KV, and its last-row frontier is the true
+    // valid length in BOTH the causal and non-causal form — at the last query
+    // row the causal frontier and the padding frontier coincide. The CUDA
+    // `Attention` kernel derives both the score-loop extent AND the causal offset
+    // from that on-device length (`standard_attention.rs`, `dev_len`), so the
+    // cache extent is pure capacity regardless of the op's `is_causal` attribute.
+    // A mask-less `Attention` still reads the cache extent as the valid length.
     node.is_default_domain()
         && node.op_type == "Attention"
         && matches!(input_index, 4 | 5)
         && node.inputs.get(3).is_some_and(Option::is_some)
-        && node
-            .attr("is_causal")
-            .and_then(|attr| attr.as_int())
-            .unwrap_or(0)
-            == 0
         || (
             // `pkg.nxrt::IndexShare` mirrors the mask-driven Attention treatment.
             // Its capacity form emits the 3-output present that ALIASES the
@@ -486,21 +486,21 @@ pub(super) fn is_additive_mask_builder_op(node: &Node) -> bool {
 
 /// Whether `node`/`input_index` is the additive-mask input (input 3) of a
 /// capacity-form `Attention`: a default-domain `Attention` whose KV cache
-/// (inputs 4/5) is already bound at physical capacity — i.e. a present mask and
-/// no `is_causal` attribute, so it derives the valid length from the mask
-/// frontier (see [`kernel_input_uses_physical_capacity`]). Such a node is
-/// *designed* to consume a physical-width additive mask, so it is a valid leaf
-/// for the frozen-mask (padded-capacity) classification.
+/// (inputs 4/5) is already bound at physical capacity — i.e. a present mask so
+/// it derives the valid length from the mask frontier (see
+/// [`kernel_input_uses_physical_capacity`]), in either the causal or non-causal
+/// form. Such a node is *designed* to consume a physical-width additive mask, so
+/// it is a valid leaf for the frozen-mask (padded-capacity) classification.
 ///
 /// The KV cache inputs (`past_key` = input 4, `past_value` = input 5) must both
 /// actually exist: `kernel_input_uses_physical_capacity` gates only on the mask
-/// (input 3) presence and `is_causal == 0`, but the CUDA `Attention` kernel's
-/// fixed-capacity append contract requires both past caches bound at physical
-/// capacity (`standard_attention.rs`: `has_past_key`/`has_past_value` on inputs
-/// 4/5, and `fixed_capacity_append` compares their capacity to the mask key
-/// width). A masked, non-causal `Attention` with only q/k/v/mask and no KV
-/// binding is NOT a capacity-form leaf, so require both KV inputs here rather
-/// than blessing any masked non-causal `Attention` as a valid cone terminus.
+/// (input 3) presence, but the CUDA `Attention` kernel's fixed-capacity append
+/// contract requires both past caches bound at physical capacity
+/// (`standard_attention.rs`: `has_past_key`/`has_past_value` on inputs 4/5, and
+/// `fixed_capacity_append` compares their capacity to the mask key width). A
+/// masked `Attention` with only q/k/v/mask and no KV binding is NOT a
+/// capacity-form leaf, so require both KV inputs here rather than blessing any
+/// masked `Attention` as a valid cone terminus.
 pub(super) fn is_capacity_form_attention_mask_input(node: &Node, input_index: usize) -> bool {
     input_index == 3
         && node.is_default_domain()

crates/onnx-runtime-session/src/executor/tests.rs

@@ -1995,8 +1995,8 @@ fn only_capacity_aware_inputs_keep_physical_capacity() {
 }
 
 // A capacity-form default-domain `Attention`: mask at input 3, KV cache at
-// inputs 4/5, no `is_causal` attribute — so it derives the valid length from the
-// mask frontier and binds the KV cache at physical capacity.
+// inputs 4/5 — so it derives the valid length from the mask frontier and binds
+// the KV cache at physical capacity, in either the causal or non-causal form.
 fn capacity_form_attention(id: u32, q: ValueId, mask: ValueId, out: ValueId) -> Node {
     Node::new(
         NodeId(id),
@@ -2009,8 +2009,9 @@ fn capacity_form_attention(id: u32, q: ValueId, mask: ValueId, out: ValueId) ->
 #[test]
 fn capacity_form_attention_mask_input_classifier() {
     // The mask slot (input 3) of a capacity-form `Attention` is a valid frozen-mask
-    // leaf; input 3 of a causal-attribute `Attention` (which reads the cache extent
-    // as the valid length) is not, and neither is a non-mask slot.
+    // leaf in both the non-causal and causal form (the frozen additive mask carries
+    // the valid length on-device either way); a non-mask slot is not, and neither is
+    // a masked `Attention` that lacks the KV cache bindings.
     let q = ValueId(0);
     let capacity = capacity_form_attention(0, q, q, q);
     assert!(is_capacity_form_attention_mask_input(&capacity, 3));
@@ -2022,9 +2023,20 @@ fn capacity_form_attention_mask_input_classifier() {
         .attributes
         .insert("is_causal".into(), Attribute::Int(1));
     assert!(
-        !is_capacity_form_attention_mask_input(&causal, 3),
-        "an is_causal Attention reads the cache extent as valid length, not the mask frontier"
+        is_capacity_form_attention_mask_input(&causal, 3),
+        "a frozen causal additive mask carries the valid length at its last-row frontier, \
+         so the causal capacity-form Attention is a valid frozen-mask leaf"
     );
+
+    // A masked `Attention` with no past KV bindings (inputs 4/5 absent) is not a
+    // capacity-form leaf regardless of causality.
+    let mask_only = Node::new(
+        NodeId(2),
+        "Attention",
+        vec![Some(q), Some(q), Some(q), Some(q)],
+        vec![q],
+    );
+    assert!(!is_capacity_form_attention_mask_input(&mask_only, 3));
 }
 
 // Build the standard additive causal-mask builder cone feeding a capacity-form
@@ -2190,26 +2202,20 @@ fn glm_indexer_add_mask_keeps_logical_width() {
 #[test]
 fn mask_builder_without_capacity_attention_is_rejected() {
     use onnx_runtime_ir::static_shape;
-    // A mask cone that ends at an is_causal Attention (reads cache extent as valid
-    // length), or reaches no Attention at all, must not be classified padded-safe.
+    // A mask cone that reaches no `Attention` at all (only a Cast to a graph
+    // output) must not be classified padded-safe: there is no capacity-form
+    // consumer, so the mask keeps exposing its logical valid length.
     let mut graph = Graph::new();
     graph.opset_imports.insert(String::new(), 17);
     let sh = || static_shape([1]);
     let mask = graph.create_named_value("attention_mask", DataType::Int64, sh());
     graph.add_input(mask);
-    let q = graph.create_named_value("q", DataType::Float32, sh());
     let cast = graph.create_named_value("cast", DataType::Float32, sh());
     graph.insert_node(Node::new(NodeId(0), "Cast", vec![Some(mask)], vec![cast]));
-    let attn = graph.create_named_value("attn", DataType::Float32, sh());
-    let mut causal = capacity_form_attention(1, q, cast, attn);
-    causal
-        .attributes
-        .insert("is_causal".into(), Attribute::Int(1));
-    graph.insert_node(causal);
-    graph.add_output(attn);
+    graph.add_output(cast);
     assert!(
         !mask_binding_feeds_capacity_form_attention(&graph, mask),
-        "an is_causal terminal Attention is not a capacity-form mask consumer"
+        "a mask cone reaching no capacity-form Attention is not padded-safe"
     );
 }
 
@@ -6488,9 +6494,11 @@ fn gqa_fixed_capacity_kv_seq_symbol_is_pinned_and_admits_the_node() {
 // attention op does NOT read its cache as physical capacity — must NOT be pinned,
 // so its symbol stays disqualifying and the node stays eager.
 //
-// (1) A default-domain, CAUSAL `Attention` derives past length from the growing
-//     cache extent (not from a mask frontier), so its cache is not physical
-//     capacity: not pinned.
+// (1) A default-domain CAUSAL `Attention` with NO mask input (input 3 absent)
+//     derives past length from the growing cache extent (there is no mask
+//     frontier to read), so its cache is not physical capacity: not pinned.
+//     (A causal Attention WITH a frozen additive mask input IS a capacity form —
+//     covered by the classifier tests above.)
 // (2) `CompressedSparseAttention` has NO past-KV inputs (its records grow from
 //     total_sequence_length), so it cannot be a capacity form: not pinned.
 #[test]