Pi0.5 on Jetson AGX Thor (SM110, 20 SMs, 225 GB/s DRAM, 32 MB L2) Baseline: mlir-tensorrt + Myelin compiler engine v1.6-L2 = 70.2 ms FlashRT: Hand-written C++/CUDA pipeline = 44 ms (CUDA Graph) Improvement: -26.2 ms (37.3% faster)

FlashRT fuses memory-bound ops that the compiler treats as separate kernels. On Thor, all non-GEMM ops are bandwidth-bound (D=2048 or 1024, 20 SMs). Each standalone kernel incurs DRAM read + write + launch overhead.

Total kernel launches eliminated: Compiler engine has ~5,300 layers / ~21,000 kernel launches per inference. FlashRT has ~13 kernels/decoder-layer × 180 blocks + ~500 encoder+siglip = ~2,840 kernel launches. Reduction: ~18,000 launches eliminated (85%).

The compiler engine uses dynamic FP8 quantization at every GEMM boundary:

Compiler (per GEMM):
  activation → amax_reduce → compute_scale → quantize_fp8 → GEMM → dequantize
  = 4 extra kernels + 1 device-sync (amax) per GEMM

FlashRT (per GEMM):
  activation (already FP8 from fused preceding kernel) → GEMM(descale=precomputed)
  = 0 extra kernels, scale is a compile-time constant

Calibration: Run 1 forward pass with amax measurement per quantization point, store scales. All subsequent inference uses static scales fused into preceding kernels.

Impact: Eliminates ~630 standalone quantize/dequantize kernel launches across the full pipeline.

Standard RMSNorm: y = x / rms(x) * (1 + weight). The (1 + weight) is constant per layer.

FlashRT: Fuse (1 + weight) into the following GEMM weight at load time:

# At weight load (once):
fused_qkv_weight = qkv_weight * (1 + rms_norm_weight)

# At inference:
x_norm = rms_norm_no_weight(x)  # simpler kernel, no weight multiply
qkv = x_norm @ fused_qkv_weight  # weight already absorbs norm scale

Applied to: Encoder QKV (18L), Encoder GateUp (18L), Pi0 Decoder QKV (18L), Pi0 Decoder GateUp (18L).

Impact: Eliminates per-element multiply in RMSNorm kernel. The rms_norm_fp8_noweight kernel is ~15% faster than weighted version.

Pi0.5 decoder uses AdaRMSNorm with learned modulation: y = adarms(x, style) where style = time_emb @ mod_weight + mod_bias.

Compiler (v1.0-v1.4): 370 Dense GEMMs computed inside the engine every inference.

FlashRT + Compiler v1.6-L2: Precompute all style vectors on CPU at prompt-set time. Inject as constants → 370 Dense GEMMs removed from GPU graph.

Impact: -5.5 ms (compiler v1.4→v1.6-L2), replicated in FlashRT as Python precompute in set_prompt().

The entire pipeline (SigLIP + PostLN + Encoder + Decoder) is captured as CUDA Graph.

Compiler engine: Already uses CUDA Graph internally (TRT engine execution). However, the graph contains 21,000+ nodes with Myelin's fine-grained kernel scheduling.

FlashRT: Graph contains ~2,840 nodes (7.4x fewer). Fewer nodes = faster graph instantiation + less scheduling overhead.

Graph autotune: CUDA Graph instantiation on Thor is non-deterministic (±2ms variance). FlashRT recaptures until a fast schedule is found. Typical: 0-1 retries needed.

Note: Compiler's "74 ms nsys" vs "70.2 ms CUDA events" gap is due to nsys profiling overhead. FlashRT's 44 ms is CUDA events measurement (wall-clock perf_counter with cudaStreamSynchronize).

The compiler engine benefits from Myelin's global fusion across SigLIP's 27 layers:

• QKV GEMM + bias + reshape fused into single kernel
• LayerNorm + GELU + bias fused as GEMM epilogues
• Entire 27-layer stack shares L2 cache optimally (Myelin schedules for L2 residency)

FlashRT uses cuBLASLt FP8 GEMM + separate custom kernels for norm/activation:

• 27 layers × ~6 kernels/layer = ~162 kernel launches (vs compiler's ~138 with deeper fusion)
• FP8 weight reduces DRAM traffic (15 MB/layer < 32 MB L2), partially compensating

This is the one area where the compiler approach wins. SigLIP's large batch (S=512) and simple attention (no RoPE, no cross-attention) is ideal for Myelin's fusion.

Key wins:

1. No data movement: Compiler inserts MoveConcat/Transpose for layout conversion between ops. FlashRT uses consistent row-major layout throughout — zero explicit transposes.
2. Fused norm+FP8: Compiler has separate RMSNorm → QDQ → Quantize → GEMM. FlashRT: rms_norm_fp8_static produces FP8 output directly.
3. cuBLASLt vs Myelin tactic: For Se=522 (encoder sequence), cuBLASLt selects better CUTLASS configurations than Myelin's tactic search for some shapes.

Key wins:

1. 10.1 ms from fusion: Compiler's ~9,700 kgen launches (norm, GELU, softmax, QDQ) across 180 blocks. FlashRT replaces with 3 fused kernels per layer: fused_adarms_fp8_static + gate_res_adarms_fp8_static + gate_geglu_merged_fp8_fp16. Despite GPU overlap making kgen "free" in the compiler, FlashRT's fused approach reduces total graph nodes and improves scheduling.
2. 8.4 ms from GEMM: At S=10 (decoder batch), cuBLASLt FP8 with pre-baked descale outperforms Myelin's CUTLASS FP8 + separate descale. Myelin's per-GEMM tactic search advantage diminishes at very small batch sizes where launch overhead dominates.

#   Kernel                              Time     Type
1.  AdaRMSNorm kgen                     6 μs     elementwise (Myelin fused)
2.  QKV proj FP8 GEMM (CUTLASS)        17 μs    compute
3.  RoPE kgen                           5 μs     elementwise
4.  RoPE kgen                           5 μs     elementwise
5.  Q@K^T GEMM                          4 μs     compute
6.  Transpose+Select kgen               3 μs     data movement
7.  Softmax kgen                        4 μs     elementwise
8.  Attn@V GEMM                         4 μs     compute
9.  QDQ kgen                            3 μs     quantize
10. Transpose kgen                      3 μs     data movement
11. O proj FP8 GEMM                    12 μs    compute
12. Gate mul + residual kgen            5 μs     elementwise
13. AdaRMSNorm kgen                     6 μs     elementwise
14. Gate+Up FP8 GEMM (Wab, CUTLASS)   17 μs    compute
15. GELU kgen                           8 μs     elementwise
16. Down FP8 GEMM (CUTLASS)           17 μs    compute
17. Gate mul + residual kgen            5 μs     elementwise
18-22. QDQ/Cast kernels (various)      15 μs    quantize/cast
────────────────────────────────
     ~22 kernels/layer                ~139 μs   total
     ×180 layers                      ~25 ms    (+ GPU overlap → 36ms wall)

#   Kernel                              Time     Type
1.  fused_adarms_fp8_static             8 μs     AdaRMSNorm+FP8 quantize (fused)
2.  fp8_gemm_descale (QKV)             14 μs    cuBLASLt FP8 + static descale
3.  qkv_split_rope_kvcache             5 μs     split Q/K/V + RoPE + KV cache write
4.  cublasGemmEx (Q@K^T)               4 μs     cuBLAS attention
5.  softmax_fp16                        3 μs     vectorized softmax
6.  cublasGemmEx (attn@V)              4 μs     cuBLAS attention
7.  quantize_fp8_static                 2 μs     O proj input quantize
8.  fp8_gemm_descale (O proj)          12 μs    cuBLASLt FP8
9.  gate_res_adarms_fp8_static          9 μs     gate×res + AdaRMSNorm + FP8 (fused)
10. fp8_gemm_descale (Gate+Up)         14 μs    cuBLASLt FP8
11. gate_geglu_merged_fp8_fp16          4 μs     GELU(gate)×up + FP8 (fused)
12. fp8_gemm_descale (Down)            14 μs    cuBLASLt FP8
13. gate_res_adarms_fp8_static          9 μs     residual + cross-layer norm (fused)
────────────────────────────────
     13 kernels/layer                  ~102 μs   per layer
     ×180 layers                       ~18.4 ms  (+ dispatch → 24.5ms wall)

Summary of FlashRT's advantage over compiler v1.6-L2 (70→44 ms):

Any TRT Plugin inserted into the Myelin engine creates an "opaque boundary" — Myelin cannot fuse across it. Tested with RoPE Plugin:

• Plugin alone: -18 ms (RoPE kernel 10x faster)
• Myelin fusion regression: +36 ms (surrounding GEMMs lose cooperative scheduling)
• Net effect: +10 ms slower

The decoder operates at S=10 (10 action tokens). At this batch size:

• All GEMMs are launch-bound, not compute-bound
• Kernel fusion savings (memory traffic) are negligible (< 1 KB intermediate data)
• Dispatch overhead dominates: compiler's 16,480 launches × ~2μs = ~33 ms dispatch cost
• FlashRT's 2,340 launches × ~3μs = ~7 ms dispatch cost
• Dispatch savings alone: ~26 ms (matches the total improvement)

For large shapes (encoder Se=522), cuBLASLt consistently selects better CUTLASS configurations than Myelin's tactic search. This is because:

1. cuBLASLt has access to NVIDIA's latest tactic database (updated with driver)
2. Myelin's search space is constrained by the engine builder's time budget
3. For very small shapes (S=10 decoder), both are launch-bound and the difference is minimal

Compiler (Myelin):  Global fusion view → fewer kernels → but limited to Myelin's patterns
FlashRT:           No global fusion   → more kernels → but each kernel is purpose-built

At large batch (SigLIP S=512):  Myelin's global fusion wins (1 ms vs 6 ms)
At small batch (Decoder S=10):  FlashRT's fewer launches wins (24 ms vs 36 ms)
At medium batch (Encoder S=522): FlashRT's cuBLASLt + fusion wins (20 ms vs 37 ms)