perf(cuda-ep): collapse bf16 SiLU/SwiGLU-mul glue nodes — realized +0.9% byte-exact (#899)

.squad/decisions/inbox/batty-glue-node-collapse.md

@@ -0,0 +1,43 @@
+### 2026-08-13: Graph-side glue node-collapse REALIZED — bf16 SiLU/SwiGLU-mul collapse ships (+0.9%, byte-exact)
+**By:** Batty
+**What:** Converted the §8 "+5.3% ceiling" GO (#899) into a realized, measured number on the
+production Muse-Glimmer-30B native CUDA decode graph. The single graph-side change that is
+byte-exactly collapsible on this model with a landed kernel: extend `CudaSiluFusion`
+(`crates/onnx-runtime-ep-cuda/src/optimizer.rs`) to accept **BFloat16** (was `Float16`-only). The
+standalone `Sigmoid(x)` + `Mul(x, sigmoid)` + `Mul(silu, up)` glue then collapses through
+`CudaSiluFusion` → `CudaSwiGluFusion` into the tagged decomposed `Mul[_cuda_silu_mul]`, which the
+runtime lowers to the already-landed **`decomposed_silu_mul_bf16`** epilogue (#867). Deletes 104
+standalone nodes (2/layer × 52): `Sigmoid` 104→52, `Mul` 210→158, total decoder 2458→2354.
+`CudaGateUpSwiGluFusion` needs an fp16 activation so it stays dormant → **int4 GEMVs untouched**
+(per §7/#898). Two unit tests added (`collapses_bf16_decomposed_swiglu_glue`,
+`leaves_fp32_decomposed_swiglu_glue_separate`).
+
+**Measured (H200, `CUDA_VISIBLE_DEVICES=0`, `ONNX_GENAI_CUDA_GRAPH=1`, `--steady --warmups 1
+--runs 3 --tokens 128`, 3 interleaved A/B rounds, same release binary):**
+- decode: 21.19 → 20.99 ms/token; **throughput 47.20 → 47.63 tok/s = +0.9%**.
+- node count/layer: ~22 glue → ~20 (−2/layer; −104 total).
+- greedy 24-token stream **byte/token-identical** before↔after
+  (`[24, 372, 1045, 10016, …, 1740, 2885]`).
+- CUDA EP cdylib (`libonnx_runtime_ep_cuda_plugin.so`) built; `ldd` confirms it does **not** link
+  `libonnxruntime`. `cargo fmt --check` + `cargo clippy` clean; 59 optimizer unit tests pass.
+
+**Why:** On the real bf16 graph the elementwise/norm fusions were **dormant** because every
+SiLU/SwiGLU pass was gated to `Float16`; Muse-Glimmer's stream is bf16. This change simply
+activates the byte-exact bf16 kernel Sebastian already landed. Zero-risk: no cooperative kernel, no
+`grid.sync`, no reduction reorder → no Chew gate.
+
+**Verdict: SHIP** (small-but-real, byte-exact). Honest bound vs the +5.3% ceiling: only the 2
+SiLU/SwiGLU-mul nodes/layer are byte-exactly collapsible here. The larger glue is blocked —
+(1) **6 norms/layer**: Gemma3 sandwich-norm residual is `x + norm(y)`; #854's
+`SkipSimplifiedLayerNormalization` (`norm(x+skip)`) *does* apply across the layer seam, **but only
+f32/f16 skip kernels exist** — the f16 kernel rounds the residual sum to f16 before the RMS
+reduction, whereas the f32-template path bf16 would reach accumulates over the **unrounded** fp32
+sum → not byte-exact → **FLAGGED FOR CHEW/Sebastian: needs a bf16 skip-RMSNorm kernel that rounds
+the sum**, then #854's fold can fire byte-exactly and recover the residual-Add + norm glue (the
+bigger lever). (2) **~208 constant `gamma+1` Adds/layer**: #872 `CudaFoldConstantAdd` already
+MEASURED a −2.8% regression — do not re-attempt. (3) **4 reshapes/layer**: GQA head-split metadata
+coupled to the attention kernel, not free-standing deletions. Remaining glue interleaves with the
+dominant GEMV serial cost, so realized ≤ ceiling as §8.3 predicted.
+
+Doc: `docs/research/dense-decode-megakernel-feasibility.md` §8.5 (realized subsection). PR
+references #899. Not self-merged.

crates/onnx-runtime-ep-cuda/src/optimizer.rs

@@ -77,6 +77,13 @@ const GATE_UP_SWIGLU_SUPPORTED_BITS: i64 = 4;
 pub(crate) struct CudaSwiGluFusion;
 
 /// Lower an exact `x * Sigmoid(x)` pair to the CUDA EP's fused SiLU kernel.
+///
+/// Fires on fp16 **and** bf16 activations: the fused SwiGLU epilogue has a
+/// byte-exact decomposed kernel for each (`decomposed_silu_mul_{f16,bf16}`),
+/// so collapsing the standalone `Sigmoid` + `Mul(x, sigmoid)` glue keeps greedy
+/// tokens byte-identical. bf16 decoders (e.g. Muse-Glimmer-30B) export SwiGLU as
+/// that standalone pair, so bf16 support is what lets the per-layer glue node
+/// collapse into the already-landed fused launch.
 #[derive(Clone, Copy, Debug, Default)]
 pub(crate) struct CudaSiluFusion;
 
@@ -130,7 +137,16 @@ impl OptimizationPass for CudaSiluFusion {
             let Some(x) = sigmoid.inputs[0] else {
                 continue;
             };
-            if graph.value(x).dtype != DataType::Float16 {
+            // The fused SwiGLU epilogue has a byte-exact decomposed kernel for
+            // both fp16 (`decomposed_silu_mul_f16`) and bf16
+            // (`decomposed_silu_mul_bf16`): each rounds the intermediate
+            // `sigmoid` and the `x * sigmoid` product back to the narrow dtype
+            // exactly as the standalone `Sigmoid`/`Mul` ops do, so greedy tokens
+            // stay byte-identical. bf16 decoders such as Muse-Glimmer-30B export
+            // the SwiGLU as a standalone `Sigmoid` + `Mul(x, sigmoid)` pair in
+            // their native bf16 stream; without bf16 here the ~2 glue nodes per
+            // layer never collapse into the fused launch.
+            if !matches!(graph.value(x).dtype, DataType::Float16 | DataType::BFloat16) {
                 continue;
             }
             let sigmoid_output = sigmoid.outputs[0];
@@ -2711,6 +2727,81 @@ mod tests {
         );
     }
 
+    /// `Sigmoid(x)` + `Mul(x, sigmoid)` + `Mul(silu, up)` in a narrow stream
+    /// collapses through `CudaSiluFusion` -> `CudaSwiGluFusion` into a single
+    /// tagged decomposed `Mul[_cuda_silu_mul]`, deleting the standalone
+    /// `Sigmoid` and the intermediate `Mul`. The plain (non-projection) `up`
+    /// value means `CudaGateUpSwiGluFusion` cannot fold the GEMVs, so the glue
+    /// collapse is observed in isolation.
+    fn decomposed_swiglu_glue_graph(dtype: DataType) -> Graph {
+        let mut graph = Graph::new();
+        graph.opset_imports.insert(String::new(), 17);
+        let x = value(&mut graph, "x", dtype, 7);
+        let up = value(&mut graph, "up", dtype, 7);
+        let sigmoid_out = value(&mut graph, "sigmoid", dtype, 7);
+        let silu_out = value(&mut graph, "silu", dtype, 7);
+        let output = value(&mut graph, "output", dtype, 7);
+        graph.add_input(x);
+        graph.add_input(up);
+        graph.insert_node(Node::new(
+            NodeId(0),
+            "Sigmoid",
+            vec![Some(x)],
+            vec![sigmoid_out],
+        ));
+        graph.insert_node(Node::new(
+            NodeId(0),
+            "Mul",
+            vec![Some(x), Some(sigmoid_out)],
+            vec![silu_out],
+        ));
+        graph.insert_node(Node::new(
+            NodeId(0),
+            "Mul",
+            vec![Some(silu_out), Some(up)],
+            vec![output],
+        ));
+        graph.add_output(output);
+        graph
+    }
+
+    #[test]
+    fn collapses_bf16_decomposed_swiglu_glue() {
+        let mut graph = decomposed_swiglu_glue_graph(DataType::BFloat16);
+        CudaSiluFusion.run(&mut graph, &PassContext::new()).unwrap();
+        CudaSwiGluFusion
+            .run(&mut graph, &PassContext::new())
+            .unwrap();
+
+        assert_eq!(graph.num_nodes(), 1, "Sigmoid + inner Mul must be removed");
+        let fused = graph.nodes.values().next().unwrap();
+        assert_eq!(fused.op_type, "Mul");
+        assert_eq!(
+            fused.attr(SILU_MUL_FUSION_ATTR).and_then(Attribute::as_int),
+            Some(1)
+        );
+        assert_eq!(
+            fused.attr(DECOMPOSED_SILU_ATTR).and_then(Attribute::as_int),
+            Some(1),
+            "the bf16 fused Mul must carry the decomposed marker so the runtime \
+             selects the byte-exact decomposed_silu_mul_bf16 kernel"
+        );
+        assert!(graph.validate().is_ok());
+    }
+
+    #[test]
+    fn leaves_fp32_decomposed_swiglu_glue_separate() {
+        let mut graph = decomposed_swiglu_glue_graph(DataType::Float32);
+        CudaSiluFusion.run(&mut graph, &PassContext::new()).unwrap();
+
+        // No fp32 decomposed-SiLU kernel: the Sigmoid must not lower to Silu.
+        assert_eq!(graph.num_nodes(), 3);
+        assert!(
+            graph.nodes.values().all(|node| node.op_type != "Silu"),
+            "fp32 has no byte-exact decomposed SiLU kernel; must stay unfused"
+        );
+    }
+
     // === QKV bias fold ===
 
     use onnx_runtime_ir::{TensorData, WeightRef};

docs/research/dense-decode-megakernel-feasibility.md

@@ -476,6 +476,74 @@ tax and no numerics gate.
   (not the modeled chain) to convert the +5.3% ceiling into a realized number, once Batty
   has a candidate collapse in `optimizer.rs`.
 
+### 8.5 REALIZED on the production graph — bf16 SiLU/SwiGLU-mul collapse (measured, this H200)
+
+The §8.4 ceiling (+5.3%) is a bound on collapsing *all* ~22 glue nodes/layer. Auditing the
+**real** Muse-Glimmer-30B decoder graph against the landed fused epilogues shows most of that
+glue is **not** byte-exactly collapsible here — so the realized number is bounded well below
+the ceiling. What the real graph actually contains, per 52-layer decoder (runtime op histogram,
+`ONNX_GENAI_PROFILE_OPS=1`, eager, one forward):
+
+| op_type | count | /layer | collapsible byte-exact? |
+|---|---|---|---|
+| `MatMulNBits` | 417 | 8 | no — irreducible int4 GEMV (§7, #898); out of scope |
+| `GroupQueryAttention` | 52 | 1 | no — GQA-decode; out of scope |
+| `SimplifiedLayerNormalization` | 312 | 6 | **no for bf16** — see below |
+| `Add` | 311 | 6 | ~208 are constant `gamma+1` folds (#872 `CudaFoldConstantAdd` REGRESSED −2.8%, not shipped); ~104 are residual adds |
+| `Mul` | 210 | 4 | 52 are the SwiGLU multiply → **collapsed here** |
+| `Reshape` | 208 | 4 | structural/kernel-coupled (GQA head-split); not a standalone deletion |
+| `Sigmoid` | 104 | 2 | 52 are the MLP SiLU gate → **collapsed here**; 52 are the attention-output gate `sigmoid(gate)*attn` (not self-gated SiLU) |
+| `Cast` | 2 | ~0 | already folded 834 → 2 by `CudaDropNormalizationCasts` + `CudaFoldConstantCast` |
+
+**Why the elementwise/norm fusions were dormant on this model:** every SiLU/SwiGLU fusion was
+gated to `Float16`, but Muse-Glimmer's activation stream is **bf16**, so `CudaSiluFusion` never
+fired. The architecture is Gemma3-style *sandwich* norm — the residual `Add` comes **after** the
+post-norm (`x + norm(y)`), so #854's `SkipSimplifiedLayerNormalization` (which computes
+`norm(x + skip)`) does apply across the layer seam, **but only f32/f16 skip kernels exist**: the
+f16 skip kernel rounds the residual sum to f16 *before* the RMS reduction (byte-exact), whereas
+the f32-template path (the only one bf16 could reach) accumulates over the **unrounded** fp32 sum.
+That changes the reduction vs the standalone bf16 `Add`→norm → **not byte-exact → flagged for
+Chew/Sebastian (needs a bf16 skip kernel that rounds the sum), NOT implemented here.**
+
+**What shipped (byte-exact, reuses #867):** extend `CudaSiluFusion` to accept `BFloat16`. The
+standalone `Sigmoid(x)` + `Mul(x, sigmoid)` + `Mul(silu, up)` chain then collapses through
+`CudaSiluFusion` → `CudaSwiGluFusion` into the tagged decomposed `Mul[_cuda_silu_mul]`, which the
+runtime lowers to the already-landed **`decomposed_silu_mul_bf16`** kernel. That kernel is
+documented (and here confirmed) to reproduce the standalone per-op bf16 rounding exactly — sigmoid
+and the products each rounded via `__float2bfloat16_rn`. `CudaGateUpSwiGluFusion` requires an fp16
+activation, so it stays dormant on bf16: **the int4 GEMVs are left untouched** (per §7 / #898), and
+only the glue nodes collapse.
+
+**Measured (H200, `CUDA_VISIBLE_DEVICES=0`, `ONNX_GENAI_CUDA_GRAPH=1`, `--steady --warmups 1
+--runs 3 --tokens 128`, 3 interleaved A/B rounds on the same release binary):**
+
+| metric | before | after |
+|---|---|---|
+| decode | 21.19 ms/token | 20.99 ms/token |
+| throughput | **47.20 tok/s** (mean; 47.18/47.24/47.19) | **47.63 tok/s** (mean; 47.55/47.74/47.60) |
+| glue nodes/layer | ~22 | ~20 |
+| `Sigmoid` (per forward) | 104 | **52** |
+| `Mul` (per forward) | 210 | **158** |
+| total decoder nodes | 2458 | **2354** (−104; −2/layer) |
+| greedy tokens (24-tok, deterministic) | `[24, 372, 1045, 10016, …, 1740, 2885]` | **identical (byte/token-exact)** |
+
+**Realized decode delta: 47.20 → 47.63 tok/s = +0.9% (−0.19 ms/token), byte-exact.**
+
+**Interpretation — a small-but-real SHIP, far below the +5.3% ceiling, honestly bounded by what
+is collapsible.** The +5.3% assumed all ~22 glue nodes/layer collapse; on the real bf16 graph only
+the **2** SiLU/SwiGLU-mul nodes/layer are byte-exactly collapsible with a landed kernel. The
+larger glue populations are each blocked for a concrete reason: the 6 norms/layer have no
+byte-exact bf16 skip kernel (reduction-rounding mismatch → Chew), the 4 constant `gamma+1`
+adds/layer are a **measured regression** (#872), and the 4 reshapes/layer are GQA head-split
+metadata coupled to the attention kernel, not free-standing deletions. Removing 104 dispatch-bound
+nodes recovers ~0.19 ms/token — consistent with the §8.1 ~0.9 µs/node replay floor plus the small
+real memory traffic of the 52 removed `Sigmoid` launches over the 19968-wide intermediate. The
+glue that remains interleaves with the dominant GEMV serial cost (§8.3 caveat), so realized ≤
+ceiling as predicted. **Verdict: SHIP** — zero-risk, byte-exact, no GEMV touch, no numerics
+reorder, no Chew gate; it simply activates the landed #867 bf16 kernel that the `Float16`-only pass
+gate was hiding. Further glue collapse on this model requires a **bf16 skip-RMSNorm kernel that
+rounds the residual sum** (Sebastian/Chew) before #854's fold can fire byte-exactly on bf16.
+
 ---
 
 - Env: `source /home/justinchu/onnx-genai/.cudaenv.sh`; `CUDA_VISIBLE_DEVICES=0 ONNX_GENAI_CUDA_DEVICE=0`.