From 02a096ba6398bca77c6208bdbde3720ef1d3d275 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Wed, 21 Oct 2026 15:57:10 +0800 Subject: [PATCH 1/2] perf(attention): fuse short bf16 decode output gates on gfx1201 --- benchmarks/bench_decode_output_gate.py | 158 ++++++++++++++++ docs/triton-decode-output-gate.md | 156 ++++++++++++++++ python/freetoken/attention/base.py | 29 +++ python/freetoken/attention/triton.py | 58 ++++++ python/freetoken/kernel/triton/attention.py | 66 ++++++- .../freetoken/models/qwen3_5_moe/attention.py | 8 +- tests/kernels/test_triton_attention.py | 171 ++++++++++++++++++ 7 files changed, 635 insertions(+), 11 deletions(-) create mode 100644 benchmarks/bench_decode_output_gate.py create mode 100644 docs/triton-decode-output-gate.md diff --git a/benchmarks/bench_decode_output_gate.py b/benchmarks/bench_decode_output_gate.py new file mode 100644 index 0000000000..6f871e2f6d --- /dev/null +++ b/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() diff --git a/docs/triton-decode-output-gate.md b/docs/triton-decode-output-gate.md new file mode 100644 index 0000000000..00c258a921 --- /dev/null +++ b/docs/triton-decode-output-gate.md @@ -0,0 +1,156 @@ +# Triton decode output gate + +Models using `attention_output * sigmoid(gate)` can call +`BaseAttnBackend.forward_gated`. The default implementation applies the Torch +sigmoid and multiply after the existing backend. Hybrid backends delegate to the +selected prefill/decode backend. Qwen3.5/3.6 full-attention layers use this interface. + +## Automatic dispatch + +The Triton backend fuses the gate into its existing decode stage2 only for: + +- ROCm gfx1201, BF16 query/gate/KV, query shape `[1,16,256]`, two KV heads; +- Graph capture, contiguous query/gate/KV, CSR index buffer capacity exactly 512; +- the default full-attention specification (no sinks, sliding window, custom scale, + or bidirectional multimodal blocks). + +The flag `FREETOKEN_TRITON_FUSE_OUTPUT_GATE=0` disables automatic fusion. +Eager, prefill, larger Graph capacities, other batch sizes, dtypes, GPU architectures and +attention backends use the default implementation. Eligibility is checked before +storing KV, so each route stores the new rows once. This change includes no HIP +extension, native build, model-name check or software-version fingerprint. + +Capacity is the host-visible allocated CSR buffer size, not the current effective +length. Setting max sequence capacity to 4096 retains the fallback even for a +short current request. The initial automatic scope is deliberately limited to the +B1/512 configuration with real-model validation. The low-level kernel accepts +broader BF16 shapes for explicit testing, including the B4 kernel benchmark; +these are not automatically enabled by the serving backend. + +## Arithmetic + +Attention stage1 and the stage2 FP32 reduction are retained. Stage2 rounds the +attention result to BF16, applies the BF16 sigmoid value, then rounds the product +as Torch does. A per-device 128 KiB table covers every BF16 input bit pattern and +is computed by Torch once before Graph capture. It preserves the active runtime's +sigmoid rounding. Explicit integer round-to-nearest-even for the product avoids +compiler narrowing to gfx12 BF16 dot2 arithmetic. An approximate sigmoid or +skipping the intermediate BF16 rounding can change FP8 MoE model logits. + +The measured window starts with prepared Q/K/V and CSR metadata and ends at the +gated attention output. QKV projection, rotary, KV writes and output projection +are outside the kernel measurement. Kernel speedup does not imply the same +end-to-end speedup. + +## Reproduce + +Kernel pairs, including the same buffers and an alternating Graph order: + +```bash +PYTHONPATH=python python benchmarks/bench_decode_output_gate.py \ + --baseline-path /path/to/main/python/freetoken/kernel/triton/attention.py \ + --json kernel.json +``` + +Run the following server command on main and this branch on the same free GPU, +with the same checkpoint, environment and options: + +```bash +PYTHONPATH=python python -m freetoken.cli serve \ + --model-path "$MODEL" --dtype bfloat16 --attention-backend triton \ + --moe-strategy offload --moe-cache-mem 20GiB --moe-cache-scope unified \ + --moe-cache-split equal --num-tokens 544 --page-size 1 \ + --max-seq-len-override 512 --max-prefill-length 512 \ + --max-running-requests 1 --cuda-graph-max-bs 1 --max-output-tokens 128 \ + --memory-ratio 0.9 --expert-load serial --cache-type naive \ + --sampling-defaults none --reasoning-parser off --text-model-only \ + --served-model-name gate-bench --host 127.0.0.1 --port 19741 +``` + +From this branch, measure either running server: + +```bash +PYTHONPATH=python python benchmarks/bench_decode_output_gate.py \ + --url http://127.0.0.1:19741 --model "$MODEL" \ + --input-tokens 128 --output-tokens 128 --repeats 1 --json server.json +``` + +The client warms once, checks exact usage, then measures from the first content +SSE to the finish frame for decode throughput. It separately reports TTFT and +request-wide throughput and saves event timestamps and an output hash. For the +larger-capacity control use 4096 for both max lengths, 4128 tokens, and a 512-token +input. Alternate main/branch order and restart the server for each arm. + +Tests: + +```bash +PYTHONPATH=python pytest tests/kernels/test_triton_attention.py \ + tests/engine/test_attention_backend_arch.py \ + tests/engine/test_attention_backend_matrix.py tests/daemon/test_daemon_import_safety.py +``` + +## 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. + +| Shape | Main us, round 1 / 2 | Branch us, round 1 / 2 | Speedup, round 1 / 2 | +|---|---:|---:|---:| +| B1 | 14.971 / 14.844 | 10.431 / 10.416 | 1.435 / 1.425x | +| B4 | 15.886 / 15.656 | 11.551 / 11.517 | 1.375 / 1.359x | + +### 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. + +| Graph capacity | Input/output tokens | Main / branch decode tok/s | Median paired decode gain | Median paired request-wide gain | Main / branch TTFT ms | +|---|---|---:|---:|---:|---:| +| 512 | 128 / 128 | 77.179 / 77.562 | +0.474% | -0.031% | 895.63 / 898.21 | +| 4096 | 512 / 128 | 76.723 / 76.723 | +0.001% | +0.047% | 895.77 / 903.04 | + +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. + +| Capacity | Round | Decode ratio | Request-wide ratio | +|---|---:|---:|---:| +| 512 | 1 | 1.004735x | 1.006108x | +| 512 | 2 | 1.003748x | 0.999692x | +| 512 | 3 | 1.009045x | 0.999500x | +| 4096 | 1 | 0.997650x | 0.991592x | +| 4096 | 2 | 1.003383x | 1.001867x | +| 4096 | 3 | 1.000008x | 1.000473x | + +### 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. diff --git a/python/freetoken/attention/base.py b/python/freetoken/attention/base.py index 846e99447a..0fe3d79f45 100644 --- a/python/freetoken/attention/base.py +++ b/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) diff --git a/python/freetoken/attention/triton.py b/python/freetoken/attention/triton.py index 7f29c481bb..9841fb4aaf 100644 --- a/python/freetoken/attention/triton.py +++ b/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, diff --git a/python/freetoken/kernel/triton/attention.py b/python/freetoken/kernel/triton/attention.py index d19457cffc..50d6a30867 100644 --- a/python/freetoken/kernel/triton/attention.py +++ b/python/freetoken/kernel/triton/attention.py @@ -19,6 +19,15 @@ _STAGE2_TILE_ELEMS = 8192 +@functools.lru_cache(maxsize=None) +def _decode_gate_sigmoid_table(device_index: int) -> torch.Tensor: + # Initialize before capture: Torch supplies its exact BF16 sigmoid rounding. + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("Initialize the decode gate sigmoid table before Graph capture") + bits = torch.arange(65536, device=torch.device("cuda", device_index), dtype=torch.int32).to(torch.int16) + return bits.view(torch.bfloat16).sigmoid() + + def _stage2_launch(max_kv_splits: int, block_dv: int) -> tuple[int, int]: """(splits merged per step, num_warps) for stage 2; the tile is at most the split cap rounded up to a power of two.""" tile = min(triton.next_power_of_2(max_kv_splits), max(4, _STAGE2_TILE_ELEMS // block_dv), 64) @@ -317,6 +326,10 @@ def _decode_stage2_kernel( q_pos_ptr, max_kv_splits, sinks_ptr, + gate_ptr, + sigmoid_ptr, + stride_gt, + stride_gh, stride_mid_ob, stride_mid_oh, stride_mid_os, @@ -332,6 +345,7 @@ def _decode_stage2_kernel( DV: tl.constexpr, SLIDING_WINDOW: tl.constexpr, HAS_SINKS: tl.constexpr, + HAS_GATE: tl.constexpr, SPLIT_TILE: tl.constexpr, ): batch_id = tl.program_id(0) @@ -384,11 +398,28 @@ def _decode_stage2_kernel( m_i = m_new out = tl.where(l_i == 0.0, 0.0, acc / l_i) - tl.store( - o_ptr + batch_id * stride_ot + q_head * stride_oh + offs_d, - out.to(o_ptr.dtype.element_ty), - mask=mask_d, - ) + if HAS_GATE: + rounded_attention = out.to(o_ptr.dtype.element_ty).to(tl.float32) + gate_bits = tl.load( + gate_ptr + batch_id * stride_gt + q_head * stride_gh + offs_d, + mask=mask_d, other=0, + ).to(tl.int32) & 65535 + sigmoid = tl.load(sigmoid_ptr + gate_bits).to(tl.float32) + product = rounded_attention * sigmoid + # Explicit RNE prevents narrowing the BF16 product to a gfx12 dot2. + bits = product.to(tl.uint32, bitcast=True) + rounded = (bits + 0x7fff + ((bits >> 16) & 1)) >> 16 + rounded = tl.where((bits & 0x7fffffff) > 0x7f800000, (bits >> 16) | 0x40, rounded) + tl.store( + o_ptr.to(tl.pointer_type(tl.uint16)) + batch_id * stride_ot + q_head * stride_oh + offs_d, + rounded.to(tl.uint16), mask=mask_d, + ) + else: + tl.store( + o_ptr + batch_id * stride_ot + q_head * stride_oh + offs_d, + out.to(o_ptr.dtype.element_ty), + mask=mask_d, + ) def decode_paged_attention( @@ -405,11 +436,17 @@ def decode_paged_attention( sinks: torch.Tensor | None = None, out: torch.Tensor | None = None, max_kv_splits: int | None = None, + *, + output_gate: torch.Tensor | None = None, ) -> torch.Tensor: """SGLang-style split-k grouped decode attention for one query per request. Each request takes ``max(ceil(keys / 128), min(8, ceil(keys / 32)))`` splits, capped by ``max_kv_splits`` (default :func:`decode_kv_split_cap`) and by the split dim of ``attn_logits`` / ``attn_lse``. The kernels compute the split count on the device, so a captured graph serves any context length. + + ``output_gate`` optionally fuses ``out * sigmoid(output_gate)`` into stage 2. + It must be a contiguous BF16 tensor shaped like q; initialize + ``_decode_gate_sigmoid_table(q.device.index)`` before Graph capture. """ assert q.is_cuda and k_cache.is_cuda and v_cache.is_cuda @@ -440,6 +477,20 @@ def decode_paged_attention( assert sinks.numel() >= num_q_heads sinks = sinks.contiguous() + gate_bits = q + sigmoid_table = q + stride_gt = stride_gh = 0 + if output_gate is not None: + if (q.dtype != torch.bfloat16 or output_gate.dtype != torch.bfloat16 + or output_gate.shape != q.shape or output_gate.device != q.device + or not output_gate.is_contiguous()): + raise ValueError("output_gate must be contiguous BF16 on q's device with q's shape") + if out is not None and out.dtype != torch.bfloat16: + raise ValueError("gated decode output must be BF16") + sigmoid_table = _decode_gate_sigmoid_table(q.device.index) + gate_bits = output_gate.view(torch.int16) + stride_gt, stride_gh = output_gate.stride(0), output_gate.stride(1) + o = out if out is not None else torch.empty_like(q) sinks_arg = sinks if sinks is not None else q group = num_q_heads // num_kv_heads @@ -511,6 +562,10 @@ def decode_paged_attention( q_positions, max_kv_splits, sinks_arg, + gate_bits, + sigmoid_table, + stride_gt, + stride_gh, attn_logits.stride(0), attn_logits.stride(1), attn_logits.stride(2), @@ -526,6 +581,7 @@ def decode_paged_attention( DV=head_dim, SLIDING_WINDOW=sliding_window or 0, HAS_SINKS=sinks is not None, + HAS_GATE=output_gate is not None, SPLIT_TILE=split_tile, num_warps=stage2_warps, ) diff --git a/python/freetoken/models/qwen3_5_moe/attention.py b/python/freetoken/models/qwen3_5_moe/attention.py index 0a736a1042..4379a100df 100644 --- a/python/freetoken/models/qwen3_5_moe/attention.py +++ b/python/freetoken/models/qwen3_5_moe/attention.py @@ -81,16 +81,12 @@ def _project(self, x: torch.Tensor): q, k = self.rotary.forward(positions, q, k) return q.view(-1, self.num_q, self.head_dim), k, v, gate - def _combine(self, attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - gated = attn_out.reshape(-1, self.qo_attn_dim) * torch.sigmoid(gate) - return self.o_proj.forward(gated) - @nvtx_annotate("MHA") def forward(self, x: torch.Tensor) -> torch.Tensor: ctx = get_global_ctx() q, k, v, gate = self._project(x) - o = ctx.attn_backend.forward(q, k, v, self.layer_id, ctx.batch) - return self._combine(o, gate) + gated = ctx.attn_backend.forward_gated(q, k, v, self.layer_id, ctx.batch, gate) + return self.o_proj.forward(gated.reshape(-1, self.qo_attn_dim)) __all__ = ["Qwen3_5Attention"] diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index 60a5e27477..ae23d2c972 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -923,3 +923,174 @@ def test_triton_metadata_keeps_full_indices_and_optional_swa_indices(monkeypatch assert metadata.indices.tolist() == [10, 11, 20, 21, 22] assert metadata.swa_indices is not None assert metadata.swa_indices.tolist() == [110, 111, 120, 121, 122] + + +def test_output_gate_fallback_and_hybrid_apply_gate_once(): + from freetoken.attention.base import AttentionSpec, BaseAttnBackend, HybridBackend + + class DenseBackend(BaseAttnBackend): + def __init__(self, value): + self.value = value + self.calls = [] + + def forward(self, q, k, v, layer_id, batch, attn_spec=None): + self.calls.append(attn_spec) + return self.value + + def prepare_metadata(self, batch): pass + def init_capture_graph(self, max_seq_len, bs_list): pass + def prepare_for_capture(self, batch): pass + def prepare_for_replay(self, batch): pass + + torch.manual_seed(29) + out = torch.randn(2, 3, 8) + gate = torch.randn(2, 24) + prefill, decode = DenseBackend(out), DenseBackend(out + 1) + backend = HybridBackend(prefill, decode) + spec = AttentionSpec(sm_scale=0.25) + for is_prefill, selected in [(True, prefill), (False, decode)]: + batch = SimpleNamespace(is_prefill=is_prefill) + actual = backend.forward_gated(out, out, out, 0, batch, gate, attn_spec=spec) + expected = (selected.value.reshape(2, 24) * gate.sigmoid()).reshape_as(out) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert selected.calls == [spec] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +@pytest.mark.parametrize( + ("lens", "heads", "kv_heads", "dim", "window", "use_sinks"), + [([512], 16, 2, 256, None, False), + ([512, 495, 478, 461], 16, 2, 256, None, False), + ([512], 16, 4, 256, None, False), + ([7, 129], 24, 4, 256, None, False), + ([1, 1000], 8, 2, 128, 300, True), + ([0, 65], 4, 1, 256, None, False)], +) +def test_decode_output_gate_matches_unfused_and_reference(lens, heads, kv_heads, dim, window, use_sinks): + from freetoken.kernel.triton.attention import decode_paged_attention, _decode_gate_sigmoid_table + + torch.manual_seed(37) + batch, total = len(lens), sum(lens) + q = torch.randn(batch, heads, dim, device="cuda", dtype=torch.bfloat16) * 0.2 + k = torch.randn(total, kv_heads, dim, device="cuda", dtype=torch.bfloat16) * 0.2 + v = torch.randn_like(k) * 0.2 + gate = torch.randn_like(q) * 8 + gate.flatten()[:3] = torch.tensor([0, float("inf"), -float("inf")], device="cuda") + ip = torch.tensor([0, *lens], device="cuda", dtype=torch.int32).cumsum(0).to(torch.int32) + ix = torch.randperm(total, device="cuda").to(torch.int32) + pos = torch.tensor([n - 1 for n in lens], device="cuda", dtype=torch.int64) + partial = torch.empty(batch, heads, 8, dim, device="cuda", dtype=torch.float32) + lse = torch.empty(batch, heads, 8, device="cuda", dtype=torch.float32) + sinks = torch.randn(heads, device="cuda") if use_sinks else None + args = (q, k, v, ip, ix, pos, partial, lse, dim ** -0.5) + kwargs = dict(sliding_window=window, sinks=sinks) + _decode_gate_sigmoid_table(q.device.index) + expected = decode_paged_attention(*args, **kwargs) * gate.sigmoid() + output = torch.empty_like(q) + actual = decode_paged_attention(*args, **kwargs, out=output, output_gate=gate) + assert actual.data_ptr() == output.data_ptr() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + if all(lens): + oracle = _reference_paged_attention( + q, k, v, ip, ix, torch.arange(batch, device="cuda"), pos, + dim ** -0.5, window, sinks=sinks, + ) * gate.sigmoid() + torch.testing.assert_close(actual.float(), oracle.float(), rtol=2e-2, atol=2e-2) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + replay_output = decode_paged_attention(*args, **kwargs, output_gate=gate) + # Change the inputs and effective CSR lengths after capture. + q.mul_(0.5) + gate.neg_() + pos.sub_(1) + expected = decode_paged_attention(*args, **kwargs) * gate.sigmoid() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(replay_output, expected, rtol=0, atol=0) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") +def test_output_gate_backend_graph_and_eager_fallback(monkeypatch): + from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata + from freetoken.kernel.triton import attention as kernel + + device = torch.device("cuda", torch.cuda.current_device()) + torch.manual_seed(41) + q = torch.randn(1, 16, 256, device=device, dtype=torch.bfloat16) * 0.2 + gate = torch.randn(1, 4096, device=device, dtype=torch.bfloat16) + k = torch.randn(512, 2, 256, device=device, dtype=torch.bfloat16) * 0.2 + v = torch.randn_like(k) + stores, calls = [], [] + cache = SimpleNamespace(device=device, k_cache=lambda _: k, v_cache=lambda _: v, + store_kv=lambda *args: stores.append(True)) + monkeypatch.setattr("freetoken.attention.triton.get_global_ctx", lambda: SimpleNamespace(kv_cache=cache)) + monkeypatch.delenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", raising=False) + config = SimpleNamespace(num_qo_heads=16, num_kv_heads=2, head_dim=256) + backend = TritonAttentionBackend(config) + ip = torch.tensor([0, 512], device=device, dtype=torch.int32) + ix = torch.arange(512, device=device, dtype=torch.int32) + pos = torch.tensor([511], device=device, dtype=torch.int64) + metadata = TritonMetadata(ip, ip, ix, torch.zeros_like(pos), pos, True, torch.zeros_like(pos), 1) + batch = SimpleNamespace(attn_metadata=metadata, out_loc=pos, mm_block_ends=None) + original = kernel.decode_paged_attention + + def observed(*args, **kwargs): + calls.append(kwargs.get("output_gate") is not None) + return original(*args, **kwargs) + + monkeypatch.setattr(kernel, "decode_paged_attention", observed) + expected = backend.forward(q, k[-1:], v[-1:], 0, batch) * gate.sigmoid().view_as(q) + eager = backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) + torch.testing.assert_close(eager, expected, rtol=0, atol=0) + assert len(stores) == 2 and calls == [False, False] + # Compile the optional kernel before capture on platforms without automatic admission. + kernel._decode_gate_sigmoid_table(device.index) + original(q, k, v, ip, ix, pos, metadata.attn_logits, metadata.attn_lse, + 256 ** -0.5, output_gate=gate.view_as(q)) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + expected_fusion = bool(torch.version.hip) and getattr( + torch.cuda.get_device_properties(device), "gcnArchName", "" + ).startswith("gfx1201") + assert len(stores) == 3 and calls[-1] == expected_fusion + # A larger CSR capacity must retain FT even while capturing. + metadata.indices = torch.cat((ix, torch.zeros(3584, device=device, dtype=torch.int32))) + fallback_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(fallback_graph, stream=stream): + fallback = backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) + fallback_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(fallback, expected, rtol=0, atol=0) + assert calls[-1] is False and len(stores) == 4 + metadata.indices = ix + monkeypatch.setenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", "0") + disabled = TritonAttentionBackend(config) + disabled_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(disabled_graph, stream=stream): + disabled_out = disabled.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) + disabled_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(disabled_out, expected, rtol=0, atol=0) + assert calls[-1] is False and len(stores) == 5 + q4, gate4 = q.repeat(4, 1, 1), gate.repeat(4, 1) + metadata.indptr = torch.arange(5, device=device, dtype=torch.int32) * 512 + metadata.indices = ix.repeat(4) + metadata.q_positions = pos.repeat(4) + expected4 = backend.forward(q4, k[-1:], v[-1:], 0, batch) * gate4.sigmoid().view_as(q4) + batch4_graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(batch4_graph, stream=stream): + batch4_out = backend.forward_gated(q4, k[-1:], v[-1:], 0, batch, gate4) + batch4_graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(batch4_out, expected4, rtol=0, atol=0) + assert calls[-1] is False and len(stores) == 7 From 62cfaa68ef9aefc657dce30047fdd3e49ba9cc0b Mon Sep 17 00:00:00 2001 From: zihaomu Date: Wed, 21 Oct 2026 16:35:08 +0800 Subject: [PATCH 2/2] test(attention): streamline output gate regression setup --- docs/triton-decode-output-gate.md | 156 ------------------------- tests/kernels/test_triton_attention.py | 96 +++++++-------- 2 files changed, 39 insertions(+), 213 deletions(-) delete mode 100644 docs/triton-decode-output-gate.md diff --git a/docs/triton-decode-output-gate.md b/docs/triton-decode-output-gate.md deleted file mode 100644 index 00c258a921..0000000000 --- a/docs/triton-decode-output-gate.md +++ /dev/null @@ -1,156 +0,0 @@ -# Triton decode output gate - -Models using `attention_output * sigmoid(gate)` can call -`BaseAttnBackend.forward_gated`. The default implementation applies the Torch -sigmoid and multiply after the existing backend. Hybrid backends delegate to the -selected prefill/decode backend. Qwen3.5/3.6 full-attention layers use this interface. - -## Automatic dispatch - -The Triton backend fuses the gate into its existing decode stage2 only for: - -- ROCm gfx1201, BF16 query/gate/KV, query shape `[1,16,256]`, two KV heads; -- Graph capture, contiguous query/gate/KV, CSR index buffer capacity exactly 512; -- the default full-attention specification (no sinks, sliding window, custom scale, - or bidirectional multimodal blocks). - -The flag `FREETOKEN_TRITON_FUSE_OUTPUT_GATE=0` disables automatic fusion. -Eager, prefill, larger Graph capacities, other batch sizes, dtypes, GPU architectures and -attention backends use the default implementation. Eligibility is checked before -storing KV, so each route stores the new rows once. This change includes no HIP -extension, native build, model-name check or software-version fingerprint. - -Capacity is the host-visible allocated CSR buffer size, not the current effective -length. Setting max sequence capacity to 4096 retains the fallback even for a -short current request. The initial automatic scope is deliberately limited to the -B1/512 configuration with real-model validation. The low-level kernel accepts -broader BF16 shapes for explicit testing, including the B4 kernel benchmark; -these are not automatically enabled by the serving backend. - -## Arithmetic - -Attention stage1 and the stage2 FP32 reduction are retained. Stage2 rounds the -attention result to BF16, applies the BF16 sigmoid value, then rounds the product -as Torch does. A per-device 128 KiB table covers every BF16 input bit pattern and -is computed by Torch once before Graph capture. It preserves the active runtime's -sigmoid rounding. Explicit integer round-to-nearest-even for the product avoids -compiler narrowing to gfx12 BF16 dot2 arithmetic. An approximate sigmoid or -skipping the intermediate BF16 rounding can change FP8 MoE model logits. - -The measured window starts with prepared Q/K/V and CSR metadata and ends at the -gated attention output. QKV projection, rotary, KV writes and output projection -are outside the kernel measurement. Kernel speedup does not imply the same -end-to-end speedup. - -## Reproduce - -Kernel pairs, including the same buffers and an alternating Graph order: - -```bash -PYTHONPATH=python python benchmarks/bench_decode_output_gate.py \ - --baseline-path /path/to/main/python/freetoken/kernel/triton/attention.py \ - --json kernel.json -``` - -Run the following server command on main and this branch on the same free GPU, -with the same checkpoint, environment and options: - -```bash -PYTHONPATH=python python -m freetoken.cli serve \ - --model-path "$MODEL" --dtype bfloat16 --attention-backend triton \ - --moe-strategy offload --moe-cache-mem 20GiB --moe-cache-scope unified \ - --moe-cache-split equal --num-tokens 544 --page-size 1 \ - --max-seq-len-override 512 --max-prefill-length 512 \ - --max-running-requests 1 --cuda-graph-max-bs 1 --max-output-tokens 128 \ - --memory-ratio 0.9 --expert-load serial --cache-type naive \ - --sampling-defaults none --reasoning-parser off --text-model-only \ - --served-model-name gate-bench --host 127.0.0.1 --port 19741 -``` - -From this branch, measure either running server: - -```bash -PYTHONPATH=python python benchmarks/bench_decode_output_gate.py \ - --url http://127.0.0.1:19741 --model "$MODEL" \ - --input-tokens 128 --output-tokens 128 --repeats 1 --json server.json -``` - -The client warms once, checks exact usage, then measures from the first content -SSE to the finish frame for decode throughput. It separately reports TTFT and -request-wide throughput and saves event timestamps and an output hash. For the -larger-capacity control use 4096 for both max lengths, 4128 tokens, and a 512-token -input. Alternate main/branch order and restart the server for each arm. - -Tests: - -```bash -PYTHONPATH=python pytest tests/kernels/test_triton_attention.py \ - tests/engine/test_attention_backend_arch.py \ - tests/engine/test_attention_backend_matrix.py tests/daemon/test_daemon_import_safety.py -``` - -## 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. - -| Shape | Main us, round 1 / 2 | Branch us, round 1 / 2 | Speedup, round 1 / 2 | -|---|---:|---:|---:| -| B1 | 14.971 / 14.844 | 10.431 / 10.416 | 1.435 / 1.425x | -| B4 | 15.886 / 15.656 | 11.551 / 11.517 | 1.375 / 1.359x | - -### 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. - -| Graph capacity | Input/output tokens | Main / branch decode tok/s | Median paired decode gain | Median paired request-wide gain | Main / branch TTFT ms | -|---|---|---:|---:|---:|---:| -| 512 | 128 / 128 | 77.179 / 77.562 | +0.474% | -0.031% | 895.63 / 898.21 | -| 4096 | 512 / 128 | 76.723 / 76.723 | +0.001% | +0.047% | 895.77 / 903.04 | - -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. - -| Capacity | Round | Decode ratio | Request-wide ratio | -|---|---:|---:|---:| -| 512 | 1 | 1.004735x | 1.006108x | -| 512 | 2 | 1.003748x | 0.999692x | -| 512 | 3 | 1.009045x | 0.999500x | -| 4096 | 1 | 0.997650x | 0.991592x | -| 4096 | 2 | 1.003383x | 1.001867x | -| 4096 | 3 | 1.000008x | 1.000473x | - -### 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. diff --git a/tests/kernels/test_triton_attention.py b/tests/kernels/test_triton_attention.py index ae23d2c972..294ac0a072 100644 --- a/tests/kernels/test_triton_attention.py +++ b/tests/kernels/test_triton_attention.py @@ -925,6 +925,15 @@ def test_triton_metadata_keeps_full_indices_and_optional_swa_indices(monkeypatch assert metadata.swa_indices.tolist() == [110, 111, 120, 121, 122] +def _capture_output_gate_call(call): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = call() + return graph, output + + def test_output_gate_fallback_and_hybrid_apply_gate_once(): from freetoken.attention.base import AttentionSpec, BaseAttnBackend, HybridBackend @@ -998,11 +1007,9 @@ def test_decode_output_gate_matches_unfused_and_reference(lens, heads, kv_heads, ) * gate.sigmoid() torch.testing.assert_close(actual.float(), oracle.float(), rtol=2e-2, atol=2e-2) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - replay_output = decode_paged_attention(*args, **kwargs, output_gate=gate) + graph, replay_output = _capture_output_gate_call( + lambda: decode_paged_attention(*args, **kwargs, output_gate=gate) + ) # Change the inputs and effective CSR lengths after capture. q.mul_(0.5) gate.neg_() @@ -1014,27 +1021,36 @@ def test_decode_output_gate_matches_unfused_and_reference(lens, heads, kv_heads, @pytest.mark.skipif(not torch.cuda.is_available(), reason="Triton attention needs CUDA") -def test_output_gate_backend_graph_and_eager_fallback(monkeypatch): +@pytest.mark.parametrize( + ("capacity", "batch_size", "enabled"), + [(512, 1, True), (4096, 1, True), (512, 1, False), (512, 4, True)], +) +def test_output_gate_backend_graph_and_eager_fallback(monkeypatch, capacity, batch_size, enabled): from freetoken.attention.triton import TritonAttentionBackend, TritonMetadata from freetoken.kernel.triton import attention as kernel device = torch.device("cuda", torch.cuda.current_device()) torch.manual_seed(41) - q = torch.randn(1, 16, 256, device=device, dtype=torch.bfloat16) * 0.2 - gate = torch.randn(1, 4096, device=device, dtype=torch.bfloat16) + q = torch.randn(batch_size, 16, 256, device=device, dtype=torch.bfloat16) * 0.2 + gate = torch.randn(batch_size, 4096, device=device, dtype=torch.bfloat16) k = torch.randn(512, 2, 256, device=device, dtype=torch.bfloat16) * 0.2 v = torch.randn_like(k) stores, calls = [], [] cache = SimpleNamespace(device=device, k_cache=lambda _: k, v_cache=lambda _: v, store_kv=lambda *args: stores.append(True)) monkeypatch.setattr("freetoken.attention.triton.get_global_ctx", lambda: SimpleNamespace(kv_cache=cache)) - monkeypatch.delenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", raising=False) + if enabled: + monkeypatch.delenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", raising=False) + else: + monkeypatch.setenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", "0") config = SimpleNamespace(num_qo_heads=16, num_kv_heads=2, head_dim=256) backend = TritonAttentionBackend(config) - ip = torch.tensor([0, 512], device=device, dtype=torch.int32) - ix = torch.arange(512, device=device, dtype=torch.int32) - pos = torch.tensor([511], device=device, dtype=torch.int64) - metadata = TritonMetadata(ip, ip, ix, torch.zeros_like(pos), pos, True, torch.zeros_like(pos), 1) + cu_q = torch.arange(batch_size + 1, device=device, dtype=torch.int32) + ip = cu_q * 512 + ix = torch.cat((torch.arange(512, device=device, dtype=torch.int32), + torch.zeros(capacity - 512, device=device, dtype=torch.int32))).repeat(batch_size) + pos = torch.full((batch_size,), 511, device=device, dtype=torch.int64) + metadata = TritonMetadata(cu_q, ip, ix, torch.arange(batch_size, device=device), pos, True, torch.zeros_like(pos), 1) batch = SimpleNamespace(attn_metadata=metadata, out_loc=pos, mm_block_ends=None) original = kernel.decode_paged_attention @@ -1047,50 +1063,16 @@ def observed(*args, **kwargs): eager = backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) torch.testing.assert_close(eager, expected, rtol=0, atol=0) assert len(stores) == 2 and calls == [False, False] - # Compile the optional kernel before capture on platforms without automatic admission. - kernel._decode_gate_sigmoid_table(device.index) - original(q, k, v, ip, ix, pos, metadata.attn_logits, metadata.attn_lse, - 256 ** -0.5, output_gate=gate.view_as(q)) - stream = torch.cuda.Stream() - stream.wait_stream(torch.cuda.current_stream()) - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=stream): - actual = backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) + expected_fusion = enabled and capacity == 512 and batch_size == 1 and bool(torch.version.hip) and getattr( + torch.cuda.get_device_properties(device), "gcnArchName", "" + ).split(":")[0] == "gfx1201" + if expected_fusion: + original(q, k, v, ip, ix, pos, metadata.attn_logits, metadata.attn_lse, + 256 ** -0.5, output_gate=gate.view_as(q)) + graph, actual = _capture_output_gate_call( + lambda: backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) + ) graph.replay() torch.cuda.synchronize() torch.testing.assert_close(actual, expected, rtol=0, atol=0) - expected_fusion = bool(torch.version.hip) and getattr( - torch.cuda.get_device_properties(device), "gcnArchName", "" - ).startswith("gfx1201") - assert len(stores) == 3 and calls[-1] == expected_fusion - # A larger CSR capacity must retain FT even while capturing. - metadata.indices = torch.cat((ix, torch.zeros(3584, device=device, dtype=torch.int32))) - fallback_graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(fallback_graph, stream=stream): - fallback = backend.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) - fallback_graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(fallback, expected, rtol=0, atol=0) - assert calls[-1] is False and len(stores) == 4 - metadata.indices = ix - monkeypatch.setenv("FREETOKEN_TRITON_FUSE_OUTPUT_GATE", "0") - disabled = TritonAttentionBackend(config) - disabled_graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(disabled_graph, stream=stream): - disabled_out = disabled.forward_gated(q, k[-1:], v[-1:], 0, batch, gate) - disabled_graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(disabled_out, expected, rtol=0, atol=0) - assert calls[-1] is False and len(stores) == 5 - q4, gate4 = q.repeat(4, 1, 1), gate.repeat(4, 1) - metadata.indptr = torch.arange(5, device=device, dtype=torch.int32) * 512 - metadata.indices = ix.repeat(4) - metadata.q_positions = pos.repeat(4) - expected4 = backend.forward(q4, k[-1:], v[-1:], 0, batch) * gate4.sigmoid().view_as(q4) - batch4_graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(batch4_graph, stream=stream): - batch4_out = backend.forward_gated(q4, k[-1:], v[-1:], 0, batch, gate4) - batch4_graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(batch4_out, expected4, rtol=0, atol=0) - assert calls[-1] is False and len(stores) == 7 + assert len(stores) == 3 and calls == [False, False, expected_fusion]