[Perf] Fuse the BF16 attention output gate in short gfx1201 decode graphs #780

Performance opportunity

FreeToken's gated full-attention path launches sigmoid and multiply kernels after decode attention. For short BF16 decode Graphs these can be fused into the existing stage2 while retaining the attention arithmetic and BF16 rounding.

Proposed scope

Automatic dispatch only for gfx1201, BF16 B1/Hq16/Hkv2/D256, CSR capacity 512, default full attention. Eager, prefill, larger capacities and other devices/shapes keep Torch gating. No cross-platform 1.4x claim is made.

Measured on Radeon AI PRO R9700

Base: cc4aa085b6eaeda01c0441fb3ecb9cc843b3e465. gfx1201, 32 GiB;
AMD EPYC 9334; Linux 7.0.0-38-generic, amdgpu 7.1.9.31600000;
Torch 2.14.0+rocm10.1.0, HIP 7.16.26385, Triton 3.8.0, Python 3.13.
Only this GPU platform was measured. These figures do not establish a speedup on
other architectures or checkpoints.

Kernel window

Two independent rounds; seven alternating samples per round; each Graph contains
four calls, with 50 replays per sample after warmup. Unit: microseconds per call.
The reference loads the upstream kernel module and applies Torch sigmoid/multiply.
Both B1 and B4 outputs match the reference elementwise exactly in these cases.
B4 lengths are 512/495/478/461; B4 is a low-level kernel result, not automatic
serving admission.

End to end: gate fusion alone

Checkpoint: Qwen/Qwen3.6-35B-A3B-FP8 (locally converted to FTW), BF16
activations/KV, FP8-block experts. The command above uses a 20 GiB offload expert
cache. Convolution and GDN are unchanged on both arms. Three fresh-process pairs
per capacity, one warmup then one measured request per process, alternating order
main/branch, branch/main, main/branch. CPU affinity 16-23,80-87 (NUMA2), GPU NUMA0;
the host is shared and clocks are not locked.

Throughput and TTFT columns are independent medians. Paired gains are computed
within each round then summarized, so they need not equal ratios of those medians.
All six pairs have equal output text hashes and exactly 128 generated tokens.
The short Graph has a small decode improvement; these measurements do not establish
a reliable request-wide gain. The 4096-capacity control retains the Torch gate.

Correctness and actual execution

• Attention/backend/import regression tests: 128 passed. They include the Torch
  attention oracle, exact fused-vs-unfused comparison, sinks/window and non-power
  GQA groups, changed Graph replay inputs, fallback, opt-out and Hybrid delegation.
• Runtime cache rebuild and Graph re-capture on Qwen/Qwen3-1.7B: one test passed, including growth, rejected oversized targets, and continued generation.
• Per capacity: four unforced QA requests plus 64 matched-prefix logit positions
  (four requests x sixteen positions). All sampled logits are bitwise equal to
  main and finite; max KL is zero; no new first-answer failures.
• Short Graph capture records ten gated attention calls. Its diagnostic first
  replay has 1157 -> 1137 GPU kernels, removing two gate kernels per full-attention
  layer. At capacity 4096, capture records zero fused calls and both traces have
  1157 kernels. These are diagnostic first-replay counts, not a steady-state
  time breakdown.

Code and reproduction: https://github.com/zihaomu/FreeToken/tree/perf/triton-decode-output-gate