perf(attention): fuse short bf16 decode output gates on gfx1201

benchmarks/bench_decode_output_gate.py

@@ -0,0 +1,158 @@
+"""Paired decode+gate kernel timings, or a warmed request to a running FT server.
+
+Kernel mode: PYTHONPATH=python python benchmarks/bench_decode_output_gate.py --json kernel.json
+Server mode: ... --url http://127.0.0.1:8000 --model /path/to/tokenizer --json server.json
+"""
+from __future__ import annotations
+
+import argparse
+import hashlib
+import importlib.util
+import json
+import statistics
+import sys
+import time
+import urllib.request
+from pathlib import Path
+
+
+def server(args):
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(args.model, local_files_only=True)
+    chat = tokenizer.apply_chat_template(
+        [{"role": "user", "content": "List the positive integers in order, starting from 1, separated by commas."}],
+        tokenize=False, add_generation_prompt=True, enable_thinking=False,
+    )
+    ids = tokenizer.encode(chat, add_special_tokens=False)
+    filler = tokenizer.encode(" hello", add_special_tokens=False)
+    assert len(filler) == 1 and len(ids) <= args.input_tokens
+    prompt = tokenizer.decode(filler * (args.input_tokens - len(ids)) + ids, skip_special_tokens=False)
+    assert len(tokenizer.encode(prompt, add_special_tokens=False)) == args.input_tokens
+    body = dict(model=args.served_model_name, prompt=prompt, max_tokens=args.output_tokens, temperature=0,
+                stream=True, ignore_eos=True, stream_options={"include_usage": True})
+    rows = []
+    for iteration in range(args.repeats + 1):
+        request = urllib.request.Request(args.url.rstrip("/") + "/v1/completions",
+                                        data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
+        started = time.perf_counter()
+        events, text, usage, finished = [], "", None, None
+        with urllib.request.urlopen(request, timeout=600) as response:
+            for raw in response:
+                if not raw.startswith(b"data: "):
+                    continue
+                payload = raw[6:].strip()
+                if payload == b"[DONE]":
+                    break
+                value = json.loads(payload)
+                elapsed = time.perf_counter() - started
+                content = "".join(c.get("text", "") for c in value.get("choices", []))
+                if content:
+                    text += content
+                    events.append(elapsed)
+                if any(c.get("finish_reason") for c in value.get("choices", [])):
+                    finished = elapsed
+                if value.get("usage"):
+                    usage = value["usage"]
+        elapsed = time.perf_counter() - started
+        assert usage and usage["completion_tokens"] == args.output_tokens and events and finished
+        rows.append(dict(warmup=iteration == 0, usage=usage, ttft_s=events[0], elapsed_s=elapsed,
+                         decode_tok_s=(args.output_tokens - 1) / (finished - events[0]),
+                         e2e_tok_s=args.output_tokens / elapsed, nonempty_events=len(events),
+                         events_s=events, finished_s=finished,
+                         output_sha256=hashlib.sha256(text.encode()).hexdigest()))
+    return dict(mode="server", request=body, rows=rows)
+
+
+def kernel(args):
+    import torch
+    from freetoken.kernel.triton.attention import (
+        _decode_gate_sigmoid_table, decode_kv_split_cap, decode_paged_attention,
+    )
+
+    reference = decode_paged_attention
+    if args.baseline_path:
+        spec = importlib.util.spec_from_file_location("gate_bench_upstream", args.baseline_path)
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[spec.name] = module
+        spec.loader.exec_module(module)
+        reference = module.decode_paged_attention
+    torch.manual_seed(20261021)
+    torch.set_num_threads(4)
+    _decode_gate_sigmoid_table(torch.cuda.current_device())
+    rows = []
+    for batch in (1, 4):
+        lengths = [512 - i * 17 for i in range(batch)]
+        total = sum(lengths)
+        q = torch.randn(batch, 16, 256, device="cuda", dtype=torch.bfloat16) * 0.2
+        k = torch.randn(total + 7, 2, 256, device="cuda", dtype=torch.bfloat16) * 0.2
+        v = torch.randn_like(k) * 0.2
+        gate = torch.randn_like(q)
+        indptr = torch.tensor([0, *lengths], device="cuda", dtype=torch.int32).cumsum(0).to(torch.int32)
+        indices = torch.randperm(total + 7, device="cuda").to(torch.int32)[:total]
+        positions = torch.tensor([n - 1 for n in lengths], device="cuda", dtype=torch.int32)
+        splits = decode_kv_split_cap(batch, 16, 2, torch.cuda.current_device())
+        partial = torch.empty(batch, 16, splits, 256, device="cuda")
+        lse = torch.empty(batch, 16, splits, device="cuda")
+        params = (q, k, v, indptr, indices, positions, partial, lse, 256 ** -0.5)
+        functions = [lambda: reference(*params) * gate.sigmoid(),
+                     lambda: decode_paged_attention(*params, output_gate=gate)]
+        expected, actual = functions[0](), functions[1]()
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+        graphs = []
+        for function in functions:
+            stream = torch.cuda.Stream()
+            stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(stream):
+                for _ in range(5):
+                    function()
+            stream.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                for _ in range(4):
+                    function()
+            graphs.append(graph)
+        for graph in graphs:
+            for _ in range(1000):
+                graph.replay()
+        torch.cuda.synchronize()
+        samples = [[], []]
+        for group in range(7):
+            for index in ((0, 1) if group % 2 == 0 else (1, 0)):
+                start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
+                start.record()
+                for _ in range(args.iterations):
+                    graphs[index].replay()
+                end.record()
+                end.synchronize()
+                samples[index].append(start.elapsed_time(end) * 1000 / (args.iterations * 4))
+        baseline_us, candidate_us = [statistics.median(s) for s in samples]
+        rows.append(dict(batch=batch, lengths=lengths, head_dim=256, query_heads=16, kv_heads=2,
+                         baseline_us=baseline_us, candidate_us=candidate_us,
+                         speedup=baseline_us / candidate_us, samples_us=samples, bitwise_equal=True))
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    return dict(mode="kernel", torch=str(torch.__version__), hip=torch.version.hip, gpu=props.name,
+                baseline_path=args.baseline_path, rows=rows)
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--url", help="running server; omit for kernel measurements")
+    parser.add_argument("--model", help="local tokenizer path, required with --url")
+    parser.add_argument("--served-model-name", default="gate-bench")
+    parser.add_argument("--input-tokens", type=int, default=128)
+    parser.add_argument("--output-tokens", type=int, default=128)
+    parser.add_argument("--repeats", type=int, default=1, help="measured requests after one warmup")
+    parser.add_argument("--iterations", type=int, default=50, help="Graph replays per kernel sample")
+    parser.add_argument("--baseline-path", help="upstream kernel/triton/attention.py for kernel A/B")
+    parser.add_argument("--json", type=Path, required=True)
+    args = parser.parse_args()
+    if args.url and not args.model:
+        parser.error("--model is required with --url")
+    result = server(args) if args.url else kernel(args)
+    args.json.write_text(json.dumps(result, indent=2) + "\n")
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    main()

python/freetoken/attention/base.py

@@ -66,6 +66,22 @@ def forward(
         attn_spec: AttentionSpec | None = None,
     ) -> torch.Tensor: ...
 
+    def forward_gated(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        gate: torch.Tensor,
+        attn_spec: AttentionSpec | None = None,
+    ) -> torch.Tensor:
+        """Attention with the model's elementwise sigmoid output gate."""
+        import torch
+
+        out = self.forward(q, k, v, layer_id, batch, attn_spec=attn_spec)
+        return out * torch.sigmoid(gate).reshape_as(out)
+
     @abstractmethod
     def prepare_metadata(self, batch: Batch) -> None: ...
 
@@ -112,6 +128,19 @@ def forward(
         backend = self.prefill_backend if batch.is_prefill else self.decode_backend
         return backend.forward(q, k, v, layer_id, batch, attn_spec=attn_spec)
 
+    def forward_gated(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        gate: torch.Tensor,
+        attn_spec: AttentionSpec | None = None,
+    ) -> torch.Tensor:
+        backend = self.prefill_backend if batch.is_prefill else self.decode_backend
+        return backend.forward_gated(q, k, v, layer_id, batch, gate, attn_spec=attn_spec)
+
     def prepare_metadata(self, batch: Batch) -> None:
         backend = self.prefill_backend if batch.is_prefill else self.decode_backend
         return backend.prepare_metadata(batch)

python/freetoken/attention/triton.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, List
 
@@ -81,6 +82,18 @@ def __init__(self, config: ModelConfig):
             default=int(getattr(config, "head_dim", 1)),
         )
         self.kv_heads = sorted({group.num_kv_heads for group in kv_groups}) or [int(getattr(config, "num_kv_heads", 1))]
+        self._output_gate_fusion = (
+            os.getenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", "1") != "0"
+            and self.device.type == "cuda"
+            and bool(torch.version.hip)
+            and self.num_q_heads == 16 and self.max_head_dim == 256 and 2 in self.kv_heads
+            and getattr(torch.cuda.get_device_properties(self.device), "gcnArchName", "").split(":")[0] == "gfx1201"
+        )
+        if self._output_gate_fusion:
+            from freetoken.kernel.triton.attention import _decode_gate_sigmoid_table
+
+            _decode_gate_sigmoid_table(self.device.index if self.device.index is not None else torch.cuda.current_device())
+
 
     def _split_cap(self, bs: int) -> int:
         """The most decode KV splits any layer group takes at this batch size."""
@@ -115,6 +128,51 @@ def _ensure_decode_scratch(
             device=self.device,
         )
 
+    def forward_gated(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        layer_id: int,
+        batch: Batch,
+        gate: torch.Tensor,
+        attn_spec: AttentionSpec | None = None,
+    ) -> torch.Tensor:
+        # Automatic fusion is limited to the B1/512 Graph validated end to end.
+        metadata = batch.attn_metadata
+        spec = attn_spec or AttentionSpec()
+        eligible = (
+            self._output_gate_fusion
+            and isinstance(metadata, TritonMetadata) and metadata.is_decode
+            and torch.cuda.is_current_stream_capturing()
+            and q.dtype == gate.dtype == torch.bfloat16
+            and q.shape == (1, 16, 256) and metadata.indices.numel() == 512
+            and gate.shape in (q.shape, (q.shape[0], 4096))
+            and gate.device == q.device and q.is_contiguous() and gate.is_contiguous()
+            and spec.sliding_window is None and spec.sinks is None
+            and spec.sm_scale is None and not spec.bidirectional_mm_blocks
+        )
+        if eligible:
+            k_raw, v_raw = self.kvcache.k_cache(layer_id), self.kvcache.v_cache(layer_id)
+            eligible = (k_raw.shape[-2:] == v_raw.shape[-2:] == (2, 256)
+                        and k_raw.dtype == v_raw.dtype == torch.bfloat16
+                        and k_raw.device == v_raw.device == q.device
+                        and k_raw.is_contiguous() and v_raw.is_contiguous())
+        if not eligible:
+            return super().forward_gated(q, k, v, layer_id, batch, gate, attn_spec=attn_spec)
+
+        from freetoken.kernel.triton.attention import decode_paged_attention
+
+        self._ensure_decode_scratch(metadata, q.shape[0], q.shape[1], q.shape[2])
+        # Decide before storing: the fallback must also write each KV row once.
+        self.kvcache.store_kv(k, v, batch.out_loc, layer_id)
+        return decode_paged_attention(
+            q, k_raw.view(-1, 2, 256), v_raw.view(-1, 2, 256),
+            metadata.indptr, metadata.indices, metadata.q_positions,
+            metadata.attn_logits[:q.shape[0]], metadata.attn_lse[:q.shape[0]],
+            q.shape[-1] ** -0.5, output_gate=gate.view_as(q),
+        )
+
     def forward(
         self,
         q: torch.Tensor,