Opt Merge Multi Launch Kernels

Purpose

When a deployment goes the kernel-first path in Phase 4/5 (model has new ops or stitching boundaries differ from llama), build new model-specific fused ELFs by stitching multiple air.launch kernels into a single func.func. This skill is the procedural recipe.

The reference llama3 deployment went 10 launches/layer → 3 launches/layer using this recipe — the dominant prefill perf win.

Success criteria

A merge is "successful" when ALL hold:

1. Merged ELF compiles without errors
2. Output cosine ≥ 0.99 vs the unmerged baseline (run the same kernels as separate XRT calls, compare output). Match the BF16 tolerance convention from the kernel registry — rtol=1e-3 is too tight for K=2048 BF16; use cosine + max_abs/max_rel informational logging (see phase-1-kernel-validation PASS criteria for the convention).
3. Wall-clock time is lower than the unmerged baseline

If (1) fails → invoke debug-multi-launch-merge. If (2) fails → bisect by un-merging the last-added kernel. If (3) fails (compile slower than unmerged) → not a correctness bug; document and accept or revert.

Knowledge base references

• programming_examples/llms/llama_kernel_builder/stitching.py — the helpers this recipe uses (_rename_all, _fix_launch_func_args, _wrap_ir_in_launch, _rename_all_with_externs, _rename_all_gemv)
• programming_examples/llms/llama32_1b/multi_launch_builder/rms_gemms_rope_multi.py — reference 6-launch prefill merge (RMSNorm + Q/K/V GEMM + RoPE Q/K)
• programming_examples/llms/llama32_1b/multi_launch_builder/o_ffn_multi.py — reference 8-launch prefill merge (O + add + RMSNorm + Gate/Up + SwiGLU + Down + add)
• programming_examples/llms/llama32_1b/multi_launch_builder/o_gemv_ffn_multi.py — decode merge with 2-K extern rename (the pattern to extend to 3-K when n_heads*head_dim != emb_dim)
• programming_examples/kernel_registry/supported_kernels.md — per-kernel constraints; the FA section notes FlashAttention stays a separate XRT call (does NOT merge into the fused block)

Workflow

Step 1: Identify merge candidates

List the per-layer kernel sequence. Mark each as one of:

• Mergeable — pure compute (RMSNorm, GEMM, GEMV, RoPE, eltwise)
• Hard-stop — has data-dependent control flow or unsupported merge pattern. FlashAttention is the canonical hard-stop (see full_block.md)
• Conditional — mergeable but check herd-shape compatibility with neighbors (e.g., GEMM uses 8×4 herd; eltwise uses 8×1; coexisting in one segment is OK but verify)

Step 2: Pick merge boundaries

Group consecutive mergeable launches between hard-stops. For a typical decoder-only LLM:

• Group A: RMSNorm + Q/K/V proj + RoPE Q + RoPE K (6 launches — rms_gemms_rope_multi.py shape)
• Hard-stop: FlashAttention (separate XRT call)
• Group B: O proj + Add + RMSNorm + Gate + Up + SiLU+Mul + Down + Add (8 launches — o_ffn_multi.py shape)

For decode, replace GEMM with GEMV throughout.

Step 3: Author the multi-launch builder

For each group, create <model>/multi_launch_builder/<group_name>_multi.py that:

1. Builds each sub-kernel's IR via @module_builder
2. Imports stitching helpers: from llama_kernel_builder.stitching import (_rename_all, _fix_launch_func_args, _wrap_ir_in_launch, ...) (resolved against the shared programming_examples/llms/llama_kernel_builder/ via sys.path)
3. For each sub-kernel: extract its function body via _extract_between_func_and_return(ir_text)
4. Rename SSA values with a per-kernel prefix via _rename_all(body, prefix=...) to avoid collisions across the merged module
5. Remap function arguments to the merged module's args via _fix_launch_func_args(body, prefix, arg_map)
6. Concatenate the renamed bodies into a single func.func. If any sub-kernel emits a bare air.herd (RMSNorm herd_x>1, Eltwise Add at herd_x>1), wrap each via _wrap_ir_in_launch(...) BEFORE stitching — otherwise the lowering's airrt-to-npu pass drops the bare herd
7. For multi-K GEMV in one ELF (decode kernel-first): use _rename_all_with_externs with per-launch extern allowlists to keep different mv_*.o symbols distinct (extend the 2-K rename in o_gemv_ffn_multi.py to 3-K when n_heads*head_dim != emb_dim)

The canonical reference is rms_gemms_rope_multi.py — copy from there.