diff --git a/benchmarks/bench_cake_fused_qk_rope_fp8_append.py b/benchmarks/bench_cake_fused_qk_rope_fp8_append.py new file mode 100644 index 00000000000..a572105a6ac --- /dev/null +++ b/benchmarks/bench_cake_fused_qk_rope_fp8_append.py @@ -0,0 +1,387 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Benchmark: Cake fused QK RMSNorm + NeoX RoPE + FP8 quantize + paged KV append. + +Arms (identical inputs, caller-owned buffers re-created per arm): + +- ``cake``: :func:`flashinfer.cake_fused_qk_rope_fp8_append.cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache` + (one launch). +- ``stable``: the step-wise pipeline from stable FlashInfer / torch ops + (``flashinfer.norm.rmsnorm`` -> ``apply_rope_pos_ids_inplace`` -> torch FP8 quantize + -> ``flashinfer.page.append_paged_kv_cache``), report only: its BF16 intermediates + cannot meet the FP8 contract tolerance and it does not clear the last-page tail. + +Timing: ``flashinfer.testing.bench_gpu_time`` (CUPTI, cold L2) over the whole call (the +union of every kernel one public call issues), median over ``--iters`` after ``--warmup``, +arms interleaved per group so clock drift affects all arms alike. + + python benchmarks/bench_cake_fused_qk_rope_fp8_append.py --rows P1,P2,H1 --json out.json +""" + +from __future__ import annotations + +import argparse +import json +import statistics +from dataclasses import dataclass + +import torch + +import flashinfer +from flashinfer.cake_fused_qk_rope_fp8_append import ( + cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache, +) +from flashinfer.testing.utils import bench_gpu_time + +HEAD_DIM = 128 +FP8_MAX = 448.0 + + +@dataclass(frozen=True) +class Row: + name: str + hq: int + hkv: int + batch: int + q_len: int + ctx: int + quant: int = 1 + norm: int = 2 + page_size: int = 64 + + +ROWS = { + r.name: r + for r in ( + # decode (B32 / Q1) and prefill (B8 / Q16), both head configurations, both quant policies + Row("D1", 8, 1, 32, 1, 2048, quant=1), + Row("D1n0", 8, 1, 32, 1, 2048, quant=1, norm=0), + Row("D1n1", 8, 1, 32, 1, 2048, quant=1, norm=1), + Row("D2", 8, 1, 32, 1, 2048, quant=2), + Row("P1", 8, 1, 8, 16, 2048, quant=1), + Row("P2", 8, 1, 8, 16, 2048, quant=2), + Row("H1", 64, 8, 32, 1, 2048, quant=1), + Row("H2", 64, 8, 32, 1, 2048, quant=2), + Row("HP1", 64, 8, 8, 16, 2048, quant=1), + Row("HP2", 64, 8, 8, 16, 2048, quant=2), + # dynamic-scale padding boundary (max_seqlen = Q) + Row("E127", 8, 1, 4, 127, 2048, quant=1), + Row("E128", 8, 1, 4, 128, 2048, quant=1), + Row("E129", 8, 1, 4, 129, 2048, quant=1), + # report-only extremes + Row("X1", 8, 1, 1, 1, 2048), + Row("X128", 8, 1, 128, 1, 2048), + Row("X256", 8, 1, 256, 1, 2048), + Row("XP1024", 8, 1, 1, 1024, 2048), + Row("Xpage16", 8, 1, 32, 1, 2048, page_size=16), + Row("Xpage128", 8, 1, 32, 1, 2048, page_size=128), + Row("Xctx32k", 8, 1, 32, 1, 32768), + ) +} + + +def rotary_cos_sin_table( + max_positions: int, device, base: float = 10000.0 +) -> torch.Tensor: + half = HEAD_DIM // 2 + inv_freq = 1.0 / ( + base ** (torch.arange(0, half, dtype=torch.float64, device=device) / half) + ) + pos = torch.arange(max_positions, dtype=torch.float64, device=device) + freqs = torch.outer(pos, inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], dim=1).to(torch.float32) + + +def build_inputs(row: Row, device, seed: int = 893): + g = torch.Generator(device="cpu").manual_seed(seed) + B = row.batch + q_lens = [row.q_len] * B + seq_lens = [row.ctx + row.q_len] * B + pages_per_req = [(s + row.page_size - 1) // row.page_size for s in seq_lens] + max_pages = max(pages_per_req) + total_pages = sum(pages_per_req) + 4 + perm = torch.randperm(total_pages, generator=g) + page_indices = torch.zeros((B, max_pages), dtype=torch.int32) + cursor = 0 + for b in range(B): + n = pages_per_req[b] + page_indices[b, :n] = perm[cursor : cursor + n].to(torch.int32) + cursor += n + q_indptr = torch.zeros(B + 1, dtype=torch.int32) + q_indptr[1:] = torch.cumsum(torch.tensor(q_lens, dtype=torch.int32), 0) + T = int(q_indptr[-1]) + width = (row.hq + 2 * row.hkv) * HEAD_DIM + qkv = ( + torch.randn((T, width), generator=g, dtype=torch.float32) + .to(torch.bfloat16) + .to(device) + ) + q_w = (torch.rand(HEAD_DIM, generator=g) + 0.5).to(torch.bfloat16) + k_w = (torch.rand(HEAD_DIM, generator=g) + 0.5).to(torch.bfloat16) + positions = torch.cat( + [ + torch.arange(row.ctx, row.ctx + row.q_len, dtype=torch.int32) + for _ in range(B) + ] + ) + batch_indices = torch.repeat_interleave( + torch.arange(B, dtype=torch.int32), row.q_len + ) + kv_indptr = torch.zeros(B + 1, dtype=torch.int32) + kv_indptr[1:] = torch.cumsum(torch.tensor(pages_per_req, dtype=torch.int32), 0) + kv_indices = torch.cat([page_indices[b, : pages_per_req[b]] for b in range(B)]) + kv_last_page_len = torch.tensor( + [(s - 1) % row.page_size + 1 for s in seq_lens], dtype=torch.int32 + ) + return { + "row": row, + "T": T, + "is_prefill": row.q_len > 1, + "max_seqlen": row.q_len if row.q_len > 1 else 0, + "qkv": qkv, + "cos_sin": rotary_cos_sin_table(max(seq_lens) + 1, device), + "seq_lens": torch.tensor(seq_lens, dtype=torch.int32, device=device), + "q_indptr": q_indptr.to(device), + "page_indices": page_indices.to(device), + "q_norm_weight": q_w.float().to(device), + "k_norm_weight": k_w.float().to(device), + "q_norm_weight_bf16": q_w.to(device), + "k_norm_weight_bf16": k_w.to(device), + "k_scale": torch.tensor([0.02], dtype=torch.float32, device=device), + "v_scale": torch.tensor([0.03], dtype=torch.float32, device=device), + "q_scale_inv": torch.tensor([1.0 / 0.02], dtype=torch.float32, device=device), + "positions": positions.to(device), + "batch_indices": batch_indices.to(device), + "kv_indptr": kv_indptr.to(device), + "kv_indices": kv_indices.to(device), + "kv_last_page_len": kv_last_page_len.to(device), + "total_pages": total_pages, + } + + +def fresh_buffers(inp, device): + row = inp["row"] + shape = (inp["total_pages"], row.page_size, row.hkv, HEAD_DIM) + kc = torch.zeros(shape, dtype=torch.float8_e4m3fn, device=device) + vc = torch.zeros(shape, dtype=torch.float8_e4m3fn, device=device) + oq = torch.empty( + (inp["T"], row.hq, HEAD_DIM), dtype=torch.float8_e4m3fn, device=device + ) + if row.quant == 2: + qs = torch.empty(0, dtype=torch.float32, device=device) + elif inp["is_prefill"]: + aligned = (inp["max_seqlen"] + 127) // 128 * 128 + qs = torch.empty( + (row.batch, row.hq, aligned), dtype=torch.float32, device=device + ) + else: + qs = torch.empty((inp["T"], row.hq), dtype=torch.float32, device=device) + flags = torch.zeros((row.batch, row.hkv), dtype=torch.int32, device=device) + return kc, vc, oq, qs, flags + + +def _fused_kwargs(inp, bufs): + row = inp["row"] + kc, vc, oq, qs, flags = bufs + return dict( + args=( + inp["qkv"], + inp["cos_sin"], + inp["seq_lens"], + inp["q_indptr"], + inp["page_indices"], + (kc, vc), + inp["is_prefill"], + inp["k_scale"], + inp["v_scale"], + row.quant, + ), + kwargs=dict( + max_seqlen=inp["max_seqlen"], + upper_max=FP8_MAX, + q_scale_inv=inp["q_scale_inv"] if row.quant == 2 else None, + q_norm_weight=inp["q_norm_weight"] if row.norm else None, + k_norm_weight=inp["k_norm_weight"] if row.norm else None, + qk_norm_policy=row.norm, + out_q=oq, + q_scale=qs, + split_k_flag=flags, + ), + ) + + +def make_cake_call(inp, device): + bufs = fresh_buffers(inp, device) + bound = _fused_kwargs(inp, bufs) + + def call(): + cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache( + *bound["args"], **bound["kwargs"] + ) + + return call, bufs + + +def make_stable_call(inp, device): + """rmsnorm + apply_rope_pos_ids_inplace + torch FP8 quantize + append_paged_kv_cache.""" + row = inp["row"] + kc, vc, oq, qs, flags = bufs = fresh_buffers(inp, device) + T = inp["T"] + q = inp["qkv"][:, : row.hq * HEAD_DIM].view(T, row.hq, HEAD_DIM) + k = inp["qkv"][:, row.hq * HEAD_DIM : (row.hq + row.hkv) * HEAD_DIM].view( + T, row.hkv, HEAD_DIM + ) + v = inp["qkv"][:, (row.hq + row.hkv) * HEAD_DIM :].view(T, row.hkv, HEAD_DIM) + tmp_q, tmp_k = torch.empty_like(q), torch.empty_like(k) + norm_q, norm_k = torch.empty_like(q), torch.empty_like(k) + k_fp8 = torch.empty_like(k, dtype=torch.float8_e4m3fn) + v_fp8 = torch.empty_like(v, dtype=torch.float8_e4m3fn) + qw, kw = inp["q_norm_weight_bf16"], inp["k_norm_weight_bf16"] + paged = (kc, vc) + k_inv, v_inv = 1.0 / inp["k_scale"], 1.0 / inp["v_scale"] + + def call(): + if row.norm == 2: + flashinfer.norm.rmsnorm( + q.reshape(-1, HEAD_DIM), qw, out=norm_q.view(-1, HEAD_DIM) + ) + flashinfer.norm.rmsnorm( + k.reshape(-1, HEAD_DIM), kw, out=norm_k.view(-1, HEAD_DIM) + ) + flashinfer.rope.apply_rope_pos_ids_inplace( + norm_q, norm_k, inp["positions"], interleave=False + ) + q_src, k_src = norm_q, norm_k + elif row.norm == 1: + tmp_q.copy_(q) + tmp_k.copy_(k) + flashinfer.rope.apply_rope_pos_ids_inplace( + tmp_q, tmp_k, inp["positions"], interleave=False + ) + flashinfer.norm.rmsnorm( + tmp_q.view(-1, HEAD_DIM), qw, out=norm_q.view(-1, HEAD_DIM) + ) + flashinfer.norm.rmsnorm( + tmp_k.view(-1, HEAD_DIM), kw, out=norm_k.view(-1, HEAD_DIM) + ) + q_src, k_src = norm_q, norm_k + else: + tmp_q.copy_(q) + tmp_k.copy_(k) + flashinfer.rope.apply_rope_pos_ids_inplace( + tmp_q, tmp_k, inp["positions"], interleave=False + ) + q_src, k_src = tmp_q, tmp_k + qf = q_src.float() + if row.quant == 1: + scale = qf.abs().amax(dim=-1, keepdim=True).clamp_min(1e-6) / FP8_MAX + oq.copy_((qf / scale).clamp(-FP8_MAX, FP8_MAX)) + if inp["is_prefill"]: + qs.view(row.batch, row.hq, -1)[:, :, : row.q_len].copy_( + scale.view(row.batch, row.q_len, row.hq).permute(0, 2, 1) + ) + else: + qs.copy_(scale.view(T, row.hq)) + else: + oq.copy_((qf * inp["q_scale_inv"]).clamp(-FP8_MAX, FP8_MAX)) + k_fp8.copy_((k_src.float() * k_inv).clamp(-FP8_MAX, FP8_MAX)) + v_fp8.copy_((v.float() * v_inv).clamp(-FP8_MAX, FP8_MAX)) + flashinfer.page.append_paged_kv_cache( + k_fp8, + v_fp8, + inp["batch_indices"], + inp["positions"], + paged, + inp["kv_indices"], + inp["kv_indptr"], + inp["kv_last_page_len"], + kv_layout="NHD", + ) + + return call, bufs + + +def time_arm(call, *, warmup: int, iters: int): + times = bench_gpu_time( + call, + dry_run_iters=warmup, + repeat_iters=iters, + cold_l2_cache=True, + enable_cupti=True, + ) + return float(statistics.median(times)) * 1e3 # ms -> us + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--rows", default=",".join(ROWS), help="comma list of row names") + ap.add_argument("--groups", type=int, default=3) + ap.add_argument("--iters", type=int, default=30) + ap.add_argument("--warmup", type=int, default=10) + ap.add_argument( + "--no-stable", action="store_true", help="skip the stable pipeline arm" + ) + ap.add_argument("--json", default=None) + args = ap.parse_args() + device = torch.device("cuda:0") + arms = ["cake"] + ([] if args.no_stable else ["stable"]) + results = [] + print(f"device {torch.cuda.get_device_name(device)} arms {arms}") + print( + f"{'row':10s} {'cfg':28s} " + + " ".join(f"{a:>10s}" for a in arms) + + " speedup vs others" + ) + for name in args.rows.split(","): + row = ROWS[name] + inp = build_inputs(row, device) + samples = {a: [] for a in arms} + for _ in range(args.groups): + for arm in arms: + if arm == "cake": + call, _bufs = make_cake_call(inp, device) + else: + call, _bufs = make_stable_call(inp, device) + samples[arm].append( + time_arm(call, warmup=args.warmup, iters=args.iters) + ) + med = {a: statistics.median(samples[a]) for a in arms} + cfg = f"({row.hq},{row.hkv}) B{row.batch} Q{row.q_len} ctx{row.ctx} q{row.quant} n{row.norm} p{row.page_size}" + ratios = " ".join( + f"{a}/cake={med[a] / med['cake']:.3f}" for a in arms if a != "cake" + ) + print( + f"{name:10s} {cfg:28s} " + + " ".join(f"{med[a]:10.2f}" for a in arms) + + f" {ratios}" + ) + results.append( + {"row": name, "config": cfg, "median_us": med, "samples_us": samples} + ) + if args.json: + with open(args.json, "w") as f: + json.dump( + { + "device": torch.cuda.get_device_name(device), + "arms": arms, + "rows": results, + }, + f, + indent=1, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/cake_fused_qk_rope_fp8_append/cake_fused_qk_rope_fp8_append_manifest.json b/csrc/cake_fused_qk_rope_fp8_append/cake_fused_qk_rope_fp8_append_manifest.json new file mode 100644 index 00000000000..de7b4ea12dd --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/cake_fused_qk_rope_fp8_append_manifest.json @@ -0,0 +1,714 @@ +{ + "artifact_kind": "source_only", + "build_contract": { + "architecture_source": "modules[].arch", + "binary_payloads": false, + "source_license_header_sha256": "dcf2449c2d70596e2c32bf1fa23720023965ff05857b85772decee99d830bf2f", + "target_infrastructure": { + "binding_runtime": "flashinfer_tvm_ffi_utils", + "headers_owned_by_target": true, + "required_headers": [ + "tvm_ffi_utils.h" + ] + }, + "translation_unit_model": "separate_device_and_binding" + }, + "contract": { + "arches": [ + "sm_90a", + "sm_100a", + "sm_103a" + ], + "integration_files": [ + "flashinfer/cake_fused_qk_rope_fp8_append.py", + "flashinfer/jit/cake_fused_qk_rope_fp8_append.py" + ], + "operator": "fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache", + "stages": [ + "fp8_e4m3" + ] + }, + "files": [ + { + "bytes": 12770, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_binding.cu", + "sha256": "23ebf2522ad4038b25df09ed5ee8030cc2b8aee2a1c960a699ba0e7b7c943669" + }, + { + "bytes": 31903, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_kernel.cu", + "sha256": "fc633c8fa81d145368fa175aef3be4734b5a0459ff34eab2b7f2b218ac8ede6b" + }, + { + "bytes": 12770, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_binding.cu", + "sha256": "47f5f9235a742595e0f5431b6eee5ef1582133a0c8a37ba798d342f6fb86bd3f" + }, + { + "bytes": 31903, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_kernel.cu", + "sha256": "130e94500ec57ad30b1e53627080a49837a7c1877c0732f5917815759c0949b0" + }, + { + "bytes": 12770, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_binding.cu", + "sha256": "7cb4035a255194f4d20eddcf62bed9f49f2a67e9a471f254e270c1b07d01f3ea" + }, + { + "bytes": 31903, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_kernel.cu", + "sha256": "695774941695ef5ba913db34bd803c794e5814cec6fd8424920430edefd017ae" + } + ], + "generation_inputs": { + "schema": "cake.library_export.generation_inputs.v4" + }, + "library": "flashinfer", + "modules": [ + { + "arch": "sm_100a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "page_indices" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "buffer", + "k_scale" + ], + [ + "buffer", + "v_scale" + ], + [ + "buffer", + "q_scale_inv" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "key_cache" + ], + [ + "buffer", + "value_cache" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_scale" + ], + [ + "buffer", + "split_k_flag" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "num_q_heads" + ], + [ + "parameter", + "num_kv_heads" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "max_pages_per_request" + ], + [ + "parameter", + "max_seqlen" + ], + [ + "parameter", + "max_seqlen_aligned" + ], + [ + "parameter", + "quant_policy" + ], + [ + "parameter", + "norm_policy" + ], + [ + "parameter", + "is_prefill" + ], + [ + "parameter", + "has_out_kv" + ], + [ + "parameter", + "upper_max" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "d8666e6d2e4bdc7ae0a1ae504018c4c85159b942f59d8fce125439b2583e3737", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 31903, + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_kernel.cu", + "sha256": "fc633c8fa81d145368fa175aef3be4734b5a0459ff34eab2b7f2b218ac8ede6b" + }, + { + "bytes": 12770, + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_binding.cu", + "sha256": "23ebf2522ad4038b25df09ed5ee8030cc2b8aee2a1c960a699ba0e7b7c943669" + } + ], + "closure_sha256": "462cd27c0e63d090604513ec1bddcf699aced8ea478231ca8fead9e5ed6da3cf", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed", + "launch": { + "block": [ + 96, + 1, + 1 + ], + "cluster": [ + 1, + 1, + 1 + ], + "cluster_scheduling_policy": "spread", + "cooperative": false, + "dynamic_smem_bytes": 0, + "max_persistent_clusters": null, + "persistent_ctas_per_sm": 1, + "use_pdl": false + }, + "module_ident": "cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_sm_100a", + "name": "cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "caf9bdb5445726a94bed62e062809b73bf1a05ec076dbc69e23dd53a38d77460", + "specialization": {}, + "stage": "fp8_e4m3", + "template": "flashinfer_blackwell_fused_qk_rmsnorm_rope_fp8_paged_kv_append" + }, + "source_build": { + "backend": "cuda_cpp", + "compile_options": [], + "cooperative": false, + "entry": "run", + "kind": "production_kernel_factory", + "persistent_ctas_per_sm": 1, + "physical_schedule": "resolved_final_ir", + "schedule_identity": "e9d3e97f4e1aeab628b5204d42ccd3cc7fc308460564ec9367384155e9d15cbb", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": {}, + "use_pdl": false, + "validate": false + }, + "specializations": {}, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_kernel.cu" + } + }, + { + "arch": "sm_103a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "page_indices" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "buffer", + "k_scale" + ], + [ + "buffer", + "v_scale" + ], + [ + "buffer", + "q_scale_inv" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "key_cache" + ], + [ + "buffer", + "value_cache" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_scale" + ], + [ + "buffer", + "split_k_flag" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "num_q_heads" + ], + [ + "parameter", + "num_kv_heads" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "max_pages_per_request" + ], + [ + "parameter", + "max_seqlen" + ], + [ + "parameter", + "max_seqlen_aligned" + ], + [ + "parameter", + "quant_policy" + ], + [ + "parameter", + "norm_policy" + ], + [ + "parameter", + "is_prefill" + ], + [ + "parameter", + "has_out_kv" + ], + [ + "parameter", + "upper_max" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "d8666e6d2e4bdc7ae0a1ae504018c4c85159b942f59d8fce125439b2583e3737", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 31903, + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_kernel.cu", + "sha256": "130e94500ec57ad30b1e53627080a49837a7c1877c0732f5917815759c0949b0" + }, + { + "bytes": 12770, + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_binding.cu", + "sha256": "47f5f9235a742595e0f5431b6eee5ef1582133a0c8a37ba798d342f6fb86bd3f" + } + ], + "closure_sha256": "45162914cde860df1cee131cce8fcc1b8e63edebc889d17f95c9238d2d35d953", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655", + "launch": { + "block": [ + 96, + 1, + 1 + ], + "cluster": [ + 1, + 1, + 1 + ], + "cluster_scheduling_policy": "spread", + "cooperative": false, + "dynamic_smem_bytes": 0, + "max_persistent_clusters": null, + "persistent_ctas_per_sm": 1, + "use_pdl": false + }, + "module_ident": "cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_sm_103a", + "name": "cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "f2873a52b7d6a3954655d4a956b1c60c930be8d36146c3492adb89bd9293eed7", + "specialization": {}, + "stage": "fp8_e4m3", + "template": "flashinfer_blackwell_fused_qk_rmsnorm_rope_fp8_paged_kv_append" + }, + "source_build": { + "backend": "cuda_cpp", + "compile_options": [], + "cooperative": false, + "entry": "run", + "kind": "production_kernel_factory", + "persistent_ctas_per_sm": 1, + "physical_schedule": "resolved_final_ir", + "schedule_identity": "e9d3e97f4e1aeab628b5204d42ccd3cc7fc308460564ec9367384155e9d15cbb", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": {}, + "use_pdl": false, + "validate": false + }, + "specializations": {}, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_kernel.cu" + } + }, + { + "arch": "sm_90a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "page_indices" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "buffer", + "k_scale" + ], + [ + "buffer", + "v_scale" + ], + [ + "buffer", + "q_scale_inv" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "key_cache" + ], + [ + "buffer", + "value_cache" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_scale" + ], + [ + "buffer", + "split_k_flag" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "num_q_heads" + ], + [ + "parameter", + "num_kv_heads" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "max_pages_per_request" + ], + [ + "parameter", + "max_seqlen" + ], + [ + "parameter", + "max_seqlen_aligned" + ], + [ + "parameter", + "quant_policy" + ], + [ + "parameter", + "norm_policy" + ], + [ + "parameter", + "is_prefill" + ], + [ + "parameter", + "has_out_kv" + ], + [ + "parameter", + "upper_max" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "d8666e6d2e4bdc7ae0a1ae504018c4c85159b942f59d8fce125439b2583e3737", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 31903, + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_kernel.cu", + "sha256": "695774941695ef5ba913db34bd803c794e5814cec6fd8424920430edefd017ae" + }, + { + "bytes": 12770, + "path": "csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_binding.cu", + "sha256": "7cb4035a255194f4d20eddcf62bed9f49f2a67e9a471f254e270c1b07d01f3ea" + } + ], + "closure_sha256": "d0803c188e10f118642158d38b406bef6e9586c9269ab71bcdde072b0c5527f7", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b", + "launch": { + "block": [ + 96, + 1, + 1 + ], + "cluster": [ + 1, + 1, + 1 + ], + "cluster_scheduling_policy": "spread", + "cooperative": false, + "dynamic_smem_bytes": 0, + "max_persistent_clusters": null, + "persistent_ctas_per_sm": 1, + "use_pdl": false + }, + "module_ident": "cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_sm_90a", + "name": "cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "31cb02d023f2ccb40a0b201e7a11c253a7922b81a710f29383dcc01538054ae5", + "specialization": {}, + "stage": "fp8_e4m3", + "template": "flashinfer_blackwell_fused_qk_rmsnorm_rope_fp8_paged_kv_append" + }, + "source_build": { + "backend": "cuda_cpp", + "compile_options": [], + "cooperative": false, + "entry": "run", + "kind": "production_kernel_factory", + "persistent_ctas_per_sm": 1, + "physical_schedule": "resolved_final_ir", + "schedule_identity": "e9d3e97f4e1aeab628b5204d42ccd3cc7fc308460564ec9367384155e9d15cbb", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": {}, + "use_pdl": false, + "validate": false + }, + "specializations": {}, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_kernel.cu" + } + } + ], + "name": "cake_fused_qk_rope_fp8_append", + "producer": "cake", + "producer_revision": null, + "public_namespace": "cake", + "schema": "cake.library_export.v5", + "sequences": [] +} diff --git a/csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_binding.cu b/csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_binding.cu new file mode 100644 index 00000000000..84f42a94c7b --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_binding.cu @@ -0,0 +1,237 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel +// 'kernel_cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed'. +#include +#include +#include +#include + +#include + +#include "tvm_ffi_utils.h" + +extern "C" __global__ void kernel_cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed( + const __nv_bfloat16* __restrict__ qkv, const float* __restrict__ cos_sin, + const int* __restrict__ seq_lens, const int* __restrict__ q_indptr, + const int* __restrict__ page_indices, const float* __restrict__ q_norm_weight, + const float* __restrict__ k_norm_weight, const float* __restrict__ k_scale, + const float* __restrict__ v_scale, const float* __restrict__ q_scale_inv, + uint8_t* __restrict__ out_q, uint8_t* key_cache, uint8_t* value_cache, uint8_t* out_k, + uint8_t* out_v, float* __restrict__ q_scale, int* __restrict__ split_k_flag, int num_rows, + int num_requests, int num_q_heads, int num_kv_heads, int page_size, int max_pages_per_request, + int max_seqlen, int max_seqlen_aligned, int quant_policy, int norm_policy, bool is_prefill, + bool has_out_kv, float upper_max); + +namespace cake_host_shim_069f92ac6d1683e5 { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_seq_lens, + TensorView arg_q_indptr, TensorView arg_page_indices, TensorView arg_q_norm_weight, + TensorView arg_k_norm_weight, TensorView arg_k_scale, TensorView arg_v_scale, + TensorView arg_q_scale_inv, TensorView arg_out_q, TensorView arg_key_cache, + TensorView arg_value_cache, TensorView arg_out_k, TensorView arg_out_v, + TensorView arg_q_scale, TensorView arg_split_k_flag, int64_t arg_num_rows, + int64_t arg_num_requests, int64_t arg_num_q_heads, int64_t arg_num_kv_heads, + int64_t arg_page_size, int64_t arg_max_pages_per_request, int64_t arg_max_seqlen, + int64_t arg_max_seqlen_aligned, int64_t arg_quant_policy, int64_t arg_norm_policy, + bool arg_is_prefill, bool arg_has_out_kv, double arg_upper_max, int64_t grid_x, + int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_qkv.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(arg_qkv, "qkv"); + check_dtype(arg_qkv, DLDataType{kDLBfloat, 16, 1}, "qkv"); + check_contiguous(arg_qkv, "qkv"); + check_cuda_tensor(arg_cos_sin, "cos_sin"); + check_dtype(arg_cos_sin, DLDataType{kDLFloat, 32, 1}, "cos_sin"); + check_contiguous(arg_cos_sin, "cos_sin"); + check_cuda_tensor(arg_seq_lens, "seq_lens"); + check_dtype(arg_seq_lens, DLDataType{kDLInt, 32, 1}, "seq_lens"); + check_contiguous(arg_seq_lens, "seq_lens"); + check_cuda_tensor(arg_q_indptr, "q_indptr"); + check_dtype(arg_q_indptr, DLDataType{kDLInt, 32, 1}, "q_indptr"); + check_contiguous(arg_q_indptr, "q_indptr"); + check_cuda_tensor(arg_page_indices, "page_indices"); + check_dtype(arg_page_indices, DLDataType{kDLInt, 32, 1}, "page_indices"); + check_contiguous(arg_page_indices, "page_indices"); + check_cuda_tensor(arg_q_norm_weight, "q_norm_weight"); + check_dtype(arg_q_norm_weight, DLDataType{kDLFloat, 32, 1}, "q_norm_weight"); + check_contiguous(arg_q_norm_weight, "q_norm_weight"); + check_cuda_tensor(arg_k_norm_weight, "k_norm_weight"); + check_dtype(arg_k_norm_weight, DLDataType{kDLFloat, 32, 1}, "k_norm_weight"); + check_contiguous(arg_k_norm_weight, "k_norm_weight"); + check_cuda_tensor(arg_k_scale, "k_scale"); + check_dtype(arg_k_scale, DLDataType{kDLFloat, 32, 1}, "k_scale"); + check_contiguous(arg_k_scale, "k_scale"); + check_cuda_tensor(arg_v_scale, "v_scale"); + check_dtype(arg_v_scale, DLDataType{kDLFloat, 32, 1}, "v_scale"); + check_contiguous(arg_v_scale, "v_scale"); + check_cuda_tensor(arg_q_scale_inv, "q_scale_inv"); + check_dtype(arg_q_scale_inv, DLDataType{kDLFloat, 32, 1}, "q_scale_inv"); + check_contiguous(arg_q_scale_inv, "q_scale_inv"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_key_cache, "key_cache"); + check_dtype(arg_key_cache, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "key_cache"); + check_contiguous(arg_key_cache, "key_cache"); + check_cuda_tensor(arg_value_cache, "value_cache"); + check_dtype(arg_value_cache, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "value_cache"); + check_contiguous(arg_value_cache, "value_cache"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + check_cuda_tensor(arg_q_scale, "q_scale"); + check_dtype(arg_q_scale, DLDataType{kDLFloat, 32, 1}, "q_scale"); + check_contiguous(arg_q_scale, "q_scale"); + check_cuda_tensor(arg_split_k_flag, "split_k_flag"); + check_dtype(arg_split_k_flag, DLDataType{kDLInt, 32, 1}, "split_k_flag"); + check_contiguous(arg_split_k_flag, "split_k_flag"); + TVM_FFI_CHECK(arg_num_rows >= -2147483648LL && arg_num_rows <= 2147483647LL, ValueError) + << "scalar 'num_rows' value " << arg_num_rows + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_requests >= -2147483648LL && arg_num_requests <= 2147483647LL, ValueError) + << "scalar 'num_requests' value " << arg_num_requests + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_kv_heads >= -2147483648LL && arg_num_kv_heads <= 2147483647LL, ValueError) + << "scalar 'num_kv_heads' value " << arg_num_kv_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK( + arg_max_pages_per_request >= -2147483648LL && arg_max_pages_per_request <= 2147483647LL, + ValueError) + << "scalar 'max_pages_per_request' value " << arg_max_pages_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_seqlen >= -2147483648LL && arg_max_seqlen <= 2147483647LL, ValueError) + << "scalar 'max_seqlen' value " << arg_max_seqlen + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_seqlen_aligned >= -2147483648LL && arg_max_seqlen_aligned <= 2147483647LL, + ValueError) + << "scalar 'max_seqlen_aligned' value " << arg_max_seqlen_aligned + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_quant_policy >= -2147483648LL && arg_quant_policy <= 2147483647LL, ValueError) + << "scalar 'quant_policy' value " << arg_quant_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_norm_policy >= -2147483648LL && arg_norm_policy <= 2147483647LL, ValueError) + << "scalar 'norm_policy' value " << arg_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_page_indices, arg_qkv, "page_indices", "qkv"); + check_same_device(arg_q_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "qkv"); + check_same_device(arg_k_scale, arg_qkv, "k_scale", "qkv"); + check_same_device(arg_v_scale, arg_qkv, "v_scale", "qkv"); + check_same_device(arg_q_scale_inv, arg_qkv, "q_scale_inv", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "qkv"); + check_same_device(arg_key_cache, arg_qkv, "key_cache", "qkv"); + check_same_device(arg_value_cache, arg_qkv, "value_cache", "qkv"); + check_same_device(arg_out_k, arg_qkv, "out_k", "qkv"); + check_same_device(arg_out_v, arg_qkv, "out_v", "qkv"); + check_same_device(arg_q_scale, arg_qkv, "q_scale", "qkv"); + check_same_device(arg_split_k_flag, arg_qkv, "split_k_flag", "qkv"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", " + << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + __nv_bfloat16* p_qkv = static_cast<__nv_bfloat16*>(arg_qkv.data_ptr()); + float* p_cos_sin = static_cast(arg_cos_sin.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_page_indices = static_cast(arg_page_indices.data_ptr()); + float* p_q_norm_weight = static_cast(arg_q_norm_weight.data_ptr()); + float* p_k_norm_weight = static_cast(arg_k_norm_weight.data_ptr()); + float* p_k_scale = static_cast(arg_k_scale.data_ptr()); + float* p_v_scale = static_cast(arg_v_scale.data_ptr()); + float* p_q_scale_inv = static_cast(arg_q_scale_inv.data_ptr()); + uint8_t* p_out_q = static_cast(arg_out_q.data_ptr()); + uint8_t* p_key_cache = static_cast(arg_key_cache.data_ptr()); + uint8_t* p_value_cache = static_cast(arg_value_cache.data_ptr()); + uint8_t* p_out_k = static_cast(arg_out_k.data_ptr()); + uint8_t* p_out_v = static_cast(arg_out_v.data_ptr()); + float* p_q_scale = static_cast(arg_q_scale.data_ptr()); + int* p_split_k_flag = static_cast(arg_split_k_flag.data_ptr()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_kv_heads = (int32_t)arg_num_kv_heads; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_max_pages_per_request = (int32_t)arg_max_pages_per_request; + int32_t v_max_seqlen = (int32_t)arg_max_seqlen; + int32_t v_max_seqlen_aligned = (int32_t)arg_max_seqlen_aligned; + int32_t v_quant_policy = (int32_t)arg_quant_policy; + int32_t v_norm_policy = (int32_t)arg_norm_policy; + bool v_is_prefill = (bool)arg_is_prefill; + bool v_has_out_kv = (bool)arg_has_out_kv; + float v_upper_max = (float)arg_upper_max; + void* kargs[] = {&p_qkv, + &p_cos_sin, + &p_seq_lens, + &p_q_indptr, + &p_page_indices, + &p_q_norm_weight, + &p_k_norm_weight, + &p_k_scale, + &p_v_scale, + &p_q_scale_inv, + &p_out_q, + &p_key_cache, + &p_value_cache, + &p_out_k, + &p_out_v, + &p_q_scale, + &p_split_k_flag, + &v_num_rows, + &v_num_requests, + &v_num_q_heads, + &v_num_kv_heads, + &v_page_size, + &v_max_pages_per_request, + &v_max_seqlen, + &v_max_seqlen_aligned, + &v_quant_policy, + &v_norm_policy, + &v_is_prefill, + &v_has_out_kv, + &v_upper_max}; + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(96u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed), + grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_069f92ac6d1683e5 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_069f92ac6d1683e5::Run); diff --git a/csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_kernel.cu b/csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_kernel.cu new file mode 100644 index 00000000000..ac9ce1b2ea7 --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/sm_100a/cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed_kernel.cu @@ -0,0 +1,805 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { + uint64_t opaque[16]; +}; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { + uint64_t opaque[16]; +} CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_META_OFF 0 +#define SMEM_META_STAGE_BYTES 16 +#define SMEM_META_STRIDE 16 +#define SMEM_TOTAL 0 +#define THREADS 96 + +#include + +extern "C" { + +__global__ __launch_bounds__(96) void kernel_cake_fused_qk_rope_fp8_append_caf9bdb5445726a94bed( + const __nv_bfloat16* __restrict__ qkv, const float* __restrict__ cos_sin, + const int* __restrict__ seq_lens, const int* __restrict__ q_indptr, + const int* __restrict__ page_indices, const float* __restrict__ q_norm_weight, + const float* __restrict__ k_norm_weight, const float* __restrict__ k_scale, + const float* __restrict__ v_scale, const float* __restrict__ q_scale_inv, + uint8_t* __restrict__ out_q, uint8_t* key_cache, uint8_t* value_cache, uint8_t* out_k, + uint8_t* out_v, float* __restrict__ q_scale, int* __restrict__ split_k_flag, int num_rows, + int num_requests, int num_q_heads, int num_kv_heads, int page_size, int max_pages_per_request, + int max_seqlen, int max_seqlen_aligned, int quant_policy, int norm_policy, bool is_prefill, + bool has_out_kv, float upper_max) { + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + __shared__ __align__(16) unsigned char smem_static_raw[16]; + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + int* meta = reinterpret_cast(smem_static_raw + 0); + const int meta_addr = (int)(unsigned long long)__cvta_generic_to_shared(meta); + + // === Task calls (dependency order) === + int tid_0 = tid; + int slot = tid_0 / 8; + int lane_1 = tid_0 % 8; + int grp_base = tid_0 % 32 - lane_1; + unsigned int grp_mask = (unsigned int)(255 << grp_base); + int kv_head = blockIdx.y; + int q_per_kv = num_q_heads / num_kv_heads; + int bid_1 = blockIdx.x; + bool is_row_cta = bid_1 < num_rows; + int row = ((is_row_cta) ? bid_1 : 0); + bool dyn_prefill = quant_policy == 1 && is_prefill; + int width = (num_q_heads + 2 * num_kv_heads) * 128; + bool is_q = slot < q_per_kv; + bool is_k = slot == 8; + bool is_v = slot == 9; + bool active = is_q || is_k || is_v; + int head_in_row = + ((is_q) ? kv_head * q_per_kv + slot + : ((is_k) ? num_q_heads + kv_head : num_q_heads + num_kv_heads + kv_head)); + int lane_lo = ((is_v) ? lane_1 * 16 : lane_1 * 8); + int lane_hi = ((is_v) ? lane_1 * 16 + 8 : 64 + lane_1 * 8); + unsigned int x_lo_w[4]; + unsigned int x_hi_w[4]; + float w_lo[8]; + float w_hi[8]; + float c_lo[8]; + float c_hi[8]; +#pragma unroll + for (int j = 0; j < 4; j++) { + x_lo_w[j] = 0; + x_hi_w[j] = 0; + } +#pragma unroll + for (int j_1 = 0; j_1 < 8; j_1++) { + w_lo[j_1] = 1.0f; + w_hi[j_1] = 1.0f; + c_lo[j_1] = 1.0f; + c_hi[j_1] = 0.0f; + } + if (is_row_cta && active) { + long long src_base = (long long)row * (long long)width + (long long)(head_in_row * 128); + { + const uint4* _vptr_0 = + reinterpret_cast(qkv + (src_base + (long long)lane_lo) + 0); + uint4* _vdst_0 = reinterpret_cast(&x_lo_w[0]); +#pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vdst_0[_blk].x), "=r"(_vdst_0[_blk].y), "=r"(_vdst_0[_blk].z), + "=r"(_vdst_0[_blk].w) + : "l"((const void*)(_vptr_0 + _blk)) + : "memory"); + } + } + { + const uint4* _vptr_1 = + reinterpret_cast(qkv + (src_base + (long long)lane_hi) + 0); + uint4* _vdst_1 = reinterpret_cast(&x_hi_w[0]); +#pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vdst_1[_blk].x), "=r"(_vdst_1[_blk].y), "=r"(_vdst_1[_blk].z), + "=r"(_vdst_1[_blk].w) + : "l"((const void*)(_vptr_1 + _blk)) + : "memory"); + } + } + if (norm_policy != 0 && (is_q || is_k)) { + if (is_q) { + { + unsigned _v4_2_0; + unsigned _v4_2_1; + unsigned _v4_2_2; + unsigned _v4_2_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_2_0), "=r"(_v4_2_1), "=r"(_v4_2_2), "=r"(_v4_2_3) + : "l"((const void*)(q_norm_weight + (lane_1 * 8) + (0))) + : "memory"); + w_lo[0 + 0] = __uint_as_float(_v4_2_0); + w_lo[0 + 1] = __uint_as_float(_v4_2_1); + w_lo[0 + 2] = __uint_as_float(_v4_2_2); + w_lo[0 + 3] = __uint_as_float(_v4_2_3); + } + { + unsigned _v4_3_0; + unsigned _v4_3_1; + unsigned _v4_3_2; + unsigned _v4_3_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_3_0), "=r"(_v4_3_1), "=r"(_v4_3_2), "=r"(_v4_3_3) + : "l"((const void*)(q_norm_weight + (lane_1 * 8 + 4) + (0))) + : "memory"); + w_lo[4 + 0] = __uint_as_float(_v4_3_0); + w_lo[4 + 1] = __uint_as_float(_v4_3_1); + w_lo[4 + 2] = __uint_as_float(_v4_3_2); + w_lo[4 + 3] = __uint_as_float(_v4_3_3); + } + { + unsigned _v4_4_0; + unsigned _v4_4_1; + unsigned _v4_4_2; + unsigned _v4_4_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_4_0), "=r"(_v4_4_1), "=r"(_v4_4_2), "=r"(_v4_4_3) + : "l"((const void*)(q_norm_weight + (64 + lane_1 * 8) + (0))) + : "memory"); + w_hi[0 + 0] = __uint_as_float(_v4_4_0); + w_hi[0 + 1] = __uint_as_float(_v4_4_1); + w_hi[0 + 2] = __uint_as_float(_v4_4_2); + w_hi[0 + 3] = __uint_as_float(_v4_4_3); + } + { + unsigned _v4_5_0; + unsigned _v4_5_1; + unsigned _v4_5_2; + unsigned _v4_5_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_5_0), "=r"(_v4_5_1), "=r"(_v4_5_2), "=r"(_v4_5_3) + : "l"((const void*)(q_norm_weight + (64 + lane_1 * 8 + 4) + (0))) + : "memory"); + w_hi[4 + 0] = __uint_as_float(_v4_5_0); + w_hi[4 + 1] = __uint_as_float(_v4_5_1); + w_hi[4 + 2] = __uint_as_float(_v4_5_2); + w_hi[4 + 3] = __uint_as_float(_v4_5_3); + } + } else { + { + unsigned _v4_6_0; + unsigned _v4_6_1; + unsigned _v4_6_2; + unsigned _v4_6_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_6_0), "=r"(_v4_6_1), "=r"(_v4_6_2), "=r"(_v4_6_3) + : "l"((const void*)(k_norm_weight + (lane_1 * 8) + (0))) + : "memory"); + w_lo[0 + 0] = __uint_as_float(_v4_6_0); + w_lo[0 + 1] = __uint_as_float(_v4_6_1); + w_lo[0 + 2] = __uint_as_float(_v4_6_2); + w_lo[0 + 3] = __uint_as_float(_v4_6_3); + } + { + unsigned _v4_7_0; + unsigned _v4_7_1; + unsigned _v4_7_2; + unsigned _v4_7_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_7_0), "=r"(_v4_7_1), "=r"(_v4_7_2), "=r"(_v4_7_3) + : "l"((const void*)(k_norm_weight + (lane_1 * 8 + 4) + (0))) + : "memory"); + w_lo[4 + 0] = __uint_as_float(_v4_7_0); + w_lo[4 + 1] = __uint_as_float(_v4_7_1); + w_lo[4 + 2] = __uint_as_float(_v4_7_2); + w_lo[4 + 3] = __uint_as_float(_v4_7_3); + } + { + unsigned _v4_8_0; + unsigned _v4_8_1; + unsigned _v4_8_2; + unsigned _v4_8_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_8_0), "=r"(_v4_8_1), "=r"(_v4_8_2), "=r"(_v4_8_3) + : "l"((const void*)(k_norm_weight + (64 + lane_1 * 8) + (0))) + : "memory"); + w_hi[0 + 0] = __uint_as_float(_v4_8_0); + w_hi[0 + 1] = __uint_as_float(_v4_8_1); + w_hi[0 + 2] = __uint_as_float(_v4_8_2); + w_hi[0 + 3] = __uint_as_float(_v4_8_3); + } + { + unsigned _v4_9_0; + unsigned _v4_9_1; + unsigned _v4_9_2; + unsigned _v4_9_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_9_0), "=r"(_v4_9_1), "=r"(_v4_9_2), "=r"(_v4_9_3) + : "l"((const void*)(k_norm_weight + (64 + lane_1 * 8 + 4) + (0))) + : "memory"); + w_hi[4 + 0] = __uint_as_float(_v4_9_0); + w_hi[4 + 1] = __uint_as_float(_v4_9_1); + w_hi[4 + 2] = __uint_as_float(_v4_9_2); + w_hi[4 + 3] = __uint_as_float(_v4_9_3); + } + } + } + } + float k_scale_v = 1.0f; + float v_scale_v = 1.0f; + float q_static = 1.0f; + if (is_row_cta) { + if (is_k) { + k_scale_v = k_scale[0]; + } + if (is_v) { + v_scale_v = v_scale[0]; + } + if (is_q && quant_policy == 2) { + q_static = q_scale_inv[0]; + } + } + int req_id = bid_1 - num_rows; + int clear_last_pos = -1; + if (!is_row_cta && req_id < num_requests) { + clear_last_pos = seq_lens[req_id] - 1; + } + bool any_bad = 0; +#pragma unroll 1 + for (int req = tid_0; req < num_requests; req += 96) { + int begin = q_indptr[req]; + int end = q_indptr[req + 1]; + int seq_len_r = seq_lens[req]; + bool bad = begin < 0 || end < begin || end > num_rows; + if (dyn_prefill && end - begin > max_seqlen) { + bad = 1; + } + if (bad) { + any_bad = 1; + } + if (is_row_cta) { + if (begin <= row && row < end) { + meta[0] = req; + meta[1] = row + seq_len_r - end; + meta[2] = row - begin; + } + } + } + if (tid_0 == 0) { + int first = q_indptr[0]; + int last_ip = q_indptr[num_requests]; + if (first != 0 || last_ip != num_rows) { + any_bad = 1; + } + } + uint32_t _cta_or_0 = __syncthreads_or(any_bad); + unsigned int invalid = _cta_or_0; + if (is_row_cta) { + if (invalid == 0 && active) { + int batch_i = meta[0]; + int pos = meta[1]; + int tok = meta[2]; + long long cache_row = 0; + if (is_q || is_k) { + long long cs_base = (long long)pos * 128 + (long long)(lane_1 * 8); + { + unsigned _v4_10_0; + unsigned _v4_10_1; + unsigned _v4_10_2; + unsigned _v4_10_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_10_0), "=r"(_v4_10_1), "=r"(_v4_10_2), "=r"(_v4_10_3) + : "l"((const void*)(cos_sin + cs_base + (0))) + : "memory"); + c_lo[0 + 0] = __uint_as_float(_v4_10_0); + c_lo[0 + 1] = __uint_as_float(_v4_10_1); + c_lo[0 + 2] = __uint_as_float(_v4_10_2); + c_lo[0 + 3] = __uint_as_float(_v4_10_3); + } + { + unsigned _v4_11_0; + unsigned _v4_11_1; + unsigned _v4_11_2; + unsigned _v4_11_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_11_0), "=r"(_v4_11_1), "=r"(_v4_11_2), "=r"(_v4_11_3) + : "l"((const void*)(cos_sin + (cs_base + 4) + (0))) + : "memory"); + c_lo[4 + 0] = __uint_as_float(_v4_11_0); + c_lo[4 + 1] = __uint_as_float(_v4_11_1); + c_lo[4 + 2] = __uint_as_float(_v4_11_2); + c_lo[4 + 3] = __uint_as_float(_v4_11_3); + } + { + unsigned _v4_12_0; + unsigned _v4_12_1; + unsigned _v4_12_2; + unsigned _v4_12_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_12_0), "=r"(_v4_12_1), "=r"(_v4_12_2), "=r"(_v4_12_3) + : "l"((const void*)(cos_sin + (cs_base + 64) + (0))) + : "memory"); + c_hi[0 + 0] = __uint_as_float(_v4_12_0); + c_hi[0 + 1] = __uint_as_float(_v4_12_1); + c_hi[0 + 2] = __uint_as_float(_v4_12_2); + c_hi[0 + 3] = __uint_as_float(_v4_12_3); + } + { + unsigned _v4_13_0; + unsigned _v4_13_1; + unsigned _v4_13_2; + unsigned _v4_13_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_13_0), "=r"(_v4_13_1), "=r"(_v4_13_2), "=r"(_v4_13_3) + : "l"((const void*)(cos_sin + (cs_base + 64 + 4) + (0))) + : "memory"); + c_hi[4 + 0] = __uint_as_float(_v4_13_0); + c_hi[4 + 1] = __uint_as_float(_v4_13_1); + c_hi[4 + 2] = __uint_as_float(_v4_13_2); + c_hi[4 + 3] = __uint_as_float(_v4_13_3); + } + } + if ((is_k || is_v) && !has_out_kv) { + int page_slot = pos / page_size; + int row_in_page = pos - page_slot * page_size; + int phys = page_indices[batch_i * max_pages_per_request + page_slot]; + cache_row = ((long long)phys * (long long)page_size + (long long)row_in_page) * + (long long)(num_kv_heads * 128) + + (long long)(kv_head * 128); + } + float x_lo[8]; + float x_hi[8]; +#pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&x_lo[_pair * 2])[0]), "=f"((&x_lo[_pair * 2])[1]) + : "r"(x_lo_w[_pair])); + } +#pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&x_hi[_pair * 2])[0]), "=f"((&x_hi[_pair * 2])[1]) + : "r"(x_hi_w[_pair])); + } + if (is_q || is_k) { + if (norm_policy == 2) { + float ss = 0.0f; +#pragma unroll + for (int j_2 = 0; j_2 < 8; j_2++) { + float _fma_0 = __fmaf_rn(x_lo[j_2], x_lo[j_2], ss); + ss = _fma_0; + float _fma_1 = __fmaf_rn(x_hi[j_2], x_hi[j_2], ss); + ss = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(grp_mask, ss, 4); + ss = ss + _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(grp_mask, ss, 2); + ss = ss + _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(grp_mask, ss, 1); + ss = ss + _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + 1e-06f); + float inv_rms = _rsqrt_0; +#pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x_lo[j_3] = x_lo[j_3] * (inv_rms * w_lo[j_3]); + x_hi[j_3] = x_hi[j_3] * (inv_rms * w_hi[j_3]); + } + } +#pragma unroll + for (int j_4 = 0; j_4 < 8; j_4++) { + float y_lo = x_lo[j_4] * c_lo[j_4] - x_hi[j_4] * c_hi[j_4]; + float y_hi = x_hi[j_4] * c_lo[j_4] + x_lo[j_4] * c_hi[j_4]; + x_lo[j_4] = y_lo; + x_hi[j_4] = y_hi; + } + if (norm_policy == 1) { + float ss1 = 0.0f; +#pragma unroll + for (int j_5 = 0; j_5 < 8; j_5++) { + float _fma_2 = __fmaf_rn(x_lo[j_5], x_lo[j_5], ss1); + ss1 = _fma_2; + float _fma_3 = __fmaf_rn(x_hi[j_5], x_hi[j_5], ss1); + ss1 = _fma_3; + } + float _shfl_xor_3 = __shfl_xor_sync(grp_mask, ss1, 4); + ss1 = ss1 + _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(grp_mask, ss1, 2); + ss1 = ss1 + _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(grp_mask, ss1, 1); + ss1 = ss1 + _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss1 * 0.0078125f + 1e-06f); + float inv_rms1 = _rsqrt_1; +#pragma unroll + for (int j_6 = 0; j_6 < 8; j_6++) { + x_lo[j_6] = x_lo[j_6] * (inv_rms1 * w_lo[j_6]); + x_hi[j_6] = x_hi[j_6] * (inv_rms1 * w_hi[j_6]); + } + } + } + float _rcp_0 = __frcp_rn(v_scale_v); + float mult = _rcp_0; + if (is_q) { + if (quant_policy == 1) { + float m = 1e-06f; +#pragma unroll + for (int j_7 = 0; j_7 < 8; j_7++) { + float _fabs_0 = fabsf(x_lo[j_7]); + float _fmax_0 = fmaxf(m, _fabs_0); + m = _fmax_0; + float _fabs_1 = fabsf(x_hi[j_7]); + float _fmax_1 = fmaxf(m, _fabs_1); + m = _fmax_1; + } + float _shfl_xor_6 = __shfl_xor_sync(grp_mask, m, 4); + float _fmax_2 = fmaxf(m, _shfl_xor_6); + m = _fmax_2; + float _shfl_xor_7 = __shfl_xor_sync(grp_mask, m, 2); + float _fmax_3 = fmaxf(m, _shfl_xor_7); + m = _fmax_3; + float _shfl_xor_8 = __shfl_xor_sync(grp_mask, m, 1); + float _fmax_4 = fmaxf(m, _shfl_xor_8); + m = _fmax_4; + float _fdiv_rn_0 = __fdiv_rn(m, upper_max); + float scale_val = _fdiv_rn_0; + if (lane_1 == 0) { + if (is_prefill) { + if (tok < max_seqlen_aligned) { + long long qs_idx = + ((long long)batch_i * (long long)num_q_heads + (long long)head_in_row) * + (long long)max_seqlen_aligned + + (long long)tok; + *(reinterpret_cast(q_scale + qs_idx) + (0)) = scale_val; + } + } else { + long long qs_idx2 = (long long)row * (long long)num_q_heads + (long long)head_in_row; + *(reinterpret_cast(q_scale + qs_idx2) + (0)) = scale_val; + } + } + float _rcp_1 = __frcp_rn(scale_val); + mult = _rcp_1; + } else { + mult = q_static; + } + } else if (is_k) { + float _rcp_2 = __frcp_rn(k_scale_v); + mult = _rcp_2; + } +#pragma unroll + for (int j_8 = 0; j_8 < 8; j_8++) { + x_lo[j_8] = x_lo[j_8] * mult; + x_hi[j_8] = x_hi[j_8] * mult; + } + long long row64 = row; + if (is_q) { + long long q_dst = (row64 * (long long)num_q_heads + (long long)head_in_row) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_q + (q_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_q + (q_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else if (is_k) { + if (has_out_kv) { + long long k_dst = (row64 * (long long)num_kv_heads + (long long)kv_head) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_k + (k_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_k + (k_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else { + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(key_cache + (cache_row + (long long)lane_lo)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(key_cache + (cache_row + (long long)lane_hi)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + } + } else { + if (has_out_kv) { + long long v_dst = (row64 * (long long)num_kv_heads + (long long)kv_head) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_v + (v_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_v + (v_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else { + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(value_cache + (cache_row + (long long)lane_lo)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(value_cache + (cache_row + (long long)lane_hi)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + } + } + } + } else if (req_id < num_requests) { + if (tid_0 == 0) { + int flag_val = ((invalid != 0) ? -1 : 0); + *(reinterpret_cast(split_k_flag + (req_id * num_kv_heads + kv_head)) + (0)) = flag_val; + } + if (invalid == 0) { + int last_pos = clear_last_pos; + if (last_pos >= 0) { + int last_slot = last_pos / page_size; + int zero_from = last_pos - last_slot * page_size + 1; + if (zero_from < page_size) { + int phys_last = page_indices[req_id * max_pages_per_request + last_slot]; + long long phys_last64 = phys_last; + unsigned int zero_w[4]; +#pragma unroll + for (int j_9 = 0; j_9 < 4; j_9++) { + zero_w[j_9] = 0; + } + int n_rows = page_size - zero_from; + int total_items = n_rows * 16; +#pragma unroll 1 + for (int item = tid_0; item < total_items; item += 96) { + int r = item / 16; + int sub = item - r * 16; + int chunk = sub % 8; + long long base = + (phys_last64 * (long long)page_size + (long long)zero_from + (long long)r) * + (long long)(num_kv_heads * 128) + + (long long)(kv_head * 128) + (long long)(chunk * 16); + if (sub < 8) { + reinterpret_cast(key_cache + base)[0] = reinterpret_cast(zero_w)[0]; + } else { + reinterpret_cast(value_cache + base)[0] = reinterpret_cast(zero_w)[0]; + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_binding.cu b/csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_binding.cu new file mode 100644 index 00000000000..0ae8a84844f --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_binding.cu @@ -0,0 +1,237 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel +// 'kernel_cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655'. +#include +#include +#include +#include + +#include + +#include "tvm_ffi_utils.h" + +extern "C" __global__ void kernel_cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655( + const __nv_bfloat16* __restrict__ qkv, const float* __restrict__ cos_sin, + const int* __restrict__ seq_lens, const int* __restrict__ q_indptr, + const int* __restrict__ page_indices, const float* __restrict__ q_norm_weight, + const float* __restrict__ k_norm_weight, const float* __restrict__ k_scale, + const float* __restrict__ v_scale, const float* __restrict__ q_scale_inv, + uint8_t* __restrict__ out_q, uint8_t* key_cache, uint8_t* value_cache, uint8_t* out_k, + uint8_t* out_v, float* __restrict__ q_scale, int* __restrict__ split_k_flag, int num_rows, + int num_requests, int num_q_heads, int num_kv_heads, int page_size, int max_pages_per_request, + int max_seqlen, int max_seqlen_aligned, int quant_policy, int norm_policy, bool is_prefill, + bool has_out_kv, float upper_max); + +namespace cake_host_shim_6b27c63b6a931207 { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_seq_lens, + TensorView arg_q_indptr, TensorView arg_page_indices, TensorView arg_q_norm_weight, + TensorView arg_k_norm_weight, TensorView arg_k_scale, TensorView arg_v_scale, + TensorView arg_q_scale_inv, TensorView arg_out_q, TensorView arg_key_cache, + TensorView arg_value_cache, TensorView arg_out_k, TensorView arg_out_v, + TensorView arg_q_scale, TensorView arg_split_k_flag, int64_t arg_num_rows, + int64_t arg_num_requests, int64_t arg_num_q_heads, int64_t arg_num_kv_heads, + int64_t arg_page_size, int64_t arg_max_pages_per_request, int64_t arg_max_seqlen, + int64_t arg_max_seqlen_aligned, int64_t arg_quant_policy, int64_t arg_norm_policy, + bool arg_is_prefill, bool arg_has_out_kv, double arg_upper_max, int64_t grid_x, + int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_qkv.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(arg_qkv, "qkv"); + check_dtype(arg_qkv, DLDataType{kDLBfloat, 16, 1}, "qkv"); + check_contiguous(arg_qkv, "qkv"); + check_cuda_tensor(arg_cos_sin, "cos_sin"); + check_dtype(arg_cos_sin, DLDataType{kDLFloat, 32, 1}, "cos_sin"); + check_contiguous(arg_cos_sin, "cos_sin"); + check_cuda_tensor(arg_seq_lens, "seq_lens"); + check_dtype(arg_seq_lens, DLDataType{kDLInt, 32, 1}, "seq_lens"); + check_contiguous(arg_seq_lens, "seq_lens"); + check_cuda_tensor(arg_q_indptr, "q_indptr"); + check_dtype(arg_q_indptr, DLDataType{kDLInt, 32, 1}, "q_indptr"); + check_contiguous(arg_q_indptr, "q_indptr"); + check_cuda_tensor(arg_page_indices, "page_indices"); + check_dtype(arg_page_indices, DLDataType{kDLInt, 32, 1}, "page_indices"); + check_contiguous(arg_page_indices, "page_indices"); + check_cuda_tensor(arg_q_norm_weight, "q_norm_weight"); + check_dtype(arg_q_norm_weight, DLDataType{kDLFloat, 32, 1}, "q_norm_weight"); + check_contiguous(arg_q_norm_weight, "q_norm_weight"); + check_cuda_tensor(arg_k_norm_weight, "k_norm_weight"); + check_dtype(arg_k_norm_weight, DLDataType{kDLFloat, 32, 1}, "k_norm_weight"); + check_contiguous(arg_k_norm_weight, "k_norm_weight"); + check_cuda_tensor(arg_k_scale, "k_scale"); + check_dtype(arg_k_scale, DLDataType{kDLFloat, 32, 1}, "k_scale"); + check_contiguous(arg_k_scale, "k_scale"); + check_cuda_tensor(arg_v_scale, "v_scale"); + check_dtype(arg_v_scale, DLDataType{kDLFloat, 32, 1}, "v_scale"); + check_contiguous(arg_v_scale, "v_scale"); + check_cuda_tensor(arg_q_scale_inv, "q_scale_inv"); + check_dtype(arg_q_scale_inv, DLDataType{kDLFloat, 32, 1}, "q_scale_inv"); + check_contiguous(arg_q_scale_inv, "q_scale_inv"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_key_cache, "key_cache"); + check_dtype(arg_key_cache, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "key_cache"); + check_contiguous(arg_key_cache, "key_cache"); + check_cuda_tensor(arg_value_cache, "value_cache"); + check_dtype(arg_value_cache, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "value_cache"); + check_contiguous(arg_value_cache, "value_cache"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + check_cuda_tensor(arg_q_scale, "q_scale"); + check_dtype(arg_q_scale, DLDataType{kDLFloat, 32, 1}, "q_scale"); + check_contiguous(arg_q_scale, "q_scale"); + check_cuda_tensor(arg_split_k_flag, "split_k_flag"); + check_dtype(arg_split_k_flag, DLDataType{kDLInt, 32, 1}, "split_k_flag"); + check_contiguous(arg_split_k_flag, "split_k_flag"); + TVM_FFI_CHECK(arg_num_rows >= -2147483648LL && arg_num_rows <= 2147483647LL, ValueError) + << "scalar 'num_rows' value " << arg_num_rows + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_requests >= -2147483648LL && arg_num_requests <= 2147483647LL, ValueError) + << "scalar 'num_requests' value " << arg_num_requests + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_kv_heads >= -2147483648LL && arg_num_kv_heads <= 2147483647LL, ValueError) + << "scalar 'num_kv_heads' value " << arg_num_kv_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK( + arg_max_pages_per_request >= -2147483648LL && arg_max_pages_per_request <= 2147483647LL, + ValueError) + << "scalar 'max_pages_per_request' value " << arg_max_pages_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_seqlen >= -2147483648LL && arg_max_seqlen <= 2147483647LL, ValueError) + << "scalar 'max_seqlen' value " << arg_max_seqlen + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_seqlen_aligned >= -2147483648LL && arg_max_seqlen_aligned <= 2147483647LL, + ValueError) + << "scalar 'max_seqlen_aligned' value " << arg_max_seqlen_aligned + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_quant_policy >= -2147483648LL && arg_quant_policy <= 2147483647LL, ValueError) + << "scalar 'quant_policy' value " << arg_quant_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_norm_policy >= -2147483648LL && arg_norm_policy <= 2147483647LL, ValueError) + << "scalar 'norm_policy' value " << arg_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_page_indices, arg_qkv, "page_indices", "qkv"); + check_same_device(arg_q_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "qkv"); + check_same_device(arg_k_scale, arg_qkv, "k_scale", "qkv"); + check_same_device(arg_v_scale, arg_qkv, "v_scale", "qkv"); + check_same_device(arg_q_scale_inv, arg_qkv, "q_scale_inv", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "qkv"); + check_same_device(arg_key_cache, arg_qkv, "key_cache", "qkv"); + check_same_device(arg_value_cache, arg_qkv, "value_cache", "qkv"); + check_same_device(arg_out_k, arg_qkv, "out_k", "qkv"); + check_same_device(arg_out_v, arg_qkv, "out_v", "qkv"); + check_same_device(arg_q_scale, arg_qkv, "q_scale", "qkv"); + check_same_device(arg_split_k_flag, arg_qkv, "split_k_flag", "qkv"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", " + << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + __nv_bfloat16* p_qkv = static_cast<__nv_bfloat16*>(arg_qkv.data_ptr()); + float* p_cos_sin = static_cast(arg_cos_sin.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_page_indices = static_cast(arg_page_indices.data_ptr()); + float* p_q_norm_weight = static_cast(arg_q_norm_weight.data_ptr()); + float* p_k_norm_weight = static_cast(arg_k_norm_weight.data_ptr()); + float* p_k_scale = static_cast(arg_k_scale.data_ptr()); + float* p_v_scale = static_cast(arg_v_scale.data_ptr()); + float* p_q_scale_inv = static_cast(arg_q_scale_inv.data_ptr()); + uint8_t* p_out_q = static_cast(arg_out_q.data_ptr()); + uint8_t* p_key_cache = static_cast(arg_key_cache.data_ptr()); + uint8_t* p_value_cache = static_cast(arg_value_cache.data_ptr()); + uint8_t* p_out_k = static_cast(arg_out_k.data_ptr()); + uint8_t* p_out_v = static_cast(arg_out_v.data_ptr()); + float* p_q_scale = static_cast(arg_q_scale.data_ptr()); + int* p_split_k_flag = static_cast(arg_split_k_flag.data_ptr()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_kv_heads = (int32_t)arg_num_kv_heads; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_max_pages_per_request = (int32_t)arg_max_pages_per_request; + int32_t v_max_seqlen = (int32_t)arg_max_seqlen; + int32_t v_max_seqlen_aligned = (int32_t)arg_max_seqlen_aligned; + int32_t v_quant_policy = (int32_t)arg_quant_policy; + int32_t v_norm_policy = (int32_t)arg_norm_policy; + bool v_is_prefill = (bool)arg_is_prefill; + bool v_has_out_kv = (bool)arg_has_out_kv; + float v_upper_max = (float)arg_upper_max; + void* kargs[] = {&p_qkv, + &p_cos_sin, + &p_seq_lens, + &p_q_indptr, + &p_page_indices, + &p_q_norm_weight, + &p_k_norm_weight, + &p_k_scale, + &p_v_scale, + &p_q_scale_inv, + &p_out_q, + &p_key_cache, + &p_value_cache, + &p_out_k, + &p_out_v, + &p_q_scale, + &p_split_k_flag, + &v_num_rows, + &v_num_requests, + &v_num_q_heads, + &v_num_kv_heads, + &v_page_size, + &v_max_pages_per_request, + &v_max_seqlen, + &v_max_seqlen_aligned, + &v_quant_policy, + &v_norm_policy, + &v_is_prefill, + &v_has_out_kv, + &v_upper_max}; + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(96u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655), + grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_6b27c63b6a931207 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6b27c63b6a931207::Run); diff --git a/csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_kernel.cu b/csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_kernel.cu new file mode 100644 index 00000000000..2df10b93cb0 --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/sm_103a/cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655_kernel.cu @@ -0,0 +1,805 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { + uint64_t opaque[16]; +}; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { + uint64_t opaque[16]; +} CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_META_OFF 0 +#define SMEM_META_STAGE_BYTES 16 +#define SMEM_META_STRIDE 16 +#define SMEM_TOTAL 0 +#define THREADS 96 + +#include + +extern "C" { + +__global__ __launch_bounds__(96) void kernel_cake_fused_qk_rope_fp8_append_f2873a52b7d6a3954655( + const __nv_bfloat16* __restrict__ qkv, const float* __restrict__ cos_sin, + const int* __restrict__ seq_lens, const int* __restrict__ q_indptr, + const int* __restrict__ page_indices, const float* __restrict__ q_norm_weight, + const float* __restrict__ k_norm_weight, const float* __restrict__ k_scale, + const float* __restrict__ v_scale, const float* __restrict__ q_scale_inv, + uint8_t* __restrict__ out_q, uint8_t* key_cache, uint8_t* value_cache, uint8_t* out_k, + uint8_t* out_v, float* __restrict__ q_scale, int* __restrict__ split_k_flag, int num_rows, + int num_requests, int num_q_heads, int num_kv_heads, int page_size, int max_pages_per_request, + int max_seqlen, int max_seqlen_aligned, int quant_policy, int norm_policy, bool is_prefill, + bool has_out_kv, float upper_max) { + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + __shared__ __align__(16) unsigned char smem_static_raw[16]; + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + int* meta = reinterpret_cast(smem_static_raw + 0); + const int meta_addr = (int)(unsigned long long)__cvta_generic_to_shared(meta); + + // === Task calls (dependency order) === + int tid_0 = tid; + int slot = tid_0 / 8; + int lane_1 = tid_0 % 8; + int grp_base = tid_0 % 32 - lane_1; + unsigned int grp_mask = (unsigned int)(255 << grp_base); + int kv_head = blockIdx.y; + int q_per_kv = num_q_heads / num_kv_heads; + int bid_1 = blockIdx.x; + bool is_row_cta = bid_1 < num_rows; + int row = ((is_row_cta) ? bid_1 : 0); + bool dyn_prefill = quant_policy == 1 && is_prefill; + int width = (num_q_heads + 2 * num_kv_heads) * 128; + bool is_q = slot < q_per_kv; + bool is_k = slot == 8; + bool is_v = slot == 9; + bool active = is_q || is_k || is_v; + int head_in_row = + ((is_q) ? kv_head * q_per_kv + slot + : ((is_k) ? num_q_heads + kv_head : num_q_heads + num_kv_heads + kv_head)); + int lane_lo = ((is_v) ? lane_1 * 16 : lane_1 * 8); + int lane_hi = ((is_v) ? lane_1 * 16 + 8 : 64 + lane_1 * 8); + unsigned int x_lo_w[4]; + unsigned int x_hi_w[4]; + float w_lo[8]; + float w_hi[8]; + float c_lo[8]; + float c_hi[8]; +#pragma unroll + for (int j = 0; j < 4; j++) { + x_lo_w[j] = 0; + x_hi_w[j] = 0; + } +#pragma unroll + for (int j_1 = 0; j_1 < 8; j_1++) { + w_lo[j_1] = 1.0f; + w_hi[j_1] = 1.0f; + c_lo[j_1] = 1.0f; + c_hi[j_1] = 0.0f; + } + if (is_row_cta && active) { + long long src_base = (long long)row * (long long)width + (long long)(head_in_row * 128); + { + const uint4* _vptr_0 = + reinterpret_cast(qkv + (src_base + (long long)lane_lo) + 0); + uint4* _vdst_0 = reinterpret_cast(&x_lo_w[0]); +#pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vdst_0[_blk].x), "=r"(_vdst_0[_blk].y), "=r"(_vdst_0[_blk].z), + "=r"(_vdst_0[_blk].w) + : "l"((const void*)(_vptr_0 + _blk)) + : "memory"); + } + } + { + const uint4* _vptr_1 = + reinterpret_cast(qkv + (src_base + (long long)lane_hi) + 0); + uint4* _vdst_1 = reinterpret_cast(&x_hi_w[0]); +#pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vdst_1[_blk].x), "=r"(_vdst_1[_blk].y), "=r"(_vdst_1[_blk].z), + "=r"(_vdst_1[_blk].w) + : "l"((const void*)(_vptr_1 + _blk)) + : "memory"); + } + } + if (norm_policy != 0 && (is_q || is_k)) { + if (is_q) { + { + unsigned _v4_2_0; + unsigned _v4_2_1; + unsigned _v4_2_2; + unsigned _v4_2_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_2_0), "=r"(_v4_2_1), "=r"(_v4_2_2), "=r"(_v4_2_3) + : "l"((const void*)(q_norm_weight + (lane_1 * 8) + (0))) + : "memory"); + w_lo[0 + 0] = __uint_as_float(_v4_2_0); + w_lo[0 + 1] = __uint_as_float(_v4_2_1); + w_lo[0 + 2] = __uint_as_float(_v4_2_2); + w_lo[0 + 3] = __uint_as_float(_v4_2_3); + } + { + unsigned _v4_3_0; + unsigned _v4_3_1; + unsigned _v4_3_2; + unsigned _v4_3_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_3_0), "=r"(_v4_3_1), "=r"(_v4_3_2), "=r"(_v4_3_3) + : "l"((const void*)(q_norm_weight + (lane_1 * 8 + 4) + (0))) + : "memory"); + w_lo[4 + 0] = __uint_as_float(_v4_3_0); + w_lo[4 + 1] = __uint_as_float(_v4_3_1); + w_lo[4 + 2] = __uint_as_float(_v4_3_2); + w_lo[4 + 3] = __uint_as_float(_v4_3_3); + } + { + unsigned _v4_4_0; + unsigned _v4_4_1; + unsigned _v4_4_2; + unsigned _v4_4_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_4_0), "=r"(_v4_4_1), "=r"(_v4_4_2), "=r"(_v4_4_3) + : "l"((const void*)(q_norm_weight + (64 + lane_1 * 8) + (0))) + : "memory"); + w_hi[0 + 0] = __uint_as_float(_v4_4_0); + w_hi[0 + 1] = __uint_as_float(_v4_4_1); + w_hi[0 + 2] = __uint_as_float(_v4_4_2); + w_hi[0 + 3] = __uint_as_float(_v4_4_3); + } + { + unsigned _v4_5_0; + unsigned _v4_5_1; + unsigned _v4_5_2; + unsigned _v4_5_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_5_0), "=r"(_v4_5_1), "=r"(_v4_5_2), "=r"(_v4_5_3) + : "l"((const void*)(q_norm_weight + (64 + lane_1 * 8 + 4) + (0))) + : "memory"); + w_hi[4 + 0] = __uint_as_float(_v4_5_0); + w_hi[4 + 1] = __uint_as_float(_v4_5_1); + w_hi[4 + 2] = __uint_as_float(_v4_5_2); + w_hi[4 + 3] = __uint_as_float(_v4_5_3); + } + } else { + { + unsigned _v4_6_0; + unsigned _v4_6_1; + unsigned _v4_6_2; + unsigned _v4_6_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_6_0), "=r"(_v4_6_1), "=r"(_v4_6_2), "=r"(_v4_6_3) + : "l"((const void*)(k_norm_weight + (lane_1 * 8) + (0))) + : "memory"); + w_lo[0 + 0] = __uint_as_float(_v4_6_0); + w_lo[0 + 1] = __uint_as_float(_v4_6_1); + w_lo[0 + 2] = __uint_as_float(_v4_6_2); + w_lo[0 + 3] = __uint_as_float(_v4_6_3); + } + { + unsigned _v4_7_0; + unsigned _v4_7_1; + unsigned _v4_7_2; + unsigned _v4_7_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_7_0), "=r"(_v4_7_1), "=r"(_v4_7_2), "=r"(_v4_7_3) + : "l"((const void*)(k_norm_weight + (lane_1 * 8 + 4) + (0))) + : "memory"); + w_lo[4 + 0] = __uint_as_float(_v4_7_0); + w_lo[4 + 1] = __uint_as_float(_v4_7_1); + w_lo[4 + 2] = __uint_as_float(_v4_7_2); + w_lo[4 + 3] = __uint_as_float(_v4_7_3); + } + { + unsigned _v4_8_0; + unsigned _v4_8_1; + unsigned _v4_8_2; + unsigned _v4_8_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_8_0), "=r"(_v4_8_1), "=r"(_v4_8_2), "=r"(_v4_8_3) + : "l"((const void*)(k_norm_weight + (64 + lane_1 * 8) + (0))) + : "memory"); + w_hi[0 + 0] = __uint_as_float(_v4_8_0); + w_hi[0 + 1] = __uint_as_float(_v4_8_1); + w_hi[0 + 2] = __uint_as_float(_v4_8_2); + w_hi[0 + 3] = __uint_as_float(_v4_8_3); + } + { + unsigned _v4_9_0; + unsigned _v4_9_1; + unsigned _v4_9_2; + unsigned _v4_9_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_9_0), "=r"(_v4_9_1), "=r"(_v4_9_2), "=r"(_v4_9_3) + : "l"((const void*)(k_norm_weight + (64 + lane_1 * 8 + 4) + (0))) + : "memory"); + w_hi[4 + 0] = __uint_as_float(_v4_9_0); + w_hi[4 + 1] = __uint_as_float(_v4_9_1); + w_hi[4 + 2] = __uint_as_float(_v4_9_2); + w_hi[4 + 3] = __uint_as_float(_v4_9_3); + } + } + } + } + float k_scale_v = 1.0f; + float v_scale_v = 1.0f; + float q_static = 1.0f; + if (is_row_cta) { + if (is_k) { + k_scale_v = k_scale[0]; + } + if (is_v) { + v_scale_v = v_scale[0]; + } + if (is_q && quant_policy == 2) { + q_static = q_scale_inv[0]; + } + } + int req_id = bid_1 - num_rows; + int clear_last_pos = -1; + if (!is_row_cta && req_id < num_requests) { + clear_last_pos = seq_lens[req_id] - 1; + } + bool any_bad = 0; +#pragma unroll 1 + for (int req = tid_0; req < num_requests; req += 96) { + int begin = q_indptr[req]; + int end = q_indptr[req + 1]; + int seq_len_r = seq_lens[req]; + bool bad = begin < 0 || end < begin || end > num_rows; + if (dyn_prefill && end - begin > max_seqlen) { + bad = 1; + } + if (bad) { + any_bad = 1; + } + if (is_row_cta) { + if (begin <= row && row < end) { + meta[0] = req; + meta[1] = row + seq_len_r - end; + meta[2] = row - begin; + } + } + } + if (tid_0 == 0) { + int first = q_indptr[0]; + int last_ip = q_indptr[num_requests]; + if (first != 0 || last_ip != num_rows) { + any_bad = 1; + } + } + uint32_t _cta_or_0 = __syncthreads_or(any_bad); + unsigned int invalid = _cta_or_0; + if (is_row_cta) { + if (invalid == 0 && active) { + int batch_i = meta[0]; + int pos = meta[1]; + int tok = meta[2]; + long long cache_row = 0; + if (is_q || is_k) { + long long cs_base = (long long)pos * 128 + (long long)(lane_1 * 8); + { + unsigned _v4_10_0; + unsigned _v4_10_1; + unsigned _v4_10_2; + unsigned _v4_10_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_10_0), "=r"(_v4_10_1), "=r"(_v4_10_2), "=r"(_v4_10_3) + : "l"((const void*)(cos_sin + cs_base + (0))) + : "memory"); + c_lo[0 + 0] = __uint_as_float(_v4_10_0); + c_lo[0 + 1] = __uint_as_float(_v4_10_1); + c_lo[0 + 2] = __uint_as_float(_v4_10_2); + c_lo[0 + 3] = __uint_as_float(_v4_10_3); + } + { + unsigned _v4_11_0; + unsigned _v4_11_1; + unsigned _v4_11_2; + unsigned _v4_11_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_11_0), "=r"(_v4_11_1), "=r"(_v4_11_2), "=r"(_v4_11_3) + : "l"((const void*)(cos_sin + (cs_base + 4) + (0))) + : "memory"); + c_lo[4 + 0] = __uint_as_float(_v4_11_0); + c_lo[4 + 1] = __uint_as_float(_v4_11_1); + c_lo[4 + 2] = __uint_as_float(_v4_11_2); + c_lo[4 + 3] = __uint_as_float(_v4_11_3); + } + { + unsigned _v4_12_0; + unsigned _v4_12_1; + unsigned _v4_12_2; + unsigned _v4_12_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_12_0), "=r"(_v4_12_1), "=r"(_v4_12_2), "=r"(_v4_12_3) + : "l"((const void*)(cos_sin + (cs_base + 64) + (0))) + : "memory"); + c_hi[0 + 0] = __uint_as_float(_v4_12_0); + c_hi[0 + 1] = __uint_as_float(_v4_12_1); + c_hi[0 + 2] = __uint_as_float(_v4_12_2); + c_hi[0 + 3] = __uint_as_float(_v4_12_3); + } + { + unsigned _v4_13_0; + unsigned _v4_13_1; + unsigned _v4_13_2; + unsigned _v4_13_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_13_0), "=r"(_v4_13_1), "=r"(_v4_13_2), "=r"(_v4_13_3) + : "l"((const void*)(cos_sin + (cs_base + 64 + 4) + (0))) + : "memory"); + c_hi[4 + 0] = __uint_as_float(_v4_13_0); + c_hi[4 + 1] = __uint_as_float(_v4_13_1); + c_hi[4 + 2] = __uint_as_float(_v4_13_2); + c_hi[4 + 3] = __uint_as_float(_v4_13_3); + } + } + if ((is_k || is_v) && !has_out_kv) { + int page_slot = pos / page_size; + int row_in_page = pos - page_slot * page_size; + int phys = page_indices[batch_i * max_pages_per_request + page_slot]; + cache_row = ((long long)phys * (long long)page_size + (long long)row_in_page) * + (long long)(num_kv_heads * 128) + + (long long)(kv_head * 128); + } + float x_lo[8]; + float x_hi[8]; +#pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&x_lo[_pair * 2])[0]), "=f"((&x_lo[_pair * 2])[1]) + : "r"(x_lo_w[_pair])); + } +#pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&x_hi[_pair * 2])[0]), "=f"((&x_hi[_pair * 2])[1]) + : "r"(x_hi_w[_pair])); + } + if (is_q || is_k) { + if (norm_policy == 2) { + float ss = 0.0f; +#pragma unroll + for (int j_2 = 0; j_2 < 8; j_2++) { + float _fma_0 = __fmaf_rn(x_lo[j_2], x_lo[j_2], ss); + ss = _fma_0; + float _fma_1 = __fmaf_rn(x_hi[j_2], x_hi[j_2], ss); + ss = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(grp_mask, ss, 4); + ss = ss + _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(grp_mask, ss, 2); + ss = ss + _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(grp_mask, ss, 1); + ss = ss + _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + 1e-06f); + float inv_rms = _rsqrt_0; +#pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x_lo[j_3] = x_lo[j_3] * (inv_rms * w_lo[j_3]); + x_hi[j_3] = x_hi[j_3] * (inv_rms * w_hi[j_3]); + } + } +#pragma unroll + for (int j_4 = 0; j_4 < 8; j_4++) { + float y_lo = x_lo[j_4] * c_lo[j_4] - x_hi[j_4] * c_hi[j_4]; + float y_hi = x_hi[j_4] * c_lo[j_4] + x_lo[j_4] * c_hi[j_4]; + x_lo[j_4] = y_lo; + x_hi[j_4] = y_hi; + } + if (norm_policy == 1) { + float ss1 = 0.0f; +#pragma unroll + for (int j_5 = 0; j_5 < 8; j_5++) { + float _fma_2 = __fmaf_rn(x_lo[j_5], x_lo[j_5], ss1); + ss1 = _fma_2; + float _fma_3 = __fmaf_rn(x_hi[j_5], x_hi[j_5], ss1); + ss1 = _fma_3; + } + float _shfl_xor_3 = __shfl_xor_sync(grp_mask, ss1, 4); + ss1 = ss1 + _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(grp_mask, ss1, 2); + ss1 = ss1 + _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(grp_mask, ss1, 1); + ss1 = ss1 + _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss1 * 0.0078125f + 1e-06f); + float inv_rms1 = _rsqrt_1; +#pragma unroll + for (int j_6 = 0; j_6 < 8; j_6++) { + x_lo[j_6] = x_lo[j_6] * (inv_rms1 * w_lo[j_6]); + x_hi[j_6] = x_hi[j_6] * (inv_rms1 * w_hi[j_6]); + } + } + } + float _rcp_0 = __frcp_rn(v_scale_v); + float mult = _rcp_0; + if (is_q) { + if (quant_policy == 1) { + float m = 1e-06f; +#pragma unroll + for (int j_7 = 0; j_7 < 8; j_7++) { + float _fabs_0 = fabsf(x_lo[j_7]); + float _fmax_0 = fmaxf(m, _fabs_0); + m = _fmax_0; + float _fabs_1 = fabsf(x_hi[j_7]); + float _fmax_1 = fmaxf(m, _fabs_1); + m = _fmax_1; + } + float _shfl_xor_6 = __shfl_xor_sync(grp_mask, m, 4); + float _fmax_2 = fmaxf(m, _shfl_xor_6); + m = _fmax_2; + float _shfl_xor_7 = __shfl_xor_sync(grp_mask, m, 2); + float _fmax_3 = fmaxf(m, _shfl_xor_7); + m = _fmax_3; + float _shfl_xor_8 = __shfl_xor_sync(grp_mask, m, 1); + float _fmax_4 = fmaxf(m, _shfl_xor_8); + m = _fmax_4; + float _fdiv_rn_0 = __fdiv_rn(m, upper_max); + float scale_val = _fdiv_rn_0; + if (lane_1 == 0) { + if (is_prefill) { + if (tok < max_seqlen_aligned) { + long long qs_idx = + ((long long)batch_i * (long long)num_q_heads + (long long)head_in_row) * + (long long)max_seqlen_aligned + + (long long)tok; + *(reinterpret_cast(q_scale + qs_idx) + (0)) = scale_val; + } + } else { + long long qs_idx2 = (long long)row * (long long)num_q_heads + (long long)head_in_row; + *(reinterpret_cast(q_scale + qs_idx2) + (0)) = scale_val; + } + } + float _rcp_1 = __frcp_rn(scale_val); + mult = _rcp_1; + } else { + mult = q_static; + } + } else if (is_k) { + float _rcp_2 = __frcp_rn(k_scale_v); + mult = _rcp_2; + } +#pragma unroll + for (int j_8 = 0; j_8 < 8; j_8++) { + x_lo[j_8] = x_lo[j_8] * mult; + x_hi[j_8] = x_hi[j_8] * mult; + } + long long row64 = row; + if (is_q) { + long long q_dst = (row64 * (long long)num_q_heads + (long long)head_in_row) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_q + (q_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_q + (q_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else if (is_k) { + if (has_out_kv) { + long long k_dst = (row64 * (long long)num_kv_heads + (long long)kv_head) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_k + (k_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_k + (k_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else { + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(key_cache + (cache_row + (long long)lane_lo)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(key_cache + (cache_row + (long long)lane_hi)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + } + } else { + if (has_out_kv) { + long long v_dst = (row64 * (long long)num_kv_heads + (long long)kv_head) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_v + (v_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_v + (v_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else { + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(value_cache + (cache_row + (long long)lane_lo)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(value_cache + (cache_row + (long long)lane_hi)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + } + } + } + } else if (req_id < num_requests) { + if (tid_0 == 0) { + int flag_val = ((invalid != 0) ? -1 : 0); + *(reinterpret_cast(split_k_flag + (req_id * num_kv_heads + kv_head)) + (0)) = flag_val; + } + if (invalid == 0) { + int last_pos = clear_last_pos; + if (last_pos >= 0) { + int last_slot = last_pos / page_size; + int zero_from = last_pos - last_slot * page_size + 1; + if (zero_from < page_size) { + int phys_last = page_indices[req_id * max_pages_per_request + last_slot]; + long long phys_last64 = phys_last; + unsigned int zero_w[4]; +#pragma unroll + for (int j_9 = 0; j_9 < 4; j_9++) { + zero_w[j_9] = 0; + } + int n_rows = page_size - zero_from; + int total_items = n_rows * 16; +#pragma unroll 1 + for (int item = tid_0; item < total_items; item += 96) { + int r = item / 16; + int sub = item - r * 16; + int chunk = sub % 8; + long long base = + (phys_last64 * (long long)page_size + (long long)zero_from + (long long)r) * + (long long)(num_kv_heads * 128) + + (long long)(kv_head * 128) + (long long)(chunk * 16); + if (sub < 8) { + reinterpret_cast(key_cache + base)[0] = reinterpret_cast(zero_w)[0]; + } else { + reinterpret_cast(value_cache + base)[0] = reinterpret_cast(zero_w)[0]; + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_binding.cu b/csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_binding.cu new file mode 100644 index 00000000000..028d33620ff --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_binding.cu @@ -0,0 +1,237 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +// Generated by the Cake source exporter — do not edit. +// tvm-ffi direct-source launcher for Cake kernel +// 'kernel_cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b'. +#include +#include +#include +#include + +#include + +#include "tvm_ffi_utils.h" + +extern "C" __global__ void kernel_cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b( + const __nv_bfloat16* __restrict__ qkv, const float* __restrict__ cos_sin, + const int* __restrict__ seq_lens, const int* __restrict__ q_indptr, + const int* __restrict__ page_indices, const float* __restrict__ q_norm_weight, + const float* __restrict__ k_norm_weight, const float* __restrict__ k_scale, + const float* __restrict__ v_scale, const float* __restrict__ q_scale_inv, + uint8_t* __restrict__ out_q, uint8_t* key_cache, uint8_t* value_cache, uint8_t* out_k, + uint8_t* out_v, float* __restrict__ q_scale, int* __restrict__ split_k_flag, int num_rows, + int num_requests, int num_q_heads, int num_kv_heads, int page_size, int max_pages_per_request, + int max_seqlen, int max_seqlen_aligned, int quant_policy, int norm_policy, bool is_prefill, + bool has_out_kv, float upper_max); + +namespace cake_host_shim_6d1d6d39fd4fdd96 { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_seq_lens, + TensorView arg_q_indptr, TensorView arg_page_indices, TensorView arg_q_norm_weight, + TensorView arg_k_norm_weight, TensorView arg_k_scale, TensorView arg_v_scale, + TensorView arg_q_scale_inv, TensorView arg_out_q, TensorView arg_key_cache, + TensorView arg_value_cache, TensorView arg_out_k, TensorView arg_out_v, + TensorView arg_q_scale, TensorView arg_split_k_flag, int64_t arg_num_rows, + int64_t arg_num_requests, int64_t arg_num_q_heads, int64_t arg_num_kv_heads, + int64_t arg_page_size, int64_t arg_max_pages_per_request, int64_t arg_max_seqlen, + int64_t arg_max_seqlen_aligned, int64_t arg_quant_policy, int64_t arg_norm_policy, + bool arg_is_prefill, bool arg_has_out_kv, double arg_upper_max, int64_t grid_x, + int64_t grid_y, int64_t grid_z) { + DLDevice dev = arg_qkv.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(arg_qkv, "qkv"); + check_dtype(arg_qkv, DLDataType{kDLBfloat, 16, 1}, "qkv"); + check_contiguous(arg_qkv, "qkv"); + check_cuda_tensor(arg_cos_sin, "cos_sin"); + check_dtype(arg_cos_sin, DLDataType{kDLFloat, 32, 1}, "cos_sin"); + check_contiguous(arg_cos_sin, "cos_sin"); + check_cuda_tensor(arg_seq_lens, "seq_lens"); + check_dtype(arg_seq_lens, DLDataType{kDLInt, 32, 1}, "seq_lens"); + check_contiguous(arg_seq_lens, "seq_lens"); + check_cuda_tensor(arg_q_indptr, "q_indptr"); + check_dtype(arg_q_indptr, DLDataType{kDLInt, 32, 1}, "q_indptr"); + check_contiguous(arg_q_indptr, "q_indptr"); + check_cuda_tensor(arg_page_indices, "page_indices"); + check_dtype(arg_page_indices, DLDataType{kDLInt, 32, 1}, "page_indices"); + check_contiguous(arg_page_indices, "page_indices"); + check_cuda_tensor(arg_q_norm_weight, "q_norm_weight"); + check_dtype(arg_q_norm_weight, DLDataType{kDLFloat, 32, 1}, "q_norm_weight"); + check_contiguous(arg_q_norm_weight, "q_norm_weight"); + check_cuda_tensor(arg_k_norm_weight, "k_norm_weight"); + check_dtype(arg_k_norm_weight, DLDataType{kDLFloat, 32, 1}, "k_norm_weight"); + check_contiguous(arg_k_norm_weight, "k_norm_weight"); + check_cuda_tensor(arg_k_scale, "k_scale"); + check_dtype(arg_k_scale, DLDataType{kDLFloat, 32, 1}, "k_scale"); + check_contiguous(arg_k_scale, "k_scale"); + check_cuda_tensor(arg_v_scale, "v_scale"); + check_dtype(arg_v_scale, DLDataType{kDLFloat, 32, 1}, "v_scale"); + check_contiguous(arg_v_scale, "v_scale"); + check_cuda_tensor(arg_q_scale_inv, "q_scale_inv"); + check_dtype(arg_q_scale_inv, DLDataType{kDLFloat, 32, 1}, "q_scale_inv"); + check_contiguous(arg_q_scale_inv, "q_scale_inv"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_key_cache, "key_cache"); + check_dtype(arg_key_cache, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "key_cache"); + check_contiguous(arg_key_cache, "key_cache"); + check_cuda_tensor(arg_value_cache, "value_cache"); + check_dtype(arg_value_cache, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "value_cache"); + check_contiguous(arg_value_cache, "value_cache"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLFloat8_e4m3fn, 8, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + check_cuda_tensor(arg_q_scale, "q_scale"); + check_dtype(arg_q_scale, DLDataType{kDLFloat, 32, 1}, "q_scale"); + check_contiguous(arg_q_scale, "q_scale"); + check_cuda_tensor(arg_split_k_flag, "split_k_flag"); + check_dtype(arg_split_k_flag, DLDataType{kDLInt, 32, 1}, "split_k_flag"); + check_contiguous(arg_split_k_flag, "split_k_flag"); + TVM_FFI_CHECK(arg_num_rows >= -2147483648LL && arg_num_rows <= 2147483647LL, ValueError) + << "scalar 'num_rows' value " << arg_num_rows + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_requests >= -2147483648LL && arg_num_requests <= 2147483647LL, ValueError) + << "scalar 'num_requests' value " << arg_num_requests + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_q_heads >= -2147483648LL && arg_num_q_heads <= 2147483647LL, ValueError) + << "scalar 'num_q_heads' value " << arg_num_q_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_num_kv_heads >= -2147483648LL && arg_num_kv_heads <= 2147483647LL, ValueError) + << "scalar 'num_kv_heads' value " << arg_num_kv_heads + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_page_size >= -2147483648LL && arg_page_size <= 2147483647LL, ValueError) + << "scalar 'page_size' value " << arg_page_size + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK( + arg_max_pages_per_request >= -2147483648LL && arg_max_pages_per_request <= 2147483647LL, + ValueError) + << "scalar 'max_pages_per_request' value " << arg_max_pages_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_seqlen >= -2147483648LL && arg_max_seqlen <= 2147483647LL, ValueError) + << "scalar 'max_seqlen' value " << arg_max_seqlen + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_max_seqlen_aligned >= -2147483648LL && arg_max_seqlen_aligned <= 2147483647LL, + ValueError) + << "scalar 'max_seqlen_aligned' value " << arg_max_seqlen_aligned + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_quant_policy >= -2147483648LL && arg_quant_policy <= 2147483647LL, ValueError) + << "scalar 'quant_policy' value " << arg_quant_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_norm_policy >= -2147483648LL && arg_norm_policy <= 2147483647LL, ValueError) + << "scalar 'norm_policy' value " << arg_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_page_indices, arg_qkv, "page_indices", "qkv"); + check_same_device(arg_q_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "qkv"); + check_same_device(arg_k_scale, arg_qkv, "k_scale", "qkv"); + check_same_device(arg_v_scale, arg_qkv, "v_scale", "qkv"); + check_same_device(arg_q_scale_inv, arg_qkv, "q_scale_inv", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "qkv"); + check_same_device(arg_key_cache, arg_qkv, "key_cache", "qkv"); + check_same_device(arg_value_cache, arg_qkv, "value_cache", "qkv"); + check_same_device(arg_out_k, arg_qkv, "out_k", "qkv"); + check_same_device(arg_out_v, arg_qkv, "out_v", "qkv"); + check_same_device(arg_q_scale, arg_qkv, "q_scale", "qkv"); + check_same_device(arg_split_k_flag, arg_qkv, "split_k_flag", "qkv"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << "launch grid dimensions must be positive, got (" << grid_x << ", " << grid_y << ", " + << grid_z << ")"; + + cudaStream_t stream = (cudaStream_t)TVMFFIEnvGetStream(dev.device_type, dev.device_id); + __nv_bfloat16* p_qkv = static_cast<__nv_bfloat16*>(arg_qkv.data_ptr()); + float* p_cos_sin = static_cast(arg_cos_sin.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_page_indices = static_cast(arg_page_indices.data_ptr()); + float* p_q_norm_weight = static_cast(arg_q_norm_weight.data_ptr()); + float* p_k_norm_weight = static_cast(arg_k_norm_weight.data_ptr()); + float* p_k_scale = static_cast(arg_k_scale.data_ptr()); + float* p_v_scale = static_cast(arg_v_scale.data_ptr()); + float* p_q_scale_inv = static_cast(arg_q_scale_inv.data_ptr()); + uint8_t* p_out_q = static_cast(arg_out_q.data_ptr()); + uint8_t* p_key_cache = static_cast(arg_key_cache.data_ptr()); + uint8_t* p_value_cache = static_cast(arg_value_cache.data_ptr()); + uint8_t* p_out_k = static_cast(arg_out_k.data_ptr()); + uint8_t* p_out_v = static_cast(arg_out_v.data_ptr()); + float* p_q_scale = static_cast(arg_q_scale.data_ptr()); + int* p_split_k_flag = static_cast(arg_split_k_flag.data_ptr()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_num_q_heads = (int32_t)arg_num_q_heads; + int32_t v_num_kv_heads = (int32_t)arg_num_kv_heads; + int32_t v_page_size = (int32_t)arg_page_size; + int32_t v_max_pages_per_request = (int32_t)arg_max_pages_per_request; + int32_t v_max_seqlen = (int32_t)arg_max_seqlen; + int32_t v_max_seqlen_aligned = (int32_t)arg_max_seqlen_aligned; + int32_t v_quant_policy = (int32_t)arg_quant_policy; + int32_t v_norm_policy = (int32_t)arg_norm_policy; + bool v_is_prefill = (bool)arg_is_prefill; + bool v_has_out_kv = (bool)arg_has_out_kv; + float v_upper_max = (float)arg_upper_max; + void* kargs[] = {&p_qkv, + &p_cos_sin, + &p_seq_lens, + &p_q_indptr, + &p_page_indices, + &p_q_norm_weight, + &p_k_norm_weight, + &p_k_scale, + &p_v_scale, + &p_q_scale_inv, + &p_out_q, + &p_key_cache, + &p_value_cache, + &p_out_k, + &p_out_v, + &p_q_scale, + &p_split_k_flag, + &v_num_rows, + &v_num_requests, + &v_num_q_heads, + &v_num_kv_heads, + &v_page_size, + &v_max_pages_per_request, + &v_max_seqlen, + &v_max_seqlen_aligned, + &v_quant_policy, + &v_norm_policy, + &v_is_prefill, + &v_has_out_kv, + &v_upper_max}; + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(96u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b), + grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_6d1d6d39fd4fdd96 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_6d1d6d39fd4fdd96::Run); diff --git a/csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_kernel.cu b/csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_kernel.cu new file mode 100644 index 00000000000..3e3e6d0f288 --- /dev/null +++ b/csrc/cake_fused_qk_rope_fp8_append/sm_90a/cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b_kernel.cu @@ -0,0 +1,805 @@ +/* + * Copyright (c) 2026 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +typedef signed char int8_t; +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +#if defined(__CUDACC_RTC__) +typedef unsigned long long uint64_t; +#else +typedef unsigned long uint64_t; +#endif +static_assert(sizeof(uint64_t) == 8, "Cake requires an LP64 CUDA host ABI"); +typedef signed int int32_t; +typedef short int int16_t; +struct __align__(64) CakeTensorMap64 { + uint64_t opaque[16]; +}; +static_assert(sizeof(CakeTensorMap64) == 128, "64-aligned tensor-map ABI size"); +static_assert(alignof(CakeTensorMap64) == 64, "64-aligned tensor-map ABI alignment"); + +#if defined(__CUDACC_RTC__) +typedef struct __align__(128) { + uint64_t opaque[16]; +} CUtensorMap; +#else +#include +#endif + +static_assert(sizeof(CUtensorMap) == 128, "CUtensorMap CUDA ABI must be 128 bytes"); +#include +#include + +__device__ __forceinline__ int make_warp_uniform(int x) { + int result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1F, 0xFFFFFFFF;" : "=r"(result) : "r"(x)); + return result; +} + +#define CAKE_INF CUDART_INF_F +#define NUM_MAIN_STAGES 1 +#define SMEM_META_OFF 0 +#define SMEM_META_STAGE_BYTES 16 +#define SMEM_META_STRIDE 16 +#define SMEM_TOTAL 0 +#define THREADS 96 + +#include + +extern "C" { + +__global__ __launch_bounds__(96) void kernel_cake_fused_qk_rope_fp8_append_31cb02d023f2ccb40a0b( + const __nv_bfloat16* __restrict__ qkv, const float* __restrict__ cos_sin, + const int* __restrict__ seq_lens, const int* __restrict__ q_indptr, + const int* __restrict__ page_indices, const float* __restrict__ q_norm_weight, + const float* __restrict__ k_norm_weight, const float* __restrict__ k_scale, + const float* __restrict__ v_scale, const float* __restrict__ q_scale_inv, + uint8_t* __restrict__ out_q, uint8_t* key_cache, uint8_t* value_cache, uint8_t* out_k, + uint8_t* out_v, float* __restrict__ q_scale, int* __restrict__ split_k_flag, int num_rows, + int num_requests, int num_q_heads, int num_kv_heads, int page_size, int max_pages_per_request, + int max_seqlen, int max_seqlen_aligned, int quant_policy, int norm_policy, bool is_prefill, + bool has_out_kv, float upper_max) { + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + __shared__ __align__(16) unsigned char smem_static_raw[16]; + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // Kernel setup ops + int* meta = reinterpret_cast(smem_static_raw + 0); + const int meta_addr = (int)(unsigned long long)__cvta_generic_to_shared(meta); + + // === Task calls (dependency order) === + int tid_0 = tid; + int slot = tid_0 / 8; + int lane_1 = tid_0 % 8; + int grp_base = tid_0 % 32 - lane_1; + unsigned int grp_mask = (unsigned int)(255 << grp_base); + int kv_head = blockIdx.y; + int q_per_kv = num_q_heads / num_kv_heads; + int bid_1 = blockIdx.x; + bool is_row_cta = bid_1 < num_rows; + int row = ((is_row_cta) ? bid_1 : 0); + bool dyn_prefill = quant_policy == 1 && is_prefill; + int width = (num_q_heads + 2 * num_kv_heads) * 128; + bool is_q = slot < q_per_kv; + bool is_k = slot == 8; + bool is_v = slot == 9; + bool active = is_q || is_k || is_v; + int head_in_row = + ((is_q) ? kv_head * q_per_kv + slot + : ((is_k) ? num_q_heads + kv_head : num_q_heads + num_kv_heads + kv_head)); + int lane_lo = ((is_v) ? lane_1 * 16 : lane_1 * 8); + int lane_hi = ((is_v) ? lane_1 * 16 + 8 : 64 + lane_1 * 8); + unsigned int x_lo_w[4]; + unsigned int x_hi_w[4]; + float w_lo[8]; + float w_hi[8]; + float c_lo[8]; + float c_hi[8]; +#pragma unroll + for (int j = 0; j < 4; j++) { + x_lo_w[j] = 0; + x_hi_w[j] = 0; + } +#pragma unroll + for (int j_1 = 0; j_1 < 8; j_1++) { + w_lo[j_1] = 1.0f; + w_hi[j_1] = 1.0f; + c_lo[j_1] = 1.0f; + c_hi[j_1] = 0.0f; + } + if (is_row_cta && active) { + long long src_base = (long long)row * (long long)width + (long long)(head_in_row * 128); + { + const uint4* _vptr_0 = + reinterpret_cast(qkv + (src_base + (long long)lane_lo) + 0); + uint4* _vdst_0 = reinterpret_cast(&x_lo_w[0]); +#pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vdst_0[_blk].x), "=r"(_vdst_0[_blk].y), "=r"(_vdst_0[_blk].z), + "=r"(_vdst_0[_blk].w) + : "l"((const void*)(_vptr_0 + _blk)) + : "memory"); + } + } + { + const uint4* _vptr_1 = + reinterpret_cast(qkv + (src_base + (long long)lane_hi) + 0); + uint4* _vdst_1 = reinterpret_cast(&x_hi_w[0]); +#pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vdst_1[_blk].x), "=r"(_vdst_1[_blk].y), "=r"(_vdst_1[_blk].z), + "=r"(_vdst_1[_blk].w) + : "l"((const void*)(_vptr_1 + _blk)) + : "memory"); + } + } + if (norm_policy != 0 && (is_q || is_k)) { + if (is_q) { + { + unsigned _v4_2_0; + unsigned _v4_2_1; + unsigned _v4_2_2; + unsigned _v4_2_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_2_0), "=r"(_v4_2_1), "=r"(_v4_2_2), "=r"(_v4_2_3) + : "l"((const void*)(q_norm_weight + (lane_1 * 8) + (0))) + : "memory"); + w_lo[0 + 0] = __uint_as_float(_v4_2_0); + w_lo[0 + 1] = __uint_as_float(_v4_2_1); + w_lo[0 + 2] = __uint_as_float(_v4_2_2); + w_lo[0 + 3] = __uint_as_float(_v4_2_3); + } + { + unsigned _v4_3_0; + unsigned _v4_3_1; + unsigned _v4_3_2; + unsigned _v4_3_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_3_0), "=r"(_v4_3_1), "=r"(_v4_3_2), "=r"(_v4_3_3) + : "l"((const void*)(q_norm_weight + (lane_1 * 8 + 4) + (0))) + : "memory"); + w_lo[4 + 0] = __uint_as_float(_v4_3_0); + w_lo[4 + 1] = __uint_as_float(_v4_3_1); + w_lo[4 + 2] = __uint_as_float(_v4_3_2); + w_lo[4 + 3] = __uint_as_float(_v4_3_3); + } + { + unsigned _v4_4_0; + unsigned _v4_4_1; + unsigned _v4_4_2; + unsigned _v4_4_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_4_0), "=r"(_v4_4_1), "=r"(_v4_4_2), "=r"(_v4_4_3) + : "l"((const void*)(q_norm_weight + (64 + lane_1 * 8) + (0))) + : "memory"); + w_hi[0 + 0] = __uint_as_float(_v4_4_0); + w_hi[0 + 1] = __uint_as_float(_v4_4_1); + w_hi[0 + 2] = __uint_as_float(_v4_4_2); + w_hi[0 + 3] = __uint_as_float(_v4_4_3); + } + { + unsigned _v4_5_0; + unsigned _v4_5_1; + unsigned _v4_5_2; + unsigned _v4_5_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_5_0), "=r"(_v4_5_1), "=r"(_v4_5_2), "=r"(_v4_5_3) + : "l"((const void*)(q_norm_weight + (64 + lane_1 * 8 + 4) + (0))) + : "memory"); + w_hi[4 + 0] = __uint_as_float(_v4_5_0); + w_hi[4 + 1] = __uint_as_float(_v4_5_1); + w_hi[4 + 2] = __uint_as_float(_v4_5_2); + w_hi[4 + 3] = __uint_as_float(_v4_5_3); + } + } else { + { + unsigned _v4_6_0; + unsigned _v4_6_1; + unsigned _v4_6_2; + unsigned _v4_6_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_6_0), "=r"(_v4_6_1), "=r"(_v4_6_2), "=r"(_v4_6_3) + : "l"((const void*)(k_norm_weight + (lane_1 * 8) + (0))) + : "memory"); + w_lo[0 + 0] = __uint_as_float(_v4_6_0); + w_lo[0 + 1] = __uint_as_float(_v4_6_1); + w_lo[0 + 2] = __uint_as_float(_v4_6_2); + w_lo[0 + 3] = __uint_as_float(_v4_6_3); + } + { + unsigned _v4_7_0; + unsigned _v4_7_1; + unsigned _v4_7_2; + unsigned _v4_7_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_7_0), "=r"(_v4_7_1), "=r"(_v4_7_2), "=r"(_v4_7_3) + : "l"((const void*)(k_norm_weight + (lane_1 * 8 + 4) + (0))) + : "memory"); + w_lo[4 + 0] = __uint_as_float(_v4_7_0); + w_lo[4 + 1] = __uint_as_float(_v4_7_1); + w_lo[4 + 2] = __uint_as_float(_v4_7_2); + w_lo[4 + 3] = __uint_as_float(_v4_7_3); + } + { + unsigned _v4_8_0; + unsigned _v4_8_1; + unsigned _v4_8_2; + unsigned _v4_8_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_8_0), "=r"(_v4_8_1), "=r"(_v4_8_2), "=r"(_v4_8_3) + : "l"((const void*)(k_norm_weight + (64 + lane_1 * 8) + (0))) + : "memory"); + w_hi[0 + 0] = __uint_as_float(_v4_8_0); + w_hi[0 + 1] = __uint_as_float(_v4_8_1); + w_hi[0 + 2] = __uint_as_float(_v4_8_2); + w_hi[0 + 3] = __uint_as_float(_v4_8_3); + } + { + unsigned _v4_9_0; + unsigned _v4_9_1; + unsigned _v4_9_2; + unsigned _v4_9_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_9_0), "=r"(_v4_9_1), "=r"(_v4_9_2), "=r"(_v4_9_3) + : "l"((const void*)(k_norm_weight + (64 + lane_1 * 8 + 4) + (0))) + : "memory"); + w_hi[4 + 0] = __uint_as_float(_v4_9_0); + w_hi[4 + 1] = __uint_as_float(_v4_9_1); + w_hi[4 + 2] = __uint_as_float(_v4_9_2); + w_hi[4 + 3] = __uint_as_float(_v4_9_3); + } + } + } + } + float k_scale_v = 1.0f; + float v_scale_v = 1.0f; + float q_static = 1.0f; + if (is_row_cta) { + if (is_k) { + k_scale_v = k_scale[0]; + } + if (is_v) { + v_scale_v = v_scale[0]; + } + if (is_q && quant_policy == 2) { + q_static = q_scale_inv[0]; + } + } + int req_id = bid_1 - num_rows; + int clear_last_pos = -1; + if (!is_row_cta && req_id < num_requests) { + clear_last_pos = seq_lens[req_id] - 1; + } + bool any_bad = 0; +#pragma unroll 1 + for (int req = tid_0; req < num_requests; req += 96) { + int begin = q_indptr[req]; + int end = q_indptr[req + 1]; + int seq_len_r = seq_lens[req]; + bool bad = begin < 0 || end < begin || end > num_rows; + if (dyn_prefill && end - begin > max_seqlen) { + bad = 1; + } + if (bad) { + any_bad = 1; + } + if (is_row_cta) { + if (begin <= row && row < end) { + meta[0] = req; + meta[1] = row + seq_len_r - end; + meta[2] = row - begin; + } + } + } + if (tid_0 == 0) { + int first = q_indptr[0]; + int last_ip = q_indptr[num_requests]; + if (first != 0 || last_ip != num_rows) { + any_bad = 1; + } + } + uint32_t _cta_or_0 = __syncthreads_or(any_bad); + unsigned int invalid = _cta_or_0; + if (is_row_cta) { + if (invalid == 0 && active) { + int batch_i = meta[0]; + int pos = meta[1]; + int tok = meta[2]; + long long cache_row = 0; + if (is_q || is_k) { + long long cs_base = (long long)pos * 128 + (long long)(lane_1 * 8); + { + unsigned _v4_10_0; + unsigned _v4_10_1; + unsigned _v4_10_2; + unsigned _v4_10_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_10_0), "=r"(_v4_10_1), "=r"(_v4_10_2), "=r"(_v4_10_3) + : "l"((const void*)(cos_sin + cs_base + (0))) + : "memory"); + c_lo[0 + 0] = __uint_as_float(_v4_10_0); + c_lo[0 + 1] = __uint_as_float(_v4_10_1); + c_lo[0 + 2] = __uint_as_float(_v4_10_2); + c_lo[0 + 3] = __uint_as_float(_v4_10_3); + } + { + unsigned _v4_11_0; + unsigned _v4_11_1; + unsigned _v4_11_2; + unsigned _v4_11_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_11_0), "=r"(_v4_11_1), "=r"(_v4_11_2), "=r"(_v4_11_3) + : "l"((const void*)(cos_sin + (cs_base + 4) + (0))) + : "memory"); + c_lo[4 + 0] = __uint_as_float(_v4_11_0); + c_lo[4 + 1] = __uint_as_float(_v4_11_1); + c_lo[4 + 2] = __uint_as_float(_v4_11_2); + c_lo[4 + 3] = __uint_as_float(_v4_11_3); + } + { + unsigned _v4_12_0; + unsigned _v4_12_1; + unsigned _v4_12_2; + unsigned _v4_12_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_12_0), "=r"(_v4_12_1), "=r"(_v4_12_2), "=r"(_v4_12_3) + : "l"((const void*)(cos_sin + (cs_base + 64) + (0))) + : "memory"); + c_hi[0 + 0] = __uint_as_float(_v4_12_0); + c_hi[0 + 1] = __uint_as_float(_v4_12_1); + c_hi[0 + 2] = __uint_as_float(_v4_12_2); + c_hi[0 + 3] = __uint_as_float(_v4_12_3); + } + { + unsigned _v4_13_0; + unsigned _v4_13_1; + unsigned _v4_13_2; + unsigned _v4_13_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_13_0), "=r"(_v4_13_1), "=r"(_v4_13_2), "=r"(_v4_13_3) + : "l"((const void*)(cos_sin + (cs_base + 64 + 4) + (0))) + : "memory"); + c_hi[4 + 0] = __uint_as_float(_v4_13_0); + c_hi[4 + 1] = __uint_as_float(_v4_13_1); + c_hi[4 + 2] = __uint_as_float(_v4_13_2); + c_hi[4 + 3] = __uint_as_float(_v4_13_3); + } + } + if ((is_k || is_v) && !has_out_kv) { + int page_slot = pos / page_size; + int row_in_page = pos - page_slot * page_size; + int phys = page_indices[batch_i * max_pages_per_request + page_slot]; + cache_row = ((long long)phys * (long long)page_size + (long long)row_in_page) * + (long long)(num_kv_heads * 128) + + (long long)(kv_head * 128); + } + float x_lo[8]; + float x_hi[8]; +#pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&x_lo[_pair * 2])[0]), "=f"((&x_lo[_pair * 2])[1]) + : "r"(x_lo_w[_pair])); + } +#pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&x_hi[_pair * 2])[0]), "=f"((&x_hi[_pair * 2])[1]) + : "r"(x_hi_w[_pair])); + } + if (is_q || is_k) { + if (norm_policy == 2) { + float ss = 0.0f; +#pragma unroll + for (int j_2 = 0; j_2 < 8; j_2++) { + float _fma_0 = __fmaf_rn(x_lo[j_2], x_lo[j_2], ss); + ss = _fma_0; + float _fma_1 = __fmaf_rn(x_hi[j_2], x_hi[j_2], ss); + ss = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(grp_mask, ss, 4); + ss = ss + _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(grp_mask, ss, 2); + ss = ss + _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(grp_mask, ss, 1); + ss = ss + _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + 1e-06f); + float inv_rms = _rsqrt_0; +#pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x_lo[j_3] = x_lo[j_3] * (inv_rms * w_lo[j_3]); + x_hi[j_3] = x_hi[j_3] * (inv_rms * w_hi[j_3]); + } + } +#pragma unroll + for (int j_4 = 0; j_4 < 8; j_4++) { + float y_lo = x_lo[j_4] * c_lo[j_4] - x_hi[j_4] * c_hi[j_4]; + float y_hi = x_hi[j_4] * c_lo[j_4] + x_lo[j_4] * c_hi[j_4]; + x_lo[j_4] = y_lo; + x_hi[j_4] = y_hi; + } + if (norm_policy == 1) { + float ss1 = 0.0f; +#pragma unroll + for (int j_5 = 0; j_5 < 8; j_5++) { + float _fma_2 = __fmaf_rn(x_lo[j_5], x_lo[j_5], ss1); + ss1 = _fma_2; + float _fma_3 = __fmaf_rn(x_hi[j_5], x_hi[j_5], ss1); + ss1 = _fma_3; + } + float _shfl_xor_3 = __shfl_xor_sync(grp_mask, ss1, 4); + ss1 = ss1 + _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(grp_mask, ss1, 2); + ss1 = ss1 + _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(grp_mask, ss1, 1); + ss1 = ss1 + _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss1 * 0.0078125f + 1e-06f); + float inv_rms1 = _rsqrt_1; +#pragma unroll + for (int j_6 = 0; j_6 < 8; j_6++) { + x_lo[j_6] = x_lo[j_6] * (inv_rms1 * w_lo[j_6]); + x_hi[j_6] = x_hi[j_6] * (inv_rms1 * w_hi[j_6]); + } + } + } + float _rcp_0 = __frcp_rn(v_scale_v); + float mult = _rcp_0; + if (is_q) { + if (quant_policy == 1) { + float m = 1e-06f; +#pragma unroll + for (int j_7 = 0; j_7 < 8; j_7++) { + float _fabs_0 = fabsf(x_lo[j_7]); + float _fmax_0 = fmaxf(m, _fabs_0); + m = _fmax_0; + float _fabs_1 = fabsf(x_hi[j_7]); + float _fmax_1 = fmaxf(m, _fabs_1); + m = _fmax_1; + } + float _shfl_xor_6 = __shfl_xor_sync(grp_mask, m, 4); + float _fmax_2 = fmaxf(m, _shfl_xor_6); + m = _fmax_2; + float _shfl_xor_7 = __shfl_xor_sync(grp_mask, m, 2); + float _fmax_3 = fmaxf(m, _shfl_xor_7); + m = _fmax_3; + float _shfl_xor_8 = __shfl_xor_sync(grp_mask, m, 1); + float _fmax_4 = fmaxf(m, _shfl_xor_8); + m = _fmax_4; + float _fdiv_rn_0 = __fdiv_rn(m, upper_max); + float scale_val = _fdiv_rn_0; + if (lane_1 == 0) { + if (is_prefill) { + if (tok < max_seqlen_aligned) { + long long qs_idx = + ((long long)batch_i * (long long)num_q_heads + (long long)head_in_row) * + (long long)max_seqlen_aligned + + (long long)tok; + *(reinterpret_cast(q_scale + qs_idx) + (0)) = scale_val; + } + } else { + long long qs_idx2 = (long long)row * (long long)num_q_heads + (long long)head_in_row; + *(reinterpret_cast(q_scale + qs_idx2) + (0)) = scale_val; + } + } + float _rcp_1 = __frcp_rn(scale_val); + mult = _rcp_1; + } else { + mult = q_static; + } + } else if (is_k) { + float _rcp_2 = __frcp_rn(k_scale_v); + mult = _rcp_2; + } +#pragma unroll + for (int j_8 = 0; j_8 < 8; j_8++) { + x_lo[j_8] = x_lo[j_8] * mult; + x_hi[j_8] = x_hi[j_8] * mult; + } + long long row64 = row; + if (is_q) { + long long q_dst = (row64 * (long long)num_q_heads + (long long)head_in_row) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_q + (q_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_q + (q_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else if (is_k) { + if (has_out_kv) { + long long k_dst = (row64 * (long long)num_kv_heads + (long long)kv_head) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_k + (k_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_k + (k_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else { + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(key_cache + (cache_row + (long long)lane_lo)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(key_cache + (cache_row + (long long)lane_hi)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + } + } else { + if (has_out_kv) { + long long v_dst = (row64 * (long long)num_kv_heads + (long long)kv_head) * 128; + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_v + (v_dst + (long long)lane_lo)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(out_v + (v_dst + (long long)lane_hi)) + (0)) = + *reinterpret_cast(_fp8_pk); + } + } else { + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_lo[0 + 0]), "f"(x_lo[0 + 1]), "f"(x_lo[0 + 2]), "f"(x_lo[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_lo[0 + 4]), "f"(x_lo[0 + 5]), "f"(x_lo[0 + 6]), "f"(x_lo[0 + 7])); + *reinterpret_cast( + reinterpret_cast(value_cache + (cache_row + (long long)lane_lo)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + { + unsigned int _fp8_pk[2]; + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[0]) + : "f"(x_hi[0 + 0]), "f"(x_hi[0 + 1]), "f"(x_hi[0 + 2]), "f"(x_hi[0 + 3])); + asm("{\n\t" + ".reg .b16 _lo, _hi;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _lo, %2, %1;\n\t" + "cvt.rn.satfinite.e4m3x2.f32 _hi, %4, %3;\n\t" + "mov.b32 %0, {_lo, _hi};\n\t" + "}\n" + : "=r"(_fp8_pk[1]) + : "f"(x_hi[0 + 4]), "f"(x_hi[0 + 5]), "f"(x_hi[0 + 6]), "f"(x_hi[0 + 7])); + *reinterpret_cast( + reinterpret_cast(value_cache + (cache_row + (long long)lane_hi)) + + (0)) = *reinterpret_cast(_fp8_pk); + } + } + } + } + } else if (req_id < num_requests) { + if (tid_0 == 0) { + int flag_val = ((invalid != 0) ? -1 : 0); + *(reinterpret_cast(split_k_flag + (req_id * num_kv_heads + kv_head)) + (0)) = flag_val; + } + if (invalid == 0) { + int last_pos = clear_last_pos; + if (last_pos >= 0) { + int last_slot = last_pos / page_size; + int zero_from = last_pos - last_slot * page_size + 1; + if (zero_from < page_size) { + int phys_last = page_indices[req_id * max_pages_per_request + last_slot]; + long long phys_last64 = phys_last; + unsigned int zero_w[4]; +#pragma unroll + for (int j_9 = 0; j_9 < 4; j_9++) { + zero_w[j_9] = 0; + } + int n_rows = page_size - zero_from; + int total_items = n_rows * 16; +#pragma unroll 1 + for (int item = tid_0; item < total_items; item += 96) { + int r = item / 16; + int sub = item - r * 16; + int chunk = sub % 8; + long long base = + (phys_last64 * (long long)page_size + (long long)zero_from + (long long)r) * + (long long)(num_kv_heads * 128) + + (long long)(kv_head * 128) + (long long)(chunk * 16); + if (sub < 8) { + reinterpret_cast(key_cache + base)[0] = reinterpret_cast(zero_w)[0]; + } else { + reinterpret_cast(value_cache + base)[0] = reinterpret_cast(zero_w)[0]; + } + } + } + } + } + } +} + +} // extern "C" diff --git a/flashinfer/cake_fused_qk_rope_fp8_append.py b/flashinfer/cake_fused_qk_rope_fp8_append.py new file mode 100644 index 00000000000..753f0702118 --- /dev/null +++ b/flashinfer/cake_fused_qk_rope_fp8_append.py @@ -0,0 +1,378 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Cake fused QK RMSNorm + NeoX RoPE + FP8 E4M3 quantization + paged KV append (one launch). + +:func:`cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache` takes the +packed BF16 ``qkv`` projection of a decode or prefill step and, in a single +kernel launch, + +1. unpacks Q, K and V (``[T, (Hq + 2 * Hkv) * 128]``, Q | K | V order), +2. applies an optional per-head RMSNorm to Q and K (f32 math, f32 weights, + ``eps = 1e-6``; ``qk_norm_policy`` 0 = none, 1 = RoPE then norm, + 2 = norm then RoPE), +3. applies NeoX rotary embedding (dims ``(i, i + 64)``) from a caller-provided + ``cos_sin_cache`` table indexed by absolute position, +4. quantizes Q to FP8 E4M3 (``quant_policy=1``: dynamic per-token / per-head + scale ``max(|q|, 1e-6) / upper_max`` returned in ``q_scale``; + ``quant_policy=2``: static ``q_scale_inv`` multiplier, empty ``q_scale``), +5. quantizes K and V with the scalar dequantization scales ``k_scale`` / + ``v_scale`` (payload ``x / scale``) and appends them to an NHD FP8 paged + cache ``[num_pages, page_size, Hkv, 128]`` (or into caller-owned + ``out_k`` / ``out_v`` ``[T, Hkv, 128]``), +6. zeroes the unused rows ``[(seq_lens[b] - 1) % page_size + 1, page_size)`` of + every request's last page in both caches, and +7. validates ``q_indptr`` on the device: a malformed table (``q_indptr[0] != 0``, + non-monotone, ``q_indptr[-1] != T``) or, for dynamic prefill quantization, a + request longer than ``max_seqlen`` leaves every output untouched and fills + ``split_k_flag`` with ``-1``; a valid call sets every flag to 0 (empty + requests included). No host synchronisation is involved. + +Positions: row ``r`` belongs to request ``b`` when ``q_indptr[b] <= r < +q_indptr[b + 1]`` and has position ``r + seq_lens[b] - q_indptr[b + 1]`` +(``seq_lens`` is the total length *after* the append). K/V of row ``r`` land at +``(page_indices[b, pos // page_size], pos % page_size, kv_head)``. + +Supported: ``(Hq, Hkv) in {(8, 1), (64, 8)}``, head_dim 128, BF16 input, FP8 +E4M3 output, SM90 (H100/H200), SM100 (B200) and SM103 (B300). The kernel is +generated from the Cake schedule +``flashinfer_blackwell_fused_qk_rmsnorm_rope_fp8_paged_kv_append``; sources and +manifest live under ``csrc/cake_fused_qk_rope_fp8_append`` (JIT-only). The call +is CUDA-graph capturable: every launch argument is a by-value scalar or a tensor +pointer, no host synchronisation, workspace or extra kernel is involved. +""" + +from __future__ import annotations + +import math +from typing import Any, Optional, Tuple + +import torch + +from .jit.cake_fused_qk_rope_fp8_append import ( + STAGES, + arch_for, + load_cake_fused_qk_rope_fp8_append_module, +) +from .utils import get_compute_capability + +HEAD_DIM = 128 +FP8_MAX = 448.0 +SUPPORTED_HEADS = ((8, 1), (64, 8)) +_STAGE = STAGES[0] + + +def launch_plan(num_rows: int, num_requests: int, num_kv_heads: int) -> dict[str, Any]: + """Grid of the fused kernel (mirrors the Cake kernel's ``launch_plan``). + + One CTA per (packed row, kv head group) plus one CTA per (request, kv head) + for the last-page tail clear and ``split_k_flag``. + """ + return {"grid": (num_rows + num_requests, num_kv_heads, 1)} + + +def _check(cond: bool, message: str) -> None: + if not cond: + raise ValueError(message) + + +def _heads_from_packed_qkv( + qkv: torch.Tensor, key_cache: torch.Tensor, value_cache: torch.Tensor +) -> Tuple[int, int]: + _check( + key_cache.dim() == 4 and value_cache.dim() == 4, + "key_cache and value_cache must be NHD [num_pages, page_size, num_kv_heads, 128]", + ) + num_kv_heads = int(key_cache.shape[2]) + width = int(qkv.shape[1]) + _check( + width % HEAD_DIM == 0 and width // HEAD_DIM > 2 * num_kv_heads, + "qkv must be [T, (num_q_heads + 2 * num_kv_heads) * 128]", + ) + num_q_heads = width // HEAD_DIM - 2 * num_kv_heads + _check( + (num_q_heads, num_kv_heads) in SUPPORTED_HEADS, + f"(num_q_heads, num_kv_heads) = ({num_q_heads}, {num_kv_heads}) is not supported; " + f"expected one of {SUPPORTED_HEADS}", + ) + return num_q_heads, num_kv_heads + + +def cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache( + qkv: torch.Tensor, + cos_sin_cache: torch.Tensor, + seq_lens: torch.Tensor, + q_indptr: torch.Tensor, + page_indices: torch.Tensor, + paged_kv_cache: Tuple[torch.Tensor, torch.Tensor], + is_prefill: bool, + k_scale: torch.Tensor, + v_scale: torch.Tensor, + quant_policy: int, + max_seqlen: int = 0, + upper_max: float = FP8_MAX, + q_scale_inv: Optional[torch.Tensor] = None, + q_norm_weight: Optional[torch.Tensor] = None, + k_norm_weight: Optional[torch.Tensor] = None, + qk_norm_policy: int = 0, + out_q: Optional[torch.Tensor] = None, + out_k: Optional[torch.Tensor] = None, + out_v: Optional[torch.Tensor] = None, + q_scale: Optional[torch.Tensor] = None, + split_k_flag: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + r"""Fused Q/K RMSNorm + NeoX RoPE + FP8 quantization + paged K/V append in one launch. + + Parameters + ---------- + qkv : torch.Tensor + BF16 ``[T, (num_q_heads + 2 * num_kv_heads) * 128]``, packed Q | K | V. The head + configuration is derived from this width and the cache's kv-head count. + cos_sin_cache : torch.Tensor + F32 ``[max_position, 128]``; row ``p`` is ``cos(p, 0..63) | sin(p, 0..63)``. + seq_lens : torch.Tensor + I32 ``[B]`` total sequence length of every request *after* this append. + q_indptr : torch.Tensor + I32 ``[B + 1]``; rows ``[q_indptr[b], q_indptr[b + 1])`` belong to request ``b``. + page_indices : torch.Tensor + I32 ``[B, max_pages]`` dense physical page table (entries beyond a request's + pages are ignored). + paged_kv_cache : (torch.Tensor, torch.Tensor) + FP8 E4M3 NHD ``(key_cache, value_cache)``, each + ``[num_pages, page_size, num_kv_heads, 128]`` (written in place). + is_prefill : bool + Selects the ``q_scale`` layout of ``quant_policy=1`` (see below). + k_scale, v_scale : torch.Tensor + F32 one-element dequantization scales; the stored payload is ``x / scale``. + quant_policy : int + 1 = dynamic per-token / per-Q-head scale ``max(|q|, 1e-6) / upper_max``; + 2 = static scale: payload ``q * q_scale_inv`` and an empty ``q_scale``. + max_seqlen : int + Dynamic prefill (``quant_policy=1`` and ``is_prefill``) only: must cover the + longest request; ``q_scale`` is ``[B, num_q_heads, round_up(max_seqlen, 128)]`` + addressed by the token index inside its request. + upper_max : float + Finite E4M3 magnitude in ``(0, 448]`` used by the dynamic scale. + q_scale_inv : torch.Tensor, optional + F32 one-element static Q multiplier (required for ``quant_policy=2``). + q_norm_weight, k_norm_weight : torch.Tensor, optional + F32 ``[128]`` per-head RMSNorm weights (required unless ``qk_norm_policy == 0``). + qk_norm_policy : int + 0 no norm, 1 RoPE then RMSNorm, 2 RMSNorm then RoPE. + out_q : torch.Tensor, optional + FP8 ``[T, num_q_heads, 128]`` output buffer (allocated when ``None``). + out_k, out_v : torch.Tensor, optional + FP8 ``[T, num_kv_heads, 128]``; when both are given, K / V rows are written there + instead of the caches (the last-page tail clear still runs). + q_scale : torch.Tensor, optional + F32 output buffer: decode ``[T, num_q_heads]``, dynamic prefill + ``[B, num_q_heads, round_up(max_seqlen, 128)]``, ``quant_policy=2`` empty. + split_k_flag : torch.Tensor, optional + I32 ``[B, num_kv_heads]`` output buffer; every element is written (0 valid, + -1 invalid metadata), no host-side clear is needed. + + Returns + ------- + (out_q, q_scale, split_k_flag) + The caller-owned buffers when given, otherwise freshly allocated ones. + """ + _check(len(paged_kv_cache) == 2, "paged_kv_cache must be (key_cache, value_cache)") + key_cache, value_cache = paged_kv_cache + _check(quant_policy in (1, 2), "quant_policy must be 1 (dynamic Q) or 2 (static Q)") + _check(qk_norm_policy in (0, 1, 2), "qk_norm_policy must be 0, 1 or 2") + _check( + math.isfinite(upper_max) and 0.0 < upper_max <= FP8_MAX, + "upper_max must be finite and in the interval (0, 448]", + ) + _check( + qkv.dim() == 2 and qkv.dtype == torch.bfloat16, "qkv must be a 2-D bf16 tensor" + ) + _check(qkv.is_cuda, "qkv must be a CUDA tensor") + device = qkv.device + num_q_heads, num_kv_heads = _heads_from_packed_qkv(qkv, key_cache, value_cache) + num_rows = int(qkv.shape[0]) + dyn_prefill = quant_policy == 1 and bool(is_prefill) + _check( + not dyn_prefill or max_seqlen > 0, + "max_seqlen must be positive for dynamic prefill quantization", + ) + _check( + cos_sin_cache.dtype == torch.float32 + and cos_sin_cache.dim() == 2 + and int(cos_sin_cache.shape[1]) == HEAD_DIM, + "cos_sin_cache must be f32 [max_position, 128]", + ) + for name, t in ( + ("seq_lens", seq_lens), + ("q_indptr", q_indptr), + ("page_indices", page_indices), + ): + _check(t.dtype == torch.int32, f"{name} must be int32") + num_requests = int(seq_lens.shape[0]) + _check( + seq_lens.dim() == 1 + and q_indptr.dim() == 1 + and int(q_indptr.shape[0]) == num_requests + 1, + "q_indptr must be [B + 1] and seq_lens [B]", + ) + _check( + page_indices.dim() == 2 and int(page_indices.shape[0]) == num_requests, + "page_indices must be [B, max_pages]", + ) + for name, t in (("key_cache", key_cache), ("value_cache", value_cache)): + _check( + t.dtype == torch.float8_e4m3fn + and int(t.shape[2]) == num_kv_heads + and int(t.shape[3]) == HEAD_DIM, + f"{name} must be float8_e4m3fn [num_pages, page_size, num_kv_heads, 128]", + ) + _check( + tuple(key_cache.shape) == tuple(value_cache.shape), + "key_cache and value_cache must have equal shapes", + ) + page_size = int(key_cache.shape[1]) + for name, t in (("k_scale", k_scale), ("v_scale", v_scale)): + _check( + t.dtype == torch.float32 and t.numel() == 1, + f"{name} must be a one-element f32 tensor", + ) + if quant_policy == 2: + _check(q_scale_inv is not None, "q_scale_inv is required when quant_policy=2") + _check( + q_scale_inv.dtype == torch.float32 and q_scale_inv.numel() == 1, + "q_scale_inv must be a one-element f32 tensor", + ) + else: + # Never read by the kernel under quant_policy=1: hand it an existing f32 scalar. + q_scale_inv = k_scale + if qk_norm_policy != 0: + _check( + q_norm_weight is not None and k_norm_weight is not None, + "q_norm_weight and k_norm_weight are required when qk_norm_policy != 0", + ) + if q_norm_weight is None or k_norm_weight is None: + # qk_norm_policy == 0 never reads the weights: hand the kernel an existing f32 [128] + # row instead of allocating and filling a tensor on every call (graph-safe). + placeholder = cos_sin_cache[0, :HEAD_DIM] + q_norm_weight = placeholder if q_norm_weight is None else q_norm_weight + k_norm_weight = placeholder if k_norm_weight is None else k_norm_weight + for name, w in (("q_norm_weight", q_norm_weight), ("k_norm_weight", k_norm_weight)): + _check( + w.dtype == torch.float32 and tuple(w.shape) == (HEAD_DIM,), + f"{name} must be f32 [128]", + ) + if out_q is None: + out_q = torch.empty( + (num_rows, num_q_heads, HEAD_DIM), dtype=torch.float8_e4m3fn, device=device + ) + _check( + out_q.dtype == torch.float8_e4m3fn + and tuple(out_q.shape) == (num_rows, num_q_heads, HEAD_DIM), + "out_q must be float8_e4m3fn [T, num_q_heads, 128]", + ) + _check((out_k is None) == (out_v is None), "out_k and out_v must be given together") + has_out_kv = out_k is not None + for name, t in (("out_k", out_k), ("out_v", out_v)): + if t is not None: + _check( + t.dtype == torch.float8_e4m3fn + and tuple(t.shape) == (num_rows, num_kv_heads, HEAD_DIM), + f"{name} must be float8_e4m3fn [T, num_kv_heads, 128]", + ) + aligned = (max_seqlen + 127) // 128 * 128 if dyn_prefill else 0 + if dyn_prefill: + q_scale_shape: Tuple[int, ...] = (num_requests, num_q_heads, aligned) + elif quant_policy == 1: + q_scale_shape = (num_rows, num_q_heads) + else: + q_scale_shape = (0,) + if q_scale is None: + q_scale = torch.empty(q_scale_shape, dtype=torch.float32, device=device) + _check( + q_scale.dtype == torch.float32 and tuple(q_scale.shape) == q_scale_shape, + f"q_scale must be f32 {q_scale_shape} for this quant_policy / is_prefill", + ) + if split_k_flag is None: + split_k_flag = torch.empty( + (num_requests, num_kv_heads), dtype=torch.int32, device=device + ) + _check( + split_k_flag.dtype == torch.int32 + and tuple(split_k_flag.shape) == (num_requests, num_kv_heads), + "split_k_flag must be int32 [B, num_kv_heads]", + ) + tensors = { + "qkv": qkv, + "cos_sin": cos_sin_cache, + "seq_lens": seq_lens, + "q_indptr": q_indptr, + "page_indices": page_indices, + "q_norm_weight": q_norm_weight, + "k_norm_weight": k_norm_weight, + "k_scale": k_scale, + "v_scale": v_scale, + "q_scale_inv": q_scale_inv, + "out_q": out_q, + "key_cache": key_cache, + "value_cache": value_cache, + "out_k": out_k if has_out_kv else key_cache, + "out_v": out_v if has_out_kv else value_cache, + "q_scale": q_scale, + "split_k_flag": split_k_flag, + } + for name, t in tensors.items(): + _check(t.is_contiguous(), f"{name} must be contiguous") + _check(t.device == device, f"{name} must be on {device}") + + arch = arch_for(get_compute_capability(device)) + module, record = load_cake_fused_qk_rope_fp8_append_module(_STAGE, arch) + plan = launch_plan(num_rows, num_requests, num_kv_heads) + scalars = { + "num_rows": num_rows, + "num_requests": num_requests, + "num_q_heads": num_q_heads, + "num_kv_heads": num_kv_heads, + "page_size": page_size, + "max_pages_per_request": int(page_indices.shape[1]), + "max_seqlen": int(max_seqlen), + "max_seqlen_aligned": aligned, + "quant_policy": int(quant_policy), + "norm_policy": int(qk_norm_policy), + "is_prefill": bool(is_prefill), + "has_out_kv": has_out_kv, + "upper_max": float(upper_max), + "grid_x": plan["grid"][0], + "grid_y": plan["grid"][1], + "grid_z": plan["grid"][2], + } + args = [] + for kind, name in record["arg_plan"]: + if name in tensors: + args.append(tensors[name]) + elif name in scalars: + args.append(scalars[name]) + else: + raise RuntimeError( + f"generated module expects unknown argument {name!r} ({kind})" + ) + import tvm_ffi + + with torch.cuda.device(device), tvm_ffi.use_torch_stream(): + getattr(module, str(record["ffi_entry"]))(*args) + return out_q, q_scale, split_k_flag + + +__all__ = [ + "cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache", + "launch_plan", +] diff --git a/flashinfer/jit/cake_fused_qk_rope_fp8_append.py b/flashinfer/jit/cake_fused_qk_rope_fp8_append.py new file mode 100644 index 00000000000..ea5e7ce4360 --- /dev/null +++ b/flashinfer/jit/cake_fused_qk_rope_fp8_append.py @@ -0,0 +1,183 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""JIT loader for the generated Cake fused QK RMSNorm + NeoX RoPE + FP8 quantize + paged KV append kernel. + +One source-only manifest lives under ``csrc/cake_fused_qk_rope_fp8_append`` +(``cake_fused_qk_rope_fp8_append_manifest.json``, schema ``cake.library_export.v5``): +a single stage (``fp8_e4m3``; head configuration, quantization / norm policies, +prefill vs decode and the caller-owned K/V route are launch parameters) rendered +for ``sm_90a`` (H100/H200), ``sm_100a`` (B200) and ``sm_103a`` (B300), each with +its tvm-ffi binding. The host wrapper lives in +:mod:`flashinfer.cake_fused_qk_rope_fp8_append`. +""" + +from __future__ import annotations + +import functools +import hashlib +import json +from pathlib import Path +from typing import Any + +from . import env as jit_env +from .core import ( + JitSpec, + gen_jit_spec, + logger, + sm90a_nvcc_flags, + sm100a_nvcc_flags, + sm103a_nvcc_flags, +) + +_PACKAGE = "cake_fused_qk_rope_fp8_append" +_GENERATED_ROOT = f"csrc/{_PACKAGE}" +_MANIFEST_NAME = f"{_PACKAGE}_manifest.json" + +STAGES = ("fp8_e4m3",) +ARCHES = ("sm_90a", "sm_100a", "sm_103a") +_NVCC_FLAGS = { + "sm_90a": sm90a_nvcc_flags, + "sm_100a": sm100a_nvcc_flags, + "sm_103a": sm103a_nvcc_flags, +} + + +def arch_for(compute_capability: tuple[int, int]) -> str: + major, minor = compute_capability + if (major, minor) == (9, 0): + return "sm_90a" + if (major, minor) == (10, 0): + return "sm_100a" + if (major, minor) == (10, 3): + return "sm_103a" + raise ValueError( + f"cake_fused_qk_rope_fp8_append has no cubin for compute capability {major}.{minor}; " + f"supported: sm_90a (9.0), sm_100a (10.0), sm_103a (10.3)" + ) + + +def _get_csrc_dir() -> Path: + installed = jit_env.FLASHINFER_CSRC_DIR / _PACKAGE + if (installed / _MANIFEST_NAME).is_file(): + return installed + checkout = Path(__file__).resolve().parents[2] / "csrc" / _PACKAGE + if (checkout / _MANIFEST_NAME).is_file(): + return checkout + raise FileNotFoundError("Cake fused QK RoPE append sources were not found") + + +def _get_include_dir() -> Path: + if jit_env.FLASHINFER_INCLUDE_DIR.exists(): + return jit_env.FLASHINFER_INCLUDE_DIR + checkout = Path(__file__).resolve().parents[2] / "include" + if checkout.exists(): + return checkout + raise FileNotFoundError("FlashInfer headers were not found") + + +@functools.cache +def _manifest() -> dict[str, Any]: + path = _get_csrc_dir() / _MANIFEST_NAME + value = json.loads(path.read_text(encoding="utf-8")) + if ( + value.get("schema") != "cake.library_export.v5" + or value.get("producer") != "cake" + or value.get("library") != "flashinfer" + or value.get("name") != _PACKAGE + ): + raise RuntimeError("invalid Cake fused QK RoPE append manifest") + modules = value.get("modules") + if not isinstance(modules, list) or not modules: + raise RuntimeError("empty Cake fused QK RoPE append module inventory") + return value + + +def _record(stage: str, arch: str) -> dict[str, Any]: + if stage not in STAGES: + raise ValueError(f"unknown stage {stage!r}; expected one of {STAGES}") + if arch not in ARCHES: + raise ValueError(f"unknown arch {arch!r}; expected one of {ARCHES}") + matches = [ + item + for item in _manifest()["modules"] + if item.get("arch") == arch + and dict(item.get("route", {})).get("stage") == stage + ] + if len(matches) != 1: + raise RuntimeError( + f"expected one generated Cake module for ({arch}, {stage}), got {len(matches)}" + ) + return matches[0] + + +def _source_path(relative_path: str) -> Path: + manifest = _manifest() + relative = Path(relative_path) + try: + suffix = relative.relative_to(Path(_GENERATED_ROOT)) + except ValueError as exc: + raise RuntimeError( + f"generated source escaped {_GENERATED_ROOT}: {relative_path}" + ) from exc + path = _get_csrc_dir() / suffix + inventory = {item["path"]: item for item in manifest["files"]} + receipt = inventory.get(relative.as_posix()) + if not isinstance(receipt, dict) or not path.is_file(): + raise FileNotFoundError( + f"generated source is absent from the manifest: {relative_path}" + ) + digest = hashlib.sha256(path.read_bytes()).hexdigest() + if digest != receipt.get("sha256"): + raise RuntimeError(f"generated source hash mismatch: {relative_path}") + return path + + +def cake_fused_qk_rope_fp8_append_record(stage: str, arch: str) -> dict[str, Any]: + """Manifest record of one generated module (arg plan, launch block, specialization).""" + return dict(_record(stage, arch)) + + +@functools.cache +def gen_cake_fused_qk_rope_fp8_append_module(stage: str, arch: str) -> JitSpec: + record = _record(stage, arch) + units = record["translation_units"] + spec = gen_jit_spec( + name=f"{record['name']}_{arch}", + sources=[_source_path(units["device"]), _source_path(units["binding"])], + extra_cuda_cflags=[*_NVCC_FLAGS[arch], *record["compile_flags"]], + extra_ldflags=["-lcuda"], + extra_include_paths=[_get_csrc_dir().parent, _get_include_dir()], + ) + logger.info("Generated Cake fused QK RoPE append JIT spec: %s", spec.name) + return spec + + +@functools.cache +def load_cake_fused_qk_rope_fp8_append_module(stage: str, arch: str): + """Build and load one module; return ``(module, manifest record)``.""" + spec = gen_cake_fused_qk_rope_fp8_append_module(stage, arch) + module = spec.build_and_load() + return module, _record(stage, arch) + + +def build_all_cake_fused_qk_rope_fp8_append_modules(arch: str) -> dict[str, str]: + built = {} + for stage in STAGES: + spec = gen_cake_fused_qk_rope_fp8_append_module(stage, arch) + module = spec.build_and_load() + record = _record(stage, arch) + if not hasattr(module, str(record["ffi_entry"])): + raise RuntimeError(f"built module {record['name']} lacks its FFI entry") + built[stage] = spec.name + return built diff --git a/flashinfer/rope.py b/flashinfer/rope.py index b43047f6108..75c31d092fb 100644 --- a/flashinfer/rope.py +++ b/flashinfer/rope.py @@ -23,6 +23,9 @@ from .cake_fused_qk_rope_append import ( # noqa: F401 (re-exported Cake fused route) cake_fused_qk_rmsnorm_rope_append_paged_kv_cache, ) +from .cake_fused_qk_rope_fp8_append import ( # noqa: F401 (re-exported Cake fused FP8 route) + cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache, +) from .trace.templates.rope import ( apply_llama31_rope_inplace_trace, apply_llama31_rope_pos_ids_inplace_trace, diff --git a/tests/attention/test_cake_fused_qk_rope_fp8_append.py b/tests/attention/test_cake_fused_qk_rope_fp8_append.py new file mode 100644 index 00000000000..4be7fbb3ad7 --- /dev/null +++ b/tests/attention/test_cake_fused_qk_rope_fp8_append.py @@ -0,0 +1,593 @@ +"""Correctness tests for the Cake fused QK RMSNorm + NeoX RoPE + FP8 quantize + paged KV append kernel. + +Every case is checked against an independent fp32 reference (per-head RMSNorm, NeoX RoPE, +``amax / upper_max`` or static Q scale, ``x / scale`` K/V payloads, last-page tail clear, +device-side ``q_indptr`` validation). FP8 payloads are compared against the *unrounded* +fp32 payload with ``atol = rtol = 0.1``; ``q_scale`` with ``rtol = 1e-5``; flags, cleared +tails and every untouched byte bitwise. +""" + +import pytest +import torch + +from flashinfer.cake_fused_qk_rope_fp8_append import ( + cake_fused_qk_rmsnorm_rope_quantize_fp8_append_paged_kv_cache as fused_fp8, +) +from flashinfer.utils import get_compute_capability + +HEAD_DIM = 128 +FP8_MAX = 448.0 +AMAX_FLOOR = 1e-6 +EPS = 1e-6 +POISON_BYTE = 0xFF # NaN in float8_e4m3fn: never produced by the kernel + + +def _skip_unless_supported(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + cc = get_compute_capability(torch.device("cuda:0")) + if cc not in ((9, 0), (10, 0), (10, 3)): + pytest.skip( + f"cake_fused_qk_rope_fp8_append has no cubin for compute capability {cc}" + ) + + +def rotary_cos_sin_table( + max_positions: int, device, base: float = 10000.0 +) -> torch.Tensor: + half = HEAD_DIM // 2 + inv_freq = 1.0 / ( + base ** (torch.arange(0, half, dtype=torch.float64, device=device) / half) + ) + pos = torch.arange(max_positions, dtype=torch.float64, device=device) + freqs = torch.outer(pos, inv_freq) + return torch.cat([freqs.cos(), freqs.sin()], dim=1).to(torch.float32) + + +def _rmsnorm(x, w): + inv = torch.rsqrt(x.square().sum(dim=-1, keepdim=True) / x.shape[-1] + EPS) + return x * inv * w.float() + + +def _rope(x, cos_sin, positions): + half = HEAD_DIM // 2 + cos = cos_sin[positions, :half].float().unsqueeze(1) + sin = cos_sin[positions, half:].float().unsqueeze(1) + left, right = x[..., :half], x[..., half:] + return torch.cat([left * cos - right * sin, right * cos + left * sin], dim=-1) + + +def reference(case, *, caller_owned_kv=False): + """fp32 reference; returns the expected unrounded payloads in output layout. + + ``key_cache`` / ``value_cache`` expectations are fp32 tensors of the cache shape + holding NaN where the kernel must not write, 0 on the cleared tails and the + unrounded ``x / scale`` payload on the appended rows. + """ + qkv, dev = case["qkv"], case["qkv"].device + T = qkv.shape[0] + hq, hkv = case["num_q_heads"], case["num_kv_heads"] + B = case["seq_lens"].shape[0] + page_size = case["page_size"] + qi = case["q_indptr"].tolist() + sl = case["seq_lens"].tolist() + pi = case["page_indices"].cpu() + quant, norm = case["quant_policy"], case["qk_norm_policy"] + dyn_prefill = quant == 1 and case["is_prefill"] + aligned = (case["max_seqlen"] + 127) // 128 * 128 if dyn_prefill else 0 + + valid = qi[0] == 0 and qi[-1] == T + lengths = [qi[b + 1] - qi[b] for b in range(B)] + valid = valid and all(n >= 0 for n in lengths) + if dyn_prefill: + valid = valid and all(n <= case["max_seqlen"] for n in lengths) + flags = torch.full((B, hkv), 0 if valid else -1, dtype=torch.int32, device=dev) + nan_cache = torch.full( + case["cache_shape"], float("nan"), dtype=torch.float32, device=dev + ) + exp = dict( + valid=valid, + flags=flags, + out_q=None, + q_scale=None, + key_cache=nan_cache.clone(), + value_cache=nan_cache.clone(), + out_k=None, + out_v=None, + ) + if not valid: + return exp + + batch_ids = torch.empty(T, dtype=torch.int64) + positions = torch.empty(T, dtype=torch.int64) + for b in range(B): + lo, hi = qi[b], qi[b + 1] + if hi > lo: + batch_ids[lo:hi] = b + positions[lo:hi] = torch.arange(lo, hi) + sl[b] - hi + assert T == 0 or bool((positions >= 0).all()), ( + "test cases keep every position valid" + ) + + x = qkv.float() + q = x[:, : hq * HEAD_DIM].reshape(T, hq, HEAD_DIM) + k = x[:, hq * HEAD_DIM : (hq + hkv) * HEAD_DIM].reshape(T, hkv, HEAD_DIM) + v = x[:, (hq + hkv) * HEAD_DIM :].reshape(T, hkv, HEAD_DIM) + qw, kw = case["q_norm_weight"], case["k_norm_weight"] + if norm == 2: + q, k = _rmsnorm(q, qw), _rmsnorm(k, kw) + if T > 0: + pos_dev = positions.to(dev) + q, k = _rope(q, case["cos_sin"], pos_dev), _rope(k, case["cos_sin"], pos_dev) + if norm == 1: + q, k = _rmsnorm(q, qw), _rmsnorm(k, kw) + if quant == 1: + scale = torch.clamp(q.abs().amax(dim=-1), min=AMAX_FLOOR) / case["upper_max"] + q_payload = q / scale.unsqueeze(-1) + if dyn_prefill: + q_scale = torch.full((B, hq, aligned), float("nan"), device=dev) + for b in range(B): + lo, hi = qi[b], qi[b + 1] + if hi > lo: + q_scale[b, :, : hi - lo] = scale[lo:hi].t() + else: + q_scale = scale + else: + q_payload = q * case["q_scale_inv"].float().reshape(()) + q_scale = torch.empty(0, device=dev) + k_payload = k / case["k_scale"].float().reshape(()) + v_payload = v / case["v_scale"].float().reshape(()) + exp["out_q"], exp["q_scale"] = q_payload, q_scale + if caller_owned_kv: + exp["out_k"], exp["out_v"] = k_payload, v_payload + elif T > 0: + phys = pi[batch_ids, positions // page_size].to(dev) + slot = (positions % page_size).to(dev) + exp["key_cache"][phys, slot] = k_payload + exp["value_cache"][phys, slot] = v_payload + for b in range(B): + last = sl[b] - 1 + if last < 0: + continue + page = int(pi[b, last // page_size]) + zero_from = last % page_size + 1 + if zero_from < page_size: + exp["key_cache"][page, zero_from:] = 0.0 + exp["value_cache"][page, zero_from:] = 0.0 + return exp + + +def make_case( + *, + q_lens, + ctx_lens, + num_q_heads, + num_kv_heads, + page_size, + qk_norm_policy, + quant_policy, + seed, + is_prefill=None, + max_seqlen=None, + upper_max=FP8_MAX, + device="cuda", +): + g = torch.Generator(device="cpu").manual_seed(seed) + B = len(q_lens) + seq_lens = [c + q for c, q in zip(ctx_lens, q_lens, strict=True)] + pages_per_req = [(s + page_size - 1) // page_size for s in seq_lens] + max_pages = max(max(pages_per_req), 1) + total_pages = sum(pages_per_req) + 3 + perm = torch.randperm(total_pages, generator=g) + page_indices = torch.zeros((B, max_pages), dtype=torch.int32) + cursor = 0 + for b in range(B): + n = pages_per_req[b] + page_indices[b, :n] = perm[cursor : cursor + n].to(torch.int32) + cursor += n + q_indptr = torch.zeros(B + 1, dtype=torch.int32) + q_indptr[1:] = torch.cumsum(torch.tensor(q_lens, dtype=torch.int32), 0) + T = int(q_indptr[-1]) + width = (num_q_heads + 2 * num_kv_heads) * HEAD_DIM + qkv = ( + torch.randn((T, width), generator=g, dtype=torch.float32) + .to(torch.bfloat16) + .to(device) + ) + q_w = (torch.rand(HEAD_DIM, generator=g) + 0.5).to(torch.bfloat16).float() + k_w = (torch.rand(HEAD_DIM, generator=g) + 0.5).to(torch.bfloat16).float() + if is_prefill is None: + is_prefill = max(q_lens) > 1 + if max_seqlen is None: + max_seqlen = max(q_lens) if quant_policy == 1 and is_prefill else 0 + cache_shape = (total_pages, page_size, num_kv_heads, HEAD_DIM) + return dict( + qkv=qkv, + cos_sin=rotary_cos_sin_table(max(seq_lens) + 1, device), + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), + q_indptr=q_indptr.to(device), + page_indices=page_indices.to(device), + cache_shape=cache_shape, + page_size=page_size, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + qk_norm_policy=qk_norm_policy, + quant_policy=quant_policy, + is_prefill=is_prefill, + max_seqlen=max_seqlen, + upper_max=upper_max, + q_norm_weight=q_w.to(device), + k_norm_weight=k_w.to(device), + k_scale=torch.tensor([0.02], dtype=torch.float32, device=device), + v_scale=torch.tensor([0.03], dtype=torch.float32, device=device), + q_scale_inv=torch.tensor([1.0 / 0.02], dtype=torch.float32, device=device), + ) + + +def _poisoned_fp8(shape, device): + t = torch.empty(shape, dtype=torch.float8_e4m3fn, device=device) + t.view(torch.uint8).fill_(POISON_BYTE) + return t + + +def _q_scale_buffer(case): + hq, B, T = case["num_q_heads"], case["seq_lens"].shape[0], case["qkv"].shape[0] + dev = case["qkv"].device + if case["quant_policy"] == 2: + return torch.empty(0, dtype=torch.float32, device=dev) + if case["is_prefill"]: + aligned = (case["max_seqlen"] + 127) // 128 * 128 + return torch.full((B, hq, aligned), float("nan"), device=dev) + return torch.full((T, hq), float("nan"), device=dev) + + +def launch(case, *, out_k=None, out_v=None, q_indptr=None, max_seqlen=None): + """Run the fused entry on poisoned caller-owned buffers; return (buffers, returned tuple).""" + dev = case["qkv"].device + T, hq, hkv = case["qkv"].shape[0], case["num_q_heads"], case["num_kv_heads"] + bufs = dict( + key_cache=_poisoned_fp8(case["cache_shape"], dev), + value_cache=_poisoned_fp8(case["cache_shape"], dev), + out_q=_poisoned_fp8((T, hq, HEAD_DIM), dev), + q_scale=_q_scale_buffer(case), + flags=torch.full( + (case["seq_lens"].shape[0], hkv), 7, dtype=torch.int32, device=dev + ), + ) + norm = case["qk_norm_policy"] + ret = fused_fp8( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"] if q_indptr is None else q_indptr, + case["page_indices"], + (bufs["key_cache"], bufs["value_cache"]), + case["is_prefill"], + case["k_scale"], + case["v_scale"], + case["quant_policy"], + max_seqlen=case["max_seqlen"] if max_seqlen is None else max_seqlen, + upper_max=case["upper_max"], + q_scale_inv=case["q_scale_inv"] if case["quant_policy"] == 2 else None, + q_norm_weight=case["q_norm_weight"] if norm else None, + k_norm_weight=case["k_norm_weight"] if norm else None, + qk_norm_policy=norm, + out_q=bufs["out_q"], + out_k=out_k, + out_v=out_v, + q_scale=bufs["q_scale"], + split_k_flag=bufs["flags"], + ) + torch.cuda.synchronize() + return bufs, ret + + +def _assert_fp8_matches(got_fp8, expected_f32, label): + """Written elements within the FP8 tolerance; NaN in ``expected`` means untouched poison.""" + written = ~torch.isnan(expected_f32) + torch.testing.assert_close( + got_fp8.float()[written], + expected_f32[written], + atol=0.1, + rtol=0.1, + msg=lambda m: f"{label}: {m}", + ) + untouched = got_fp8.view(torch.uint8)[~written] + assert bool((untouched == POISON_BYTE).all()), f"{label}: wrote outside its range" + + +def run_and_compare(case, *, caller_owned_kv=False): + dev = case["qkv"].device + T, hkv = case["qkv"].shape[0], case["num_kv_heads"] + ok = ov = None + if caller_owned_kv: + ok = _poisoned_fp8((T, hkv, HEAD_DIM), dev) + ov = _poisoned_fp8((T, hkv, HEAD_DIM), dev) + bufs, ret = launch(case, out_k=ok, out_v=ov) + assert ret[0] is bufs["out_q"] and ret[1] is bufs["q_scale"] + assert ret[2] is bufs["flags"] + exp = reference(case, caller_owned_kv=caller_owned_kv) + assert exp["valid"] + assert torch.equal(bufs["flags"], exp["flags"]) + _assert_fp8_matches(bufs["out_q"], exp["out_q"], "out_q") + if case["quant_policy"] == 1: + valid = ~torch.isnan(exp["q_scale"]) + torch.testing.assert_close( + bufs["q_scale"][valid], exp["q_scale"][valid], rtol=1e-5, atol=1e-6 + ) + # the padding of the dynamic-prefill layout is never written + assert bool(torch.isnan(bufs["q_scale"][~valid]).all()) + else: + assert bufs["q_scale"].numel() == 0 + _assert_fp8_matches(bufs["key_cache"], exp["key_cache"], "key_cache") + _assert_fp8_matches(bufs["value_cache"], exp["value_cache"], "value_cache") + if caller_owned_kv: + _assert_fp8_matches(ok, exp["out_k"], "out_k") + _assert_fp8_matches(ov, exp["out_v"], "out_v") + return bufs + + +@pytest.mark.parametrize("heads", [(8, 1), (64, 8)]) +@pytest.mark.parametrize("qk_norm_policy", [0, 1, 2]) +@pytest.mark.parametrize("quant_policy", [1, 2]) +@pytest.mark.parametrize("page_size", [16, 64]) +def test_decode_batch(heads, qk_norm_policy, quant_policy, page_size): + _skip_unless_supported() + case = make_case( + q_lens=[1] * 6, + ctx_lens=[0, 5, 63, 64, 130, 2047], + num_q_heads=heads[0], + num_kv_heads=heads[1], + page_size=page_size, + qk_norm_policy=qk_norm_policy, + quant_policy=quant_policy, + seed=893 + page_size, + ) + run_and_compare(case) + + +@pytest.mark.parametrize("heads", [(8, 1), (64, 8)]) +@pytest.mark.parametrize("qk_norm_policy", [0, 1, 2]) +@pytest.mark.parametrize("quant_policy", [1, 2]) +def test_prefill_chunks_cross_pages(heads, qk_norm_policy, quant_policy): + _skip_unless_supported() + case = make_case( + q_lens=[16, 0, 7, 70], + ctx_lens=[60, 3, 0, 125], + num_q_heads=heads[0], + num_kv_heads=heads[1], + page_size=64, + qk_norm_policy=qk_norm_policy, + quant_policy=quant_policy, + seed=71, + ) + run_and_compare(case) + + +@pytest.mark.parametrize("heads", [(8, 1), (64, 8)]) +@pytest.mark.parametrize("q_len", [127, 128, 129]) +def test_dynamic_prefill_scale_padding_boundary(heads, q_len): + """``q_scale`` padding (``round_up(max_seqlen, 128)``) stays untouched at the boundary.""" + _skip_unless_supported() + case = make_case( + q_lens=[q_len] * 3, + ctx_lens=[0, 64, 1000], + num_q_heads=heads[0], + num_kv_heads=heads[1], + page_size=64, + qk_norm_policy=2, + quant_policy=1, + seed=q_len, + max_seqlen=q_len, + ) + run_and_compare(case) + + +@pytest.mark.parametrize("quant_policy", [1, 2]) +def test_ragged_with_empty_requests_and_caller_owned_kv(quant_policy): + _skip_unless_supported() + case = make_case( + q_lens=[3, 0, 2, 0], + ctx_lens=[2, 5, 2, 0], + num_q_heads=8, + num_kv_heads=1, + page_size=4, + qk_norm_policy=2, + quant_policy=quant_policy, + seed=5, + ) + run_and_compare(case, caller_owned_kv=True) + run_and_compare(case, caller_owned_kv=False) + + +def test_small_upper_max_and_small_page(): + _skip_unless_supported() + case = make_case( + q_lens=[1, 1, 1], + ctx_lens=[7, 8, 9], + num_q_heads=8, + num_kv_heads=1, + page_size=4, + qk_norm_policy=1, + quant_policy=1, + seed=11, + upper_max=240.0, + ) + run_and_compare(case) + + +@pytest.mark.parametrize( + "kind", ["first_nonzero", "non_monotone", "last_mismatch", "overflow_max_seqlen"] +) +def test_invalid_metadata_leaves_outputs_untouched(kind): + _skip_unless_supported() + case = make_case( + q_lens=[4, 1, 3], + ctx_lens=[1, 3, 60], + num_q_heads=8, + num_kv_heads=1, + page_size=16, + qk_norm_policy=2, + quant_policy=1, + seed=9, + max_seqlen=4, + ) + dev = case["qkv"].device + q_indptr, max_seqlen = case["q_indptr"], None + if kind == "first_nonzero": + q_indptr = torch.tensor([1, 4, 5, 8], dtype=torch.int32, device=dev) + elif kind == "non_monotone": + q_indptr = torch.tensor([0, 5, 4, 8], dtype=torch.int32, device=dev) + elif kind == "last_mismatch": + q_indptr = torch.tensor([0, 4, 5, 9], dtype=torch.int32, device=dev) + else: + max_seqlen = 3 # request 0 has 4 rows + bufs, _ = launch(case, q_indptr=q_indptr, max_seqlen=max_seqlen) + assert bool((bufs["flags"] == -1).all()) + for name in ("out_q", "key_cache", "value_cache"): + assert bool((bufs[name].view(torch.uint8) == POISON_BYTE).all()), name + assert bool(torch.isnan(bufs["q_scale"]).all()) + + +def test_cuda_graph_replay_matches_eager(): + _skip_unless_supported() + case = make_case( + q_lens=[16] * 4, + ctx_lens=[2048, 100, 0, 63], + num_q_heads=64, + num_kv_heads=8, + page_size=64, + qk_norm_policy=2, + quant_policy=1, + seed=897, + ) + dev = case["qkv"].device + T = case["qkv"].shape[0] + kc = _poisoned_fp8(case["cache_shape"], dev) + vc = _poisoned_fp8(case["cache_shape"], dev) + oq = _poisoned_fp8((T, 64, HEAD_DIM), dev) + qs = _q_scale_buffer(case) + flags = torch.full((4, 8), 7, dtype=torch.int32, device=dev) + + def run(): + fused_fp8( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + (kc, vc), + True, + case["k_scale"], + case["v_scale"], + 1, + max_seqlen=case["max_seqlen"], + q_norm_weight=case["q_norm_weight"], + k_norm_weight=case["k_norm_weight"], + qk_norm_policy=2, + out_q=oq, + q_scale=qs, + split_k_flag=flags, + ) + + def poison(): + for t in (kc, vc, oq): + t.view(torch.uint8).fill_(POISON_BYTE) + qs.fill_(float("nan")) + flags.fill_(7) + + def snapshot(): + torch.cuda.synchronize() + return [ + t.view(torch.uint8).clone() if t.dtype == torch.float8_e4m3fn else t.clone() + for t in (oq, qs, flags, kc, vc) + ] + + poison() + run() + eager = snapshot() + assert bool((flags == 0).all()) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() # warm the JIT module / allocator on the capture stream + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + run() + torch.cuda.current_stream().wait_stream(stream) + for _ in range(3): + poison() + graph.replay() + for got, want in zip(snapshot(), eager, strict=True): + assert torch.equal( + got.nan_to_num(-1.0) if got.is_floating_point() else got, + want.nan_to_num(-1.0) if want.is_floating_point() else want, + ) + + +def test_rejects_bad_arguments(): + _skip_unless_supported() + case = make_case( + q_lens=[2, 1], + ctx_lens=[5, 70], + num_q_heads=8, + num_kv_heads=1, + page_size=64, + qk_norm_policy=2, + quant_policy=1, + seed=1, + ) + dev = case["qkv"].device + T = case["qkv"].shape[0] + kc = _poisoned_fp8(case["cache_shape"], dev) + vc = _poisoned_fp8(case["cache_shape"], dev) + + def call(**over): + kw = dict( + max_seqlen=2, + q_norm_weight=case["q_norm_weight"], + k_norm_weight=case["k_norm_weight"], + qk_norm_policy=2, + ) + pos = dict( + qkv=case["qkv"], + cos_sin_cache=case["cos_sin"], + seq_lens=case["seq_lens"], + q_indptr=case["q_indptr"], + page_indices=case["page_indices"], + paged_kv_cache=(kc, vc), + is_prefill=True, + k_scale=case["k_scale"], + v_scale=case["v_scale"], + quant_policy=1, + ) + pos.update({k: v for k, v in over.items() if k in pos}) + kw.update({k: v for k, v in over.items() if k not in pos}) + return fused_fp8(**pos, **kw) + + for upper_max in (0.0, -1.0, 448.01, float("inf"), float("nan")): + with pytest.raises(ValueError): + call(upper_max=upper_max) + with pytest.raises(ValueError): + call(quant_policy=3) + with pytest.raises(ValueError): + call(quant_policy=2) # q_scale_inv missing + with pytest.raises(ValueError): + call(qk_norm_policy=1, q_norm_weight=None) + with pytest.raises(ValueError): + call(max_seqlen=0) # dynamic prefill needs a positive max_seqlen + with pytest.raises(ValueError): + call(out_q=_poisoned_fp8((T, 4, HEAD_DIM), dev)) + with pytest.raises(ValueError): + call(q_scale=torch.empty((T, 8), device=dev)) # decode layout on prefill + with pytest.raises(ValueError): + call(split_k_flag=torch.empty((2, 2), dtype=torch.int32, device=dev)) + with pytest.raises(ValueError): + call(out_k=_poisoned_fp8((T, 1, HEAD_DIM), dev)) # out_v missing + with pytest.raises(ValueError): + call(qkv=case["qkv"][:, : 9 * HEAD_DIM].contiguous()) # (7, 1) heads + # the valid call still works after the rejections + out_q, q_scale, flags = call() + torch.cuda.synchronize() + assert out_q.shape == (T, 8, HEAD_DIM) and q_scale.shape == (2, 8, 128) + assert bool((flags == 0).all())