diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d87508cb52f..5860744c167 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -70,7 +70,7 @@ repos: # Source-export directories are checksum-attested generated artifacts; # formatting them would break their static source-registry identities. exclude: | - (?x)^(3rdparty/.*|csrc/cake_bf16_bmm/.*|csrc/fused_moe/warp_decode/.*|flashinfer/jit/aot_config.py|csrc/concat_mla/.*|csrc/cake_vsa/.*|csrc/cake_grouped_mxfp8_quantize/.*|csrc/cake_sage_block_sparse_attention/(?:sm_100a|sm_103a)/.*|csrc/dcp/.*|csrc/generated/.*|csrc/kda/(?:tf32|bf16)/.*|csrc/alphamoe_nvfp4_(?:sm100|c[0-9]+_[a-z_]+)\.cu)$ + (?x)^(3rdparty/.*|csrc/cake_bf16_bmm/.*|csrc/fused_moe/warp_decode/.*|flashinfer/jit/aot_config.py|csrc/concat_mla/.*|csrc/cake_vsa/.*|csrc/cake_grouped_mxfp8_quantize/.*|csrc/cake_fused_qk_rope_append/.*|csrc/cake_sage_block_sparse_attention/(?:sm_100a|sm_103a)/.*|csrc/dcp/.*|csrc/generated/.*|csrc/kda/(?:tf32|bf16)/.*|csrc/alphamoe_nvfp4_(?:sm100|c[0-9]+_[a-z_]+)\.cu)$ - repo: https://github.com/pre-commit/mirrors-mypy rev: 'v1.17.1' # Use the sha / tag you want to point at diff --git a/benchmarks/bench_cake_fused_qk_rope_append.py b/benchmarks/bench_cake_fused_qk_rope_append.py new file mode 100644 index 00000000000..de550aecf63 --- /dev/null +++ b/benchmarks/bench_cake_fused_qk_rope_append.py @@ -0,0 +1,392 @@ +"""Benchmark the Cake fused QK RMSNorm + NeoX RoPE + paged KV append (BF16) kernel. + +Times the Cake route (``cake_fused_qk_rmsnorm_rope_append_paged_kv_cache``) against the stable +FlashInfer pipeline (``rmsnorm`` x2 + ``apply_rope_pos_ids_inplace`` + ``append_paged_kv_cache``) and, +optionally, a baseline fused implementation imported from ``--baseline-module`` (for example the +FlashInfer PR #5405 package ``flashinfer.experimental.fused_qk_rope_append`` when running from a checkout +that contains it). All arms see the same inputs, the same BF16-rounded norm weights and their own +poison-filled caches / outputs. + +Timing: ``flashinfer.testing.bench_gpu_time`` (CUPTI, cold L2) over the whole call (union of every +launch of one arm), eager and CUDA-graph replay, paired A/B/B/A groups. + +Usage:: + + python benchmarks/bench_cake_fused_qk_rope_append.py # default rows + python benchmarks/bench_cake_fused_qk_rope_append.py --rows M1,M2,H1 --json out.json + python benchmarks/bench_cake_fused_qk_rope_append.py --baseline-module flashinfer.experimental.fused_qk_rope_append +""" + +from __future__ import annotations + +import argparse +import importlib +import json +import statistics +from dataclasses import dataclass + +import torch + +import flashinfer +from flashinfer.cake_fused_qk_rope_append import ( + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache, +) +from flashinfer.testing.utils import bench_gpu_time + +HEAD_DIM = 128 +POISON = 1234.0 + + +@dataclass(frozen=True) +class Row: + name: str + hq: int + hkv: int + batch: int + q_len: int + ctx: int + page_size: int = 64 + policy: int = 2 + + +ROWS = { + r.name: r + for r in ( + Row("M1", 8, 1, 32, 1, 2048), + Row("M1a", 8, 1, 32, 1, 2048, policy=0), + Row("M1b", 8, 1, 32, 1, 2048, policy=1), + Row("M2", 8, 1, 8, 16, 2048), + Row("H1", 64, 8, 32, 1, 2048), + Row("H2", 64, 8, 8, 16, 2048), + Row("H3", 64, 8, 128, 1, 2048), + Row("H4", 64, 8, 1, 128, 2048), + Row("D1", 8, 1, 1, 1, 2048), + Row("D3", 8, 1, 128, 1, 2048), + Row("D4", 8, 1, 256, 1, 2048), + Row("P1", 8, 1, 1, 128, 2048), + Row("P2", 8, 1, 4, 64, 2048), + Row("E1", 8, 1, 32, 1, 2048, page_size=16), + Row("E2", 8, 1, 32, 1, 2048, page_size=128), + Row("E3", 8, 1, 32, 1, 0), + Row("E4", 8, 1, 8, 16, 2047), + ) +} + + +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 = 892): + 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, + "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), + "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"] + kc = torch.full( + (inp["total_pages"], row.page_size, row.hkv, HEAD_DIM), + POISON, + dtype=torch.bfloat16, + device=device, + ) + vc = torch.full_like(kc, POISON) + oq = torch.full( + (inp["T"], row.hq, HEAD_DIM), POISON, dtype=torch.bfloat16, device=device + ) + return kc, vc, oq + + +def make_cake_call(inp, device): + row = inp["row"] + kc, vc, oq = fresh_buffers(inp, device) + policy = row.policy + + def call(): + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + inp["qkv"], + inp["cos_sin"], + inp["seq_lens"], + inp["q_indptr"], + inp["page_indices"], + kc, + vc, + num_q_heads=row.hq, + num_kv_heads=row.hkv, + qk_norm_policy=policy, + q_norm_weight=inp["q_norm_weight"] if policy else None, + k_norm_weight=inp["k_norm_weight"] if policy else None, + out_q=oq, + ) + + return call, (kc, vc, oq) + + +def make_stable_call(inp, device): + """rmsnorm(q) + rmsnorm(k) + apply_rope_pos_ids_inplace + append_paged_kv_cache (no last-page clear).""" + row = inp["row"] + kc, vc, oq = 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) + norm_k = torch.empty_like(k) + tmp_q = torch.empty_like(q) + tmp_k = torch.empty_like(k) + qw, kw = inp["q_norm_weight_bf16"], inp["k_norm_weight_bf16"] + paged = (kc, vc) + + def call(): + if row.policy == 2: + flashinfer.norm.rmsnorm( + q.reshape(-1, HEAD_DIM), qw, out=oq.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( + oq, norm_k, inp["positions"], interleave=False + ) + key_src = norm_k + elif row.policy == 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=oq.view(-1, HEAD_DIM) + ) + flashinfer.norm.rmsnorm( + tmp_k.view(-1, HEAD_DIM), kw, out=norm_k.view(-1, HEAD_DIM) + ) + key_src = norm_k + else: + oq.copy_(q) + tmp_k.copy_(k) + flashinfer.rope.apply_rope_pos_ids_inplace( + oq, tmp_k, inp["positions"], interleave=False + ) + key_src = tmp_k + flashinfer.page.append_paged_kv_cache( + key_src, + v, + inp["batch_indices"], + inp["positions"], + paged, + inp["kv_indices"], + inp["kv_indptr"], + inp["kv_last_page_len"], + kv_layout="NHD", + ) + + return call, (kc, vc, oq) + + +def make_baseline_call(inp, device, module_name: str): + """Fused baseline from an importable package exposing ``fused_qk_norm_rope_append_paged_kv_cache``. + + The pinned PR #5405 entry point is positional: + ``fn(qkv, cos_sin, seq_lens, q_indptr, page_indices, (k_cache, v_cache), is_prefill, + q_norm_weight, k_norm_weight, qk_norm_policy, out_q=..., out_k=..., out_v=...)``. + """ + mod = importlib.import_module(module_name) + fn = getattr(mod, "fused_qk_norm_rope_append_paged_kv_cache", None) + if fn is None: + fn = importlib.import_module( + module_name + ".backend" + ).fused_qk_norm_rope_append_paged_kv_cache + row = inp["row"] + kc, vc, oq = fresh_buffers(inp, device) + cache = (kc, vc) + is_prefill = row.q_len > 1 + + def call(): + fn( + inp["qkv"], + inp["cos_sin"], + inp["seq_lens"], + inp["q_indptr"], + inp["page_indices"], + cache, + is_prefill, + inp["q_norm_weight"], + inp["k_norm_weight"], + row.policy, + out_q=oq, + ) + + return call, (kc, vc, oq) + + +def time_arm(call, *, graph: bool, warmup: int, iters: int): + if graph: + s = torch.cuda.Stream() + with torch.cuda.stream(s): + for _ in range(3): + call() + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, stream=s): + call() + torch.cuda.synchronize() + fn = g.replay + else: + fn = call + times = bench_gpu_time( + fn, + 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( + "--baseline-module", + default=None, + help="importable fused baseline package (e.g. PR #5405)", + ) + ap.add_argument("--json", default=None) + args = ap.parse_args() + device = torch.device("cuda:0") + arms = ( + ["cake"] + + ([] if args.no_stable else ["stable"]) + + (["baseline"] if args.baseline_module else []) + ) + results = [] + print( + f"device {torch.cuda.get_device_name(device)} | arms {arms} | groups {args.groups} x iters {args.iters}" + ) + header = f"{'row':5s} {'geometry':30s} " + " ".join( + f"{a + '-' + m:>14s}" for a in arms for m in ("graph", "eager") + ) + print(header) + for name in args.rows.split(","): + row = ROWS[name] + inp = build_inputs(row, device) + calls = {} + for arm in arms: + if arm == "cake": + calls[arm] = make_cake_call(inp, device)[0] + elif arm == "stable": + calls[arm] = make_stable_call(inp, device)[0] + else: + calls[arm] = make_baseline_call(inp, device, args.baseline_module)[0] + samples = {(a, m): [] for a in arms for m in ("graph", "eager")} + for g in range(args.groups): + order = arms if g % 2 == 0 else list(reversed(arms)) + for arm in order: + for mode in ("graph", "eager"): + samples[(arm, mode)].append( + time_arm( + calls[arm], + graph=(mode == "graph"), + warmup=args.warmup, + iters=args.iters, + ) + ) + med = {k: statistics.median(v) for k, v in samples.items()} + geometry = f"{row.hq}/{row.hkv} B{row.batch} Q{row.q_len} ctx{row.ctx} p{row.page_size} pol{row.policy}" + print( + f"{name:5s} {geometry:30s} " + + " ".join( + f"{med[(a, m)]:>11.2f} us" for a in arms for m in ("graph", "eager") + ) + ) + results.append( + { + "row": name, + "geometry": geometry, + "medians_us": {f"{a}-{m}": med[(a, m)] for a, m in med}, + } + ) + if args.json: + with open(args.json, "w") as f: + json.dump( + {"device": torch.cuda.get_device_name(device), "rows": results}, + f, + indent=1, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/csrc/cake_fused_qk_rope_append/cake_fused_qk_rope_append_manifest.json b/csrc/cake_fused_qk_rope_append/cake_fused_qk_rope_append_manifest.json new file mode 100644 index 00000000000..fa5a29be790 --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/cake_fused_qk_rope_append_manifest.json @@ -0,0 +1,1336 @@ +{ + "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_append.py", + "flashinfer/jit/cake_fused_qk_rope_append.py" + ], + "operator": "fused_qk_rmsnorm_rope_append_paged_kv_cache_bf16", + "stages": [ + "hq8_hkv1", + "hq64_hkv8" + ] + }, + "files": [ + { + "bytes": 10123, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_binding.cu", + "sha256": "48a2ea53024e2162a46850dd459a5d6d970254e19321f4cf1ce10f523a012662" + }, + { + "bytes": 33050, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_kernel.cu", + "sha256": "ff84f296c7139712cfb07c47a78417edb7487d91323f0d699b6af000a81732ab" + }, + { + "bytes": 10123, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_binding.cu", + "sha256": "2a700ac9bdfa520c5b76a4c49b842f0cd7613f0a0d950932713f43ebf404d601" + }, + { + "bytes": 33049, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_kernel.cu", + "sha256": "e0cdfbc47cf74a7491424439923541e0dfda0db0eaa6e952f83e11a6c075c180" + }, + { + "bytes": 10123, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_binding.cu", + "sha256": "9165543e304f17597deebf04c1ac90f179e3591d68cec63a67a578b117ca26d3" + }, + { + "bytes": 33049, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_kernel.cu", + "sha256": "3ec4feae9174c1006523c580c12947e2e12815e955b79d90a27f85f6169a1a9d" + }, + { + "bytes": 10123, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_binding.cu", + "sha256": "3f884216bb85caa4ad61938584e1fecba2738af7f08ebfd2a2412bcaa546f56f" + }, + { + "bytes": 33050, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_kernel.cu", + "sha256": "f8fa130ebde8cd98ae9653f1c4efd07eb5c2c6f6307fee536290e63c6f4e8179" + }, + { + "bytes": 10123, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_binding.cu", + "sha256": "4cda5d17051e12ee2bb9d32c2b0e1071e0e515ef5de282e448c2a07705300a0e" + }, + { + "bytes": 33050, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_kernel.cu", + "sha256": "4ee378e944572b306ef812a03c553cb38e7320415f808eaee9f6ecf19aef25a2" + }, + { + "bytes": 10123, + "kind": "tvm_ffi_binding", + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_binding.cu", + "sha256": "3ccd531a0d993c15cc86f2e220cfd773a2d6f01ec502287d6f6e2198570ff875" + }, + { + "bytes": 33049, + "kind": "device_source", + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_kernel.cu", + "sha256": "28632ce3636e9c775515dbec9a8b157e453c47b5e4cbfad45d4437786b743d93" + } + ], + "generation_inputs": { + "schema": "cake.library_export.generation_inputs.v4" + }, + "library": "flashinfer", + "modules": [ + { + "arch": "sm_100a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "page_table" + ], + [ + "buffer", + "k_cache" + ], + [ + "buffer", + "v_cache" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "max_pages" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "k_page_stride" + ], + [ + "parameter", + "v_page_stride" + ], + [ + "parameter", + "num_row_ctas" + ], + [ + "parameter", + "clear_units_per_request" + ], + [ + "parameter", + "qk_norm_policy" + ], + [ + "parameter", + "eps" + ], + [ + "parameter", + "use_out_k" + ], + [ + "parameter", + "use_out_v" + ], + [ + "parameter", + "clear_last_page" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "10a1c8ff264ea5959c5af56b76bd2ee8e8ba6e5bdecaa21223964f4c9a7e2874", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 33050, + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_kernel.cu", + "sha256": "ff84f296c7139712cfb07c47a78417edb7487d91323f0d699b6af000a81732ab" + }, + { + "bytes": 10123, + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_binding.cu", + "sha256": "48a2ea53024e2162a46850dd459a5d6d970254e19321f4cf1ce10f523a012662" + } + ], + "closure_sha256": "2d489e4f1031730a801f9c2f4c7973ef14c0b4b4a3f345079830e98d876d1ff9", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_append_1a5e46c8842b05e59fbb", + "launch": { + "block": [ + 128, + 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_append_1a5e46c8842b05e59fbb_sm_100a", + "name": "cake_fused_qk_rope_append_1a5e46c8842b05e59fbb", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "1a5e46c8842b05e59fbb73b9259a4aacd7536f2f1110020cabd4d0b6cb1b4767", + "specialization": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "stage": "hq64_hkv8", + "template": "fused_qk_rmsnorm_rope_paged_kv_append_bf16_v3" + }, + "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": "97039feee98500b89317ef81641a89e5b233d95c0bbea09ae01a2aa0ae1e0a00", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "use_pdl": false, + "validate": false + }, + "specializations": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_kernel.cu" + } + }, + { + "arch": "sm_100a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "page_table" + ], + [ + "buffer", + "k_cache" + ], + [ + "buffer", + "v_cache" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "max_pages" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "k_page_stride" + ], + [ + "parameter", + "v_page_stride" + ], + [ + "parameter", + "num_row_ctas" + ], + [ + "parameter", + "clear_units_per_request" + ], + [ + "parameter", + "qk_norm_policy" + ], + [ + "parameter", + "eps" + ], + [ + "parameter", + "use_out_k" + ], + [ + "parameter", + "use_out_v" + ], + [ + "parameter", + "clear_last_page" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "10a1c8ff264ea5959c5af56b76bd2ee8e8ba6e5bdecaa21223964f4c9a7e2874", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 33049, + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_kernel.cu", + "sha256": "e0cdfbc47cf74a7491424439923541e0dfda0db0eaa6e952f83e11a6c075c180" + }, + { + "bytes": 10123, + "path": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_binding.cu", + "sha256": "2a700ac9bdfa520c5b76a4c49b842f0cd7613f0a0d950932713f43ebf404d601" + } + ], + "closure_sha256": "755971ce89595e2ddea756d718d58329bbe9dea6908c76710056271d49bc7c00", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_append_51a814b69e055ec5993b", + "launch": { + "block": [ + 128, + 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_append_51a814b69e055ec5993b_sm_100a", + "name": "cake_fused_qk_rope_append_51a814b69e055ec5993b", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "51a814b69e055ec5993bfd46b5e2a1559047de6ff0f17a67b7dd4f2bfdda19e8", + "specialization": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "stage": "hq8_hkv1", + "template": "fused_qk_rmsnorm_rope_paged_kv_append_bf16_v3" + }, + "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": "44ab2c641f35a3cf09d987c3595efc348cc20b05164d8641550b82b773cfe8a4", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "use_pdl": false, + "validate": false + }, + "specializations": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_kernel.cu" + } + }, + { + "arch": "sm_103a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "page_table" + ], + [ + "buffer", + "k_cache" + ], + [ + "buffer", + "v_cache" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "max_pages" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "k_page_stride" + ], + [ + "parameter", + "v_page_stride" + ], + [ + "parameter", + "num_row_ctas" + ], + [ + "parameter", + "clear_units_per_request" + ], + [ + "parameter", + "qk_norm_policy" + ], + [ + "parameter", + "eps" + ], + [ + "parameter", + "use_out_k" + ], + [ + "parameter", + "use_out_v" + ], + [ + "parameter", + "clear_last_page" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "10a1c8ff264ea5959c5af56b76bd2ee8e8ba6e5bdecaa21223964f4c9a7e2874", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 33049, + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_kernel.cu", + "sha256": "3ec4feae9174c1006523c580c12947e2e12815e955b79d90a27f85f6169a1a9d" + }, + { + "bytes": 10123, + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_binding.cu", + "sha256": "9165543e304f17597deebf04c1ac90f179e3591d68cec63a67a578b117ca26d3" + } + ], + "closure_sha256": "aaa6b89e4ca1a1edadbb1a57770d610dd300c9307b8f536e3092759e931f79a7", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_append_786a6a33470bb271f31a", + "launch": { + "block": [ + 128, + 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_append_786a6a33470bb271f31a_sm_103a", + "name": "cake_fused_qk_rope_append_786a6a33470bb271f31a", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "786a6a33470bb271f31a65d789c1de28d2f73653d585132ec9a3c832e5d735db", + "specialization": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "stage": "hq8_hkv1", + "template": "fused_qk_rmsnorm_rope_paged_kv_append_bf16_v3" + }, + "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": "44ab2c641f35a3cf09d987c3595efc348cc20b05164d8641550b82b773cfe8a4", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "use_pdl": false, + "validate": false + }, + "specializations": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_kernel.cu" + } + }, + { + "arch": "sm_103a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "page_table" + ], + [ + "buffer", + "k_cache" + ], + [ + "buffer", + "v_cache" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "max_pages" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "k_page_stride" + ], + [ + "parameter", + "v_page_stride" + ], + [ + "parameter", + "num_row_ctas" + ], + [ + "parameter", + "clear_units_per_request" + ], + [ + "parameter", + "qk_norm_policy" + ], + [ + "parameter", + "eps" + ], + [ + "parameter", + "use_out_k" + ], + [ + "parameter", + "use_out_v" + ], + [ + "parameter", + "clear_last_page" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "10a1c8ff264ea5959c5af56b76bd2ee8e8ba6e5bdecaa21223964f4c9a7e2874", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 33050, + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_kernel.cu", + "sha256": "f8fa130ebde8cd98ae9653f1c4efd07eb5c2c6f6307fee536290e63c6f4e8179" + }, + { + "bytes": 10123, + "path": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_binding.cu", + "sha256": "3f884216bb85caa4ad61938584e1fecba2738af7f08ebfd2a2412bcaa546f56f" + } + ], + "closure_sha256": "c9dd89f180030a80fe20e6cb95dd5fe1ef32aeb5d17d2cfe691f1e2d737b3200", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_append_d46a979444ec64e1bda3", + "launch": { + "block": [ + 128, + 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_append_d46a979444ec64e1bda3_sm_103a", + "name": "cake_fused_qk_rope_append_d46a979444ec64e1bda3", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "d46a979444ec64e1bda360a43305dfedfbd70e499cdcb3832e6fd49a8929802e", + "specialization": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "stage": "hq64_hkv8", + "template": "fused_qk_rmsnorm_rope_paged_kv_append_bf16_v3" + }, + "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": "97039feee98500b89317ef81641a89e5b233d95c0bbea09ae01a2aa0ae1e0a00", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "use_pdl": false, + "validate": false + }, + "specializations": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_kernel.cu" + } + }, + { + "arch": "sm_90a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "page_table" + ], + [ + "buffer", + "k_cache" + ], + [ + "buffer", + "v_cache" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "max_pages" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "k_page_stride" + ], + [ + "parameter", + "v_page_stride" + ], + [ + "parameter", + "num_row_ctas" + ], + [ + "parameter", + "clear_units_per_request" + ], + [ + "parameter", + "qk_norm_policy" + ], + [ + "parameter", + "eps" + ], + [ + "parameter", + "use_out_k" + ], + [ + "parameter", + "use_out_v" + ], + [ + "parameter", + "clear_last_page" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "10a1c8ff264ea5959c5af56b76bd2ee8e8ba6e5bdecaa21223964f4c9a7e2874", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 33050, + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_kernel.cu", + "sha256": "4ee378e944572b306ef812a03c553cb38e7320415f808eaee9f6ecf19aef25a2" + }, + { + "bytes": 10123, + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_binding.cu", + "sha256": "4cda5d17051e12ee2bb9d32c2b0e1071e0e515ef5de282e448c2a07705300a0e" + } + ], + "closure_sha256": "837bac668db764d323b09dfe906873da6f1f30c40e7efa0a3bc498339f87b85b", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9", + "launch": { + "block": [ + 128, + 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_append_81fc2fe5ac8bb14706e9_sm_90a", + "name": "cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "81fc2fe5ac8bb14706e9943055a86f33db40cc20288de7e4bb85fea26fe3b3fe", + "specialization": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "stage": "hq64_hkv8", + "template": "fused_qk_rmsnorm_rope_paged_kv_append_bf16_v3" + }, + "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": "97039feee98500b89317ef81641a89e5b233d95c0bbea09ae01a2aa0ae1e0a00", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "use_pdl": false, + "validate": false + }, + "specializations": { + "NUM_KV_HEADS": 8, + "NUM_Q_HEADS": 64, + "WARPS_PER_ROW": 8 + }, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_kernel.cu" + } + }, + { + "arch": "sm_90a", + "arg_plan": [ + [ + "buffer", + "qkv" + ], + [ + "buffer", + "cos_sin" + ], + [ + "buffer", + "q_indptr" + ], + [ + "buffer", + "seq_lens" + ], + [ + "buffer", + "page_table" + ], + [ + "buffer", + "k_cache" + ], + [ + "buffer", + "v_cache" + ], + [ + "buffer", + "out_q" + ], + [ + "buffer", + "out_k" + ], + [ + "buffer", + "out_v" + ], + [ + "buffer", + "q_norm_weight" + ], + [ + "buffer", + "k_norm_weight" + ], + [ + "parameter", + "num_rows" + ], + [ + "parameter", + "num_requests" + ], + [ + "parameter", + "max_pages" + ], + [ + "parameter", + "page_size" + ], + [ + "parameter", + "k_page_stride" + ], + [ + "parameter", + "v_page_stride" + ], + [ + "parameter", + "num_row_ctas" + ], + [ + "parameter", + "clear_units_per_request" + ], + [ + "parameter", + "qk_norm_policy" + ], + [ + "parameter", + "eps" + ], + [ + "parameter", + "use_out_k" + ], + [ + "parameter", + "use_out_v" + ], + [ + "parameter", + "clear_last_page" + ], + [ + "grid", + "grid_x" + ], + [ + "grid", + "grid_y" + ], + [ + "grid", + "grid_z" + ] + ], + "arg_plan_sha256": "10a1c8ff264ea5959c5af56b76bd2ee8e8ba6e5bdecaa21223964f4c9a7e2874", + "binding_infrastructure": { + "runtime": "flashinfer_tvm_ffi_utils", + "target_owned_headers": [ + "tvm_ffi_utils.h" + ] + }, + "closure": [ + { + "bytes": 33049, + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_kernel.cu", + "sha256": "28632ce3636e9c775515dbec9a8b157e453c47b5e4cbfad45d4437786b743d93" + }, + { + "bytes": 10123, + "path": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_binding.cu", + "sha256": "3ccd531a0d993c15cc86f2e220cfd773a2d6f01ec502287d6f6e2198570ff875" + } + ], + "closure_sha256": "f6f88a25ac990889b55b9b3a626ad2545eb874b6cf46fca41bc7ea6d461875c5", + "compile_flags": [], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_fused_qk_rope_append_ef57b221493533ee8018", + "launch": { + "block": [ + 128, + 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_append_ef57b221493533ee8018_sm_90a", + "name": "cake_fused_qk_rope_append_ef57b221493533ee8018", + "noncontiguous_tensors": [], + "role": "kernel", + "route": { + "source_identity": "ef57b221493533ee80185819ad156be83e180a76a1de472aace0f6722ed3112c", + "specialization": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "stage": "hq8_hkv1", + "template": "fused_qk_rmsnorm_rope_paged_kv_append_bf16_v3" + }, + "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": "44ab2c641f35a3cf09d987c3595efc348cc20b05164d8641550b82b773cfe8a4", + "source_arch": null, + "source_cluster_scheduling_policy": "spread", + "source_tma_abi": "grid_constant", + "specializations": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "use_pdl": false, + "validate": false + }, + "specializations": { + "NUM_KV_HEADS": 1, + "NUM_Q_HEADS": 8, + "WARPS_PER_ROW": 4 + }, + "tma_abi": "grid_constant", + "translation_units": { + "binding": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_binding.cu", + "compile_separately": true, + "device": "csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_kernel.cu" + } + } + ], + "name": "cake_fused_qk_rope_append", + "producer": "cake", + "producer_revision": null, + "public_namespace": "cake", + "schema": "cake.library_export.v5", + "sequences": [] +} diff --git a/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_binding.cu b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_binding.cu new file mode 100644 index 00000000000..6092f46905b --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_binding.cu @@ -0,0 +1,160 @@ +/* + * 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_append_1a5e46c8842b05e59fbb'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include + +extern "C" __global__ void kernel_cake_fused_qk_rope_append_1a5e46c8842b05e59fbb(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page); + + +namespace cake_host_shim_2fb4dbace7eb8f3f { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_q_indptr, TensorView arg_seq_lens, TensorView arg_page_table, TensorView arg_k_cache, TensorView arg_v_cache, TensorView arg_out_q, TensorView arg_out_k, TensorView arg_out_v, TensorView arg_q_norm_weight, TensorView arg_k_norm_weight, int64_t arg_num_rows, int64_t arg_num_requests, int64_t arg_max_pages, int64_t arg_page_size, int64_t arg_k_page_stride, int64_t arg_v_page_stride, int64_t arg_num_row_ctas, int64_t arg_clear_units_per_request, int64_t arg_qk_norm_policy, double arg_eps, int64_t arg_use_out_k, int64_t arg_use_out_v, int64_t arg_clear_last_page, 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_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_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_page_table, "page_table"); + check_dtype(arg_page_table, DLDataType{kDLInt, 32, 1}, "page_table"); + check_contiguous(arg_page_table, "page_table"); + check_cuda_tensor(arg_k_cache, "k_cache"); + check_dtype(arg_k_cache, DLDataType{kDLBfloat, 16, 1}, "k_cache"); + check_contiguous(arg_k_cache, "k_cache"); + check_cuda_tensor(arg_v_cache, "v_cache"); + check_dtype(arg_v_cache, DLDataType{kDLBfloat, 16, 1}, "v_cache"); + check_contiguous(arg_v_cache, "v_cache"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLBfloat, 16, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLBfloat, 16, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLBfloat, 16, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + 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"); + 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_max_pages >= -2147483648LL && arg_max_pages <= 2147483647LL, ValueError) + << "scalar 'max_pages' value " << arg_max_pages + << " 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_num_row_ctas >= -2147483648LL && arg_num_row_ctas <= 2147483647LL, ValueError) + << "scalar 'num_row_ctas' value " << arg_num_row_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_units_per_request >= -2147483648LL && arg_clear_units_per_request <= 2147483647LL, ValueError) + << "scalar 'clear_units_per_request' value " << arg_clear_units_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_norm_policy >= -2147483648LL && arg_qk_norm_policy <= 2147483647LL, ValueError) + << "scalar 'qk_norm_policy' value " << arg_qk_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_k >= -2147483648LL && arg_use_out_k <= 2147483647LL, ValueError) + << "scalar 'use_out_k' value " << arg_use_out_k + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_v >= -2147483648LL && arg_use_out_v <= 2147483647LL, ValueError) + << "scalar 'use_out_v' value " << arg_use_out_v + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_last_page >= -2147483648LL && arg_clear_last_page <= 2147483647LL, ValueError) + << "scalar 'clear_last_page' value " << arg_clear_last_page + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_page_table, arg_qkv, "page_table", "qkv"); + check_same_device(arg_k_cache, arg_qkv, "k_cache", "qkv"); + check_same_device(arg_v_cache, arg_qkv, "v_cache", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "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_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "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_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_page_table = static_cast(arg_page_table.data_ptr()); + __nv_bfloat16* p_k_cache = static_cast<__nv_bfloat16*>(arg_k_cache.data_ptr()); + __nv_bfloat16* p_v_cache = static_cast<__nv_bfloat16*>(arg_v_cache.data_ptr()); + __nv_bfloat16* p_out_q = static_cast<__nv_bfloat16*>(arg_out_q.data_ptr()); + __nv_bfloat16* p_out_k = static_cast<__nv_bfloat16*>(arg_out_k.data_ptr()); + __nv_bfloat16* p_out_v = static_cast<__nv_bfloat16*>(arg_out_v.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()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_max_pages = (int32_t)arg_max_pages; + int32_t v_page_size = (int32_t)arg_page_size; + int64_t v_k_page_stride = (int64_t)arg_k_page_stride; + int64_t v_v_page_stride = (int64_t)arg_v_page_stride; + int32_t v_num_row_ctas = (int32_t)arg_num_row_ctas; + int32_t v_clear_units_per_request = (int32_t)arg_clear_units_per_request; + int32_t v_qk_norm_policy = (int32_t)arg_qk_norm_policy; + float v_eps = (float)arg_eps; + int32_t v_use_out_k = (int32_t)arg_use_out_k; + int32_t v_use_out_v = (int32_t)arg_use_out_v; + int32_t v_clear_last_page = (int32_t)arg_clear_last_page; + void* kargs[] = {&p_qkv, &p_cos_sin, &p_q_indptr, &p_seq_lens, &p_page_table, &p_k_cache, &p_v_cache, &p_out_q, &p_out_k, &p_out_v, &p_q_norm_weight, &p_k_norm_weight, &v_num_rows, &v_num_requests, &v_max_pages, &v_page_size, &v_k_page_stride, &v_v_page_stride, &v_num_row_ctas, &v_clear_units_per_request, &v_qk_norm_policy, &v_eps, &v_use_out_k, &v_use_out_v, &v_clear_last_page}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_append_1a5e46c8842b05e59fbb), grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_append_1a5e46c8842b05e59fbb failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_2fb4dbace7eb8f3f + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_2fb4dbace7eb8f3f::Run); diff --git a/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_kernel.cu b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_kernel.cu new file mode 100644 index 00000000000..6437155365a --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_1a5e46c8842b05e59fbb_kernel.cu @@ -0,0 +1,608 @@ +/* + * 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 THREADS 128 +#define NUM_Q_HEADS 64 +#define NUM_KV_HEADS 8 +#define WARPS_PER_ROW 8 +#define SKIP_MATH 0 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_fused_qk_rope_append_1a5e46c8842b05e59fbb(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // === Task calls (dependency order) === + int lane_0 = lane; + int warp_1 = warp; + int hsel = lane_0 / 8; + int d0 = lane_0 % 8 * 8; + if (bid >= num_row_ctas) { + int cid = (bid - num_row_ctas) * 4 + warp_1; + if (cid < num_requests * clear_units_per_request && clear_last_page != 0) { + int req_c = cid / clear_units_per_request; + int u = cid % clear_units_per_request; + int cache_sel = u % 2; + int unit_begin = u / 2 * 256; + int seq_len_c = seq_lens[req_c]; + if (seq_len_c > 0) { + int last_c = seq_len_c - 1; + int start_chunk = (last_c % page_size + 1) * (NUM_KV_HEADS * 128 / 8); + int end_chunk = page_size * (NUM_KV_HEADS * 128 / 8); + if (start_chunk < unit_begin + 256 && unit_begin < end_chunk) { + int page_c32 = page_table[req_c * max_pages + last_c / page_size]; + long long page_c = (long long)page_c32; + float zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0.0f; + } + int c = 0; + long long off_c = 0; + if (cache_sel == 0) { + long long base_k = page_c * k_page_stride; + #pragma unroll + for (int k = 0; k < 8; k++) { + c = unit_begin + lane_0 + 32 * k; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (base_k + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else { + long long base_v = page_c * v_page_stride; + #pragma unroll + for (int k_1 = 0; k_1 < 8; k_1++) { + c = unit_begin + lane_0 + 32 * k_1; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(v_cache + (base_v + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } else { + int row = bid / ((WARPS_PER_ROW + 3) / 4) * ((4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW) + warp_1 / WARPS_PER_ROW; + int sub = bid % ((WARPS_PER_ROW + 3) / 4) * 4 + warp_1 % WARPS_PER_ROW; + if (row < num_rows) { + long long row_base = (long long)row * (long long)((NUM_Q_HEADS + 2 * NUM_KV_HEADS) * 128); + float xs[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 2 * 8]; + int head_of[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_h = 0; + #pragma unroll + for (int g = 0; g < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g++) { + head_of[g] = 4 * (sub + WARPS_PER_ROW * g) + hsel; + src_h = ((head_of[g] < NUM_Q_HEADS + NUM_KV_HEADS) ? head_of[g] : NUM_Q_HEADS + NUM_KV_HEADS - 1); + { + const uint4* _vptr_0 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0)) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_0[_blk].x), "=r"(_vld_0[_blk].y), "=r"(_vld_0[_blk].z), "=r"(_vld_0[_blk].w) : "l"((const void*)(_vptr_0 + _blk)) : "memory"); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #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"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0 + 64)) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_1[_blk].x), "=r"(_vld_1[_blk].y), "=r"(_vld_1[_blk].z), "=r"(_vld_1[_blk].w) : "l"((const void*)(_vptr_1 + _blk)) : "memory"); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #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"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + int vbits[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 8]; + int vhead_of[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_v = 0; + #pragma unroll + for (int g_1 = 0; g_1 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_1++) { + vhead_of[g_1] = 4 * (sub + WARPS_PER_ROW * g_1) + hsel; + src_v = ((vhead_of[g_1] < NUM_KV_HEADS) ? vhead_of[g_1] : NUM_KV_HEADS - 1); + { + const int4* _ivptr_2 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0)) + 0); + int4 _ivld_2; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_2.x), "=r"(_ivld_2.y), "=r"(_ivld_2.z), "=r"(_ivld_2.w) : "l"((const void*)(_ivptr_2)) : "memory"); + vbits[g_1 * 8 + 0] = _ivld_2.x; + vbits[g_1 * 8 + 1] = _ivld_2.y; + vbits[g_1 * 8 + 2] = _ivld_2.z; + vbits[g_1 * 8 + 3] = _ivld_2.w; + } + { + const int4* _ivptr_3 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0 + 64)) + 0); + int4 _ivld_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_3.x), "=r"(_ivld_3.y), "=r"(_ivld_3.z), "=r"(_ivld_3.w) : "l"((const void*)(_ivptr_3)) : "memory"); + vbits[g_1 * 8 + 4 + 0] = _ivld_3.x; + vbits[g_1 * 8 + 4 + 1] = _ivld_3.y; + vbits[g_1 * 8 + 4 + 2] = _ivld_3.z; + vbits[g_1 * 8 + 4 + 3] = _ivld_3.w; + } + } + float qw[16]; + float kw[16]; + if (qk_norm_policy != 0) { + { + 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 + d0 + (0))) : "memory"); + qw[0 + 0] = __uint_as_float(_v4_4_0); + qw[0 + 1] = __uint_as_float(_v4_4_1); + qw[0 + 2] = __uint_as_float(_v4_4_2); + qw[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 + (d0 + 4) + (0))) : "memory"); + qw[4 + 0] = __uint_as_float(_v4_5_0); + qw[4 + 1] = __uint_as_float(_v4_5_1); + qw[4 + 2] = __uint_as_float(_v4_5_2); + qw[4 + 3] = __uint_as_float(_v4_5_3); + } + { + 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*)(q_norm_weight + (d0 + 64) + (0))) : "memory"); + qw[8 + 0] = __uint_as_float(_v4_6_0); + qw[8 + 1] = __uint_as_float(_v4_6_1); + qw[8 + 2] = __uint_as_float(_v4_6_2); + qw[8 + 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*)(q_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + qw[12 + 0] = __uint_as_float(_v4_7_0); + qw[12 + 1] = __uint_as_float(_v4_7_1); + qw[12 + 2] = __uint_as_float(_v4_7_2); + qw[12 + 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 + d0 + (0))) : "memory"); + kw[0 + 0] = __uint_as_float(_v4_8_0); + kw[0 + 1] = __uint_as_float(_v4_8_1); + kw[0 + 2] = __uint_as_float(_v4_8_2); + kw[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 + (d0 + 4) + (0))) : "memory"); + kw[4 + 0] = __uint_as_float(_v4_9_0); + kw[4 + 1] = __uint_as_float(_v4_9_1); + kw[4 + 2] = __uint_as_float(_v4_9_2); + kw[4 + 3] = __uint_as_float(_v4_9_3); + } + { + 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*)(k_norm_weight + (d0 + 64) + (0))) : "memory"); + kw[8 + 0] = __uint_as_float(_v4_10_0); + kw[8 + 1] = __uint_as_float(_v4_10_1); + kw[8 + 2] = __uint_as_float(_v4_10_2); + kw[8 + 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*)(k_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + kw[12 + 0] = __uint_as_float(_v4_11_0); + kw[12 + 1] = __uint_as_float(_v4_11_1); + kw[12 + 2] = __uint_as_float(_v4_11_2); + kw[12 + 3] = __uint_as_float(_v4_11_3); + } + } + int req = -1; + int pos = -1; + int lb[8]; + int le[8]; + int ll[8]; + int bk = 0; + #pragma unroll + for (int k_2 = 0; k_2 < 8; k_2++) { + lb[k_2] = 0; + le[k_2] = 0; + ll[k_2] = 0; + bk = lane_0 + 32 * k_2; + if (bk < num_requests) { + lb[k_2] = q_indptr[bk]; + le[k_2] = q_indptr[bk + 1]; + ll[k_2] = seq_lens[bk]; + } + } + int hit = 0; + unsigned int hit_mask = 0; + int src_lane = 0; + #pragma unroll + for (int k_3 = 0; k_3 < 8; k_3++) { + if (req < 0 && 32 * k_3 < num_requests) { + bk = lane_0 + 32 * k_3; + hit = 0; + if (bk < num_requests) { + if (lb[k_3] <= row && row < le[k_3]) { + hit = 1; + } + } + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_0; + if (hit_mask != 0) { + int _ffs_0 = __ffs(hit_mask); + src_lane = _ffs_0 - 1; + int _shfl_0 = __shfl_sync(0xFFFFFFFF, bk, src_lane); + req = _shfl_0; + int _shfl_1 = __shfl_sync(0xFFFFFFFF, row + ll[k_3] - le[k_3], src_lane); + pos = _shfl_1; + } + } + } + if (req < 0) { + #pragma unroll 1 + for (int base = 256; base < num_requests; base += 32) { + int b = base + lane_0; + hit = 0; + int b_begin = 0; + int b_end = 0; + int b_len = 0; + if (b < num_requests) { + b_begin = q_indptr[b]; + b_end = q_indptr[b + 1]; + b_len = seq_lens[b]; + if (b_begin <= row && row < b_end) { + hit = 1; + } + } + unsigned int _vote_1 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_1; + if (hit_mask != 0) { + int _ffs_1 = __ffs(hit_mask); + src_lane = _ffs_1 - 1; + int _shfl_2 = __shfl_sync(0xFFFFFFFF, b, src_lane); + req = _shfl_2; + int _shfl_3 = __shfl_sync(0xFFFFFFFF, row + b_len - b_end, src_lane); + pos = _shfl_3; + break; + } + } + } + if (req >= 0 && pos >= 0) { + int page32 = page_table[req * max_pages + pos / page_size]; + long long page = (long long)page32; + int slot = pos % page_size; + long long k_row_base = page * k_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long v_row_base = page * v_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long cs_base = (long long)pos * 128; + float cosv[8]; + float sinv[8]; + { + 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 + (long long)d0) + (0))) : "memory"); + cosv[0 + 0] = __uint_as_float(_v4_12_0); + cosv[0 + 1] = __uint_as_float(_v4_12_1); + cosv[0 + 2] = __uint_as_float(_v4_12_2); + cosv[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 + (long long)d0 + 4) + (0))) : "memory"); + cosv[4 + 0] = __uint_as_float(_v4_13_0); + cosv[4 + 1] = __uint_as_float(_v4_13_1); + cosv[4 + 2] = __uint_as_float(_v4_13_2); + cosv[4 + 3] = __uint_as_float(_v4_13_3); + } + { + unsigned _v4_14_0; + unsigned _v4_14_1; + unsigned _v4_14_2; + unsigned _v4_14_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_14_0), "=r"(_v4_14_1), "=r"(_v4_14_2), "=r"(_v4_14_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64)) + (0))) : "memory"); + sinv[0 + 0] = __uint_as_float(_v4_14_0); + sinv[0 + 1] = __uint_as_float(_v4_14_1); + sinv[0 + 2] = __uint_as_float(_v4_14_2); + sinv[0 + 3] = __uint_as_float(_v4_14_3); + } + { + unsigned _v4_15_0; + unsigned _v4_15_1; + unsigned _v4_15_2; + unsigned _v4_15_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_15_0), "=r"(_v4_15_1), "=r"(_v4_15_2), "=r"(_v4_15_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64) + 4) + (0))) : "memory"); + sinv[4 + 0] = __uint_as_float(_v4_15_0); + sinv[4 + 1] = __uint_as_float(_v4_15_1); + sinv[4 + 2] = __uint_as_float(_v4_15_2); + sinv[4 + 3] = __uint_as_float(_v4_15_3); + } + float ss = 0.0f; + float inv = 0.0f; + float wj = 0.0f; + float x1 = 0.0f; + float x2 = 0.0f; + int head = 0; + int is_q = 0; + int kvh = 0; + long long q_dst = 0; + long long k_dst = 0; + #pragma unroll + for (int g_2 = 0; g_2 < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_2++) { + head = head_of[g_2]; + is_q = ((head < NUM_Q_HEADS) ? 1 : 0); + if (SKIP_MATH == 0 && qk_norm_policy == 2) { + ss = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ss += xs[g_2 * 2 * 8 + j_1] * xs[g_2 * 2 * 8 + j_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_0; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + wj = ((is_q != 0) ? qw[j_2] : kw[j_2]); + xs[g_2 * 2 * 8 + j_2] = xs[g_2 * 2 * 8 + j_2] * (inv * wj); + } + } + { + #pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x1 = xs[g_2 * 2 * 8 + j_3]; + x2 = xs[g_2 * 2 * 8 + 8 + j_3]; + float _fma_0 = __fmaf_rn(x1, cosv[j_3], -(x2 * sinv[j_3])); + xs[g_2 * 2 * 8 + j_3] = _fma_0; + float _fma_1 = __fmaf_rn(x2, cosv[j_3], x1 * sinv[j_3]); + xs[g_2 * 2 * 8 + 8 + j_3] = _fma_1; + } + } + if (SKIP_MATH == 0 && qk_norm_policy == 1) { + ss = 0.0f; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + ss += xs[g_2 * 2 * 8 + j_4] * xs[g_2 * 2 * 8 + j_4]; + } + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_1; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + wj = ((is_q != 0) ? qw[j_5] : kw[j_5]); + xs[g_2 * 2 * 8 + j_5] = xs[g_2 * 2 * 8 + j_5] * (inv * wj); + } + } + if (head < NUM_Q_HEADS) { + q_dst = (long long)row * (long long)(NUM_Q_HEADS * 128) + (long long)(head * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + q_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + (q_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else if (head < NUM_Q_HEADS + NUM_KV_HEADS) { + kvh = head - NUM_Q_HEADS; + if (use_out_k != 0) { + k_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + k_dst = k_row_base + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + int vh = 0; + long long v_dst = 0; + #pragma unroll + for (int g_3 = 0; g_3 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_3++) { + vh = vhead_of[g_3]; + if (vh < NUM_KV_HEADS) { + if (use_out_v != 0) { + v_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(out_v + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(out_v + (v_dst + 64) + 0) = _iv4; + } + } else { + v_dst = v_row_base + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(v_cache + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(v_cache + (v_dst + 64) + 0) = _iv4; + } + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_binding.cu b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_binding.cu new file mode 100644 index 00000000000..c5cdf9e6a9c --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_binding.cu @@ -0,0 +1,160 @@ +/* + * 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_append_51a814b69e055ec5993b'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include + +extern "C" __global__ void kernel_cake_fused_qk_rope_append_51a814b69e055ec5993b(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page); + + +namespace cake_host_shim_a2d5ce401c9cfeb0 { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_q_indptr, TensorView arg_seq_lens, TensorView arg_page_table, TensorView arg_k_cache, TensorView arg_v_cache, TensorView arg_out_q, TensorView arg_out_k, TensorView arg_out_v, TensorView arg_q_norm_weight, TensorView arg_k_norm_weight, int64_t arg_num_rows, int64_t arg_num_requests, int64_t arg_max_pages, int64_t arg_page_size, int64_t arg_k_page_stride, int64_t arg_v_page_stride, int64_t arg_num_row_ctas, int64_t arg_clear_units_per_request, int64_t arg_qk_norm_policy, double arg_eps, int64_t arg_use_out_k, int64_t arg_use_out_v, int64_t arg_clear_last_page, 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_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_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_page_table, "page_table"); + check_dtype(arg_page_table, DLDataType{kDLInt, 32, 1}, "page_table"); + check_contiguous(arg_page_table, "page_table"); + check_cuda_tensor(arg_k_cache, "k_cache"); + check_dtype(arg_k_cache, DLDataType{kDLBfloat, 16, 1}, "k_cache"); + check_contiguous(arg_k_cache, "k_cache"); + check_cuda_tensor(arg_v_cache, "v_cache"); + check_dtype(arg_v_cache, DLDataType{kDLBfloat, 16, 1}, "v_cache"); + check_contiguous(arg_v_cache, "v_cache"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLBfloat, 16, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLBfloat, 16, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLBfloat, 16, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + 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"); + 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_max_pages >= -2147483648LL && arg_max_pages <= 2147483647LL, ValueError) + << "scalar 'max_pages' value " << arg_max_pages + << " 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_num_row_ctas >= -2147483648LL && arg_num_row_ctas <= 2147483647LL, ValueError) + << "scalar 'num_row_ctas' value " << arg_num_row_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_units_per_request >= -2147483648LL && arg_clear_units_per_request <= 2147483647LL, ValueError) + << "scalar 'clear_units_per_request' value " << arg_clear_units_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_norm_policy >= -2147483648LL && arg_qk_norm_policy <= 2147483647LL, ValueError) + << "scalar 'qk_norm_policy' value " << arg_qk_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_k >= -2147483648LL && arg_use_out_k <= 2147483647LL, ValueError) + << "scalar 'use_out_k' value " << arg_use_out_k + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_v >= -2147483648LL && arg_use_out_v <= 2147483647LL, ValueError) + << "scalar 'use_out_v' value " << arg_use_out_v + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_last_page >= -2147483648LL && arg_clear_last_page <= 2147483647LL, ValueError) + << "scalar 'clear_last_page' value " << arg_clear_last_page + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_page_table, arg_qkv, "page_table", "qkv"); + check_same_device(arg_k_cache, arg_qkv, "k_cache", "qkv"); + check_same_device(arg_v_cache, arg_qkv, "v_cache", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "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_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "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_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_page_table = static_cast(arg_page_table.data_ptr()); + __nv_bfloat16* p_k_cache = static_cast<__nv_bfloat16*>(arg_k_cache.data_ptr()); + __nv_bfloat16* p_v_cache = static_cast<__nv_bfloat16*>(arg_v_cache.data_ptr()); + __nv_bfloat16* p_out_q = static_cast<__nv_bfloat16*>(arg_out_q.data_ptr()); + __nv_bfloat16* p_out_k = static_cast<__nv_bfloat16*>(arg_out_k.data_ptr()); + __nv_bfloat16* p_out_v = static_cast<__nv_bfloat16*>(arg_out_v.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()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_max_pages = (int32_t)arg_max_pages; + int32_t v_page_size = (int32_t)arg_page_size; + int64_t v_k_page_stride = (int64_t)arg_k_page_stride; + int64_t v_v_page_stride = (int64_t)arg_v_page_stride; + int32_t v_num_row_ctas = (int32_t)arg_num_row_ctas; + int32_t v_clear_units_per_request = (int32_t)arg_clear_units_per_request; + int32_t v_qk_norm_policy = (int32_t)arg_qk_norm_policy; + float v_eps = (float)arg_eps; + int32_t v_use_out_k = (int32_t)arg_use_out_k; + int32_t v_use_out_v = (int32_t)arg_use_out_v; + int32_t v_clear_last_page = (int32_t)arg_clear_last_page; + void* kargs[] = {&p_qkv, &p_cos_sin, &p_q_indptr, &p_seq_lens, &p_page_table, &p_k_cache, &p_v_cache, &p_out_q, &p_out_k, &p_out_v, &p_q_norm_weight, &p_k_norm_weight, &v_num_rows, &v_num_requests, &v_max_pages, &v_page_size, &v_k_page_stride, &v_v_page_stride, &v_num_row_ctas, &v_clear_units_per_request, &v_qk_norm_policy, &v_eps, &v_use_out_k, &v_use_out_v, &v_clear_last_page}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_append_51a814b69e055ec5993b), grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_append_51a814b69e055ec5993b failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_a2d5ce401c9cfeb0 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_a2d5ce401c9cfeb0::Run); diff --git a/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_kernel.cu b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_kernel.cu new file mode 100644 index 00000000000..47729fa31cb --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_100a/cake_fused_qk_rope_append_51a814b69e055ec5993b_kernel.cu @@ -0,0 +1,608 @@ +/* + * 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 THREADS 128 +#define NUM_Q_HEADS 8 +#define NUM_KV_HEADS 1 +#define WARPS_PER_ROW 4 +#define SKIP_MATH 0 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_fused_qk_rope_append_51a814b69e055ec5993b(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // === Task calls (dependency order) === + int lane_0 = lane; + int warp_1 = warp; + int hsel = lane_0 / 8; + int d0 = lane_0 % 8 * 8; + if (bid >= num_row_ctas) { + int cid = (bid - num_row_ctas) * 4 + warp_1; + if (cid < num_requests * clear_units_per_request && clear_last_page != 0) { + int req_c = cid / clear_units_per_request; + int u = cid % clear_units_per_request; + int cache_sel = u % 2; + int unit_begin = u / 2 * 256; + int seq_len_c = seq_lens[req_c]; + if (seq_len_c > 0) { + int last_c = seq_len_c - 1; + int start_chunk = (last_c % page_size + 1) * (NUM_KV_HEADS * 128 / 8); + int end_chunk = page_size * (NUM_KV_HEADS * 128 / 8); + if (start_chunk < unit_begin + 256 && unit_begin < end_chunk) { + int page_c32 = page_table[req_c * max_pages + last_c / page_size]; + long long page_c = (long long)page_c32; + float zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0.0f; + } + int c = 0; + long long off_c = 0; + if (cache_sel == 0) { + long long base_k = page_c * k_page_stride; + #pragma unroll + for (int k = 0; k < 8; k++) { + c = unit_begin + lane_0 + 32 * k; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (base_k + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else { + long long base_v = page_c * v_page_stride; + #pragma unroll + for (int k_1 = 0; k_1 < 8; k_1++) { + c = unit_begin + lane_0 + 32 * k_1; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(v_cache + (base_v + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } else { + int row = bid / ((WARPS_PER_ROW + 3) / 4) * ((4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW) + warp_1 / WARPS_PER_ROW; + int sub = bid % ((WARPS_PER_ROW + 3) / 4) * 4 + warp_1 % WARPS_PER_ROW; + if (row < num_rows) { + long long row_base = (long long)row * (long long)((NUM_Q_HEADS + 2 * NUM_KV_HEADS) * 128); + float xs[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 2 * 8]; + int head_of[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_h = 0; + #pragma unroll + for (int g = 0; g < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g++) { + head_of[g] = 4 * (sub + WARPS_PER_ROW * g) + hsel; + src_h = ((head_of[g] < NUM_Q_HEADS + NUM_KV_HEADS) ? head_of[g] : NUM_Q_HEADS + NUM_KV_HEADS - 1); + { + const uint4* _vptr_0 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0)) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_0[_blk].x), "=r"(_vld_0[_blk].y), "=r"(_vld_0[_blk].z), "=r"(_vld_0[_blk].w) : "l"((const void*)(_vptr_0 + _blk)) : "memory"); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #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"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0 + 64)) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_1[_blk].x), "=r"(_vld_1[_blk].y), "=r"(_vld_1[_blk].z), "=r"(_vld_1[_blk].w) : "l"((const void*)(_vptr_1 + _blk)) : "memory"); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #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"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + int vbits[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 8]; + int vhead_of[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_v = 0; + #pragma unroll + for (int g_1 = 0; g_1 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_1++) { + vhead_of[g_1] = 4 * (sub + WARPS_PER_ROW * g_1) + hsel; + src_v = ((vhead_of[g_1] < NUM_KV_HEADS) ? vhead_of[g_1] : NUM_KV_HEADS - 1); + { + const int4* _ivptr_2 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0)) + 0); + int4 _ivld_2; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_2.x), "=r"(_ivld_2.y), "=r"(_ivld_2.z), "=r"(_ivld_2.w) : "l"((const void*)(_ivptr_2)) : "memory"); + vbits[g_1 * 8 + 0] = _ivld_2.x; + vbits[g_1 * 8 + 1] = _ivld_2.y; + vbits[g_1 * 8 + 2] = _ivld_2.z; + vbits[g_1 * 8 + 3] = _ivld_2.w; + } + { + const int4* _ivptr_3 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0 + 64)) + 0); + int4 _ivld_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_3.x), "=r"(_ivld_3.y), "=r"(_ivld_3.z), "=r"(_ivld_3.w) : "l"((const void*)(_ivptr_3)) : "memory"); + vbits[g_1 * 8 + 4 + 0] = _ivld_3.x; + vbits[g_1 * 8 + 4 + 1] = _ivld_3.y; + vbits[g_1 * 8 + 4 + 2] = _ivld_3.z; + vbits[g_1 * 8 + 4 + 3] = _ivld_3.w; + } + } + float qw[16]; + float kw[16]; + if (qk_norm_policy != 0) { + { + 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 + d0 + (0))) : "memory"); + qw[0 + 0] = __uint_as_float(_v4_4_0); + qw[0 + 1] = __uint_as_float(_v4_4_1); + qw[0 + 2] = __uint_as_float(_v4_4_2); + qw[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 + (d0 + 4) + (0))) : "memory"); + qw[4 + 0] = __uint_as_float(_v4_5_0); + qw[4 + 1] = __uint_as_float(_v4_5_1); + qw[4 + 2] = __uint_as_float(_v4_5_2); + qw[4 + 3] = __uint_as_float(_v4_5_3); + } + { + 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*)(q_norm_weight + (d0 + 64) + (0))) : "memory"); + qw[8 + 0] = __uint_as_float(_v4_6_0); + qw[8 + 1] = __uint_as_float(_v4_6_1); + qw[8 + 2] = __uint_as_float(_v4_6_2); + qw[8 + 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*)(q_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + qw[12 + 0] = __uint_as_float(_v4_7_0); + qw[12 + 1] = __uint_as_float(_v4_7_1); + qw[12 + 2] = __uint_as_float(_v4_7_2); + qw[12 + 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 + d0 + (0))) : "memory"); + kw[0 + 0] = __uint_as_float(_v4_8_0); + kw[0 + 1] = __uint_as_float(_v4_8_1); + kw[0 + 2] = __uint_as_float(_v4_8_2); + kw[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 + (d0 + 4) + (0))) : "memory"); + kw[4 + 0] = __uint_as_float(_v4_9_0); + kw[4 + 1] = __uint_as_float(_v4_9_1); + kw[4 + 2] = __uint_as_float(_v4_9_2); + kw[4 + 3] = __uint_as_float(_v4_9_3); + } + { + 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*)(k_norm_weight + (d0 + 64) + (0))) : "memory"); + kw[8 + 0] = __uint_as_float(_v4_10_0); + kw[8 + 1] = __uint_as_float(_v4_10_1); + kw[8 + 2] = __uint_as_float(_v4_10_2); + kw[8 + 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*)(k_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + kw[12 + 0] = __uint_as_float(_v4_11_0); + kw[12 + 1] = __uint_as_float(_v4_11_1); + kw[12 + 2] = __uint_as_float(_v4_11_2); + kw[12 + 3] = __uint_as_float(_v4_11_3); + } + } + int req = -1; + int pos = -1; + int lb[8]; + int le[8]; + int ll[8]; + int bk = 0; + #pragma unroll + for (int k_2 = 0; k_2 < 8; k_2++) { + lb[k_2] = 0; + le[k_2] = 0; + ll[k_2] = 0; + bk = lane_0 + 32 * k_2; + if (bk < num_requests) { + lb[k_2] = q_indptr[bk]; + le[k_2] = q_indptr[bk + 1]; + ll[k_2] = seq_lens[bk]; + } + } + int hit = 0; + unsigned int hit_mask = 0; + int src_lane = 0; + #pragma unroll + for (int k_3 = 0; k_3 < 8; k_3++) { + if (req < 0 && 32 * k_3 < num_requests) { + bk = lane_0 + 32 * k_3; + hit = 0; + if (bk < num_requests) { + if (lb[k_3] <= row && row < le[k_3]) { + hit = 1; + } + } + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_0; + if (hit_mask != 0) { + int _ffs_0 = __ffs(hit_mask); + src_lane = _ffs_0 - 1; + int _shfl_0 = __shfl_sync(0xFFFFFFFF, bk, src_lane); + req = _shfl_0; + int _shfl_1 = __shfl_sync(0xFFFFFFFF, row + ll[k_3] - le[k_3], src_lane); + pos = _shfl_1; + } + } + } + if (req < 0) { + #pragma unroll 1 + for (int base = 256; base < num_requests; base += 32) { + int b = base + lane_0; + hit = 0; + int b_begin = 0; + int b_end = 0; + int b_len = 0; + if (b < num_requests) { + b_begin = q_indptr[b]; + b_end = q_indptr[b + 1]; + b_len = seq_lens[b]; + if (b_begin <= row && row < b_end) { + hit = 1; + } + } + unsigned int _vote_1 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_1; + if (hit_mask != 0) { + int _ffs_1 = __ffs(hit_mask); + src_lane = _ffs_1 - 1; + int _shfl_2 = __shfl_sync(0xFFFFFFFF, b, src_lane); + req = _shfl_2; + int _shfl_3 = __shfl_sync(0xFFFFFFFF, row + b_len - b_end, src_lane); + pos = _shfl_3; + break; + } + } + } + if (req >= 0 && pos >= 0) { + int page32 = page_table[req * max_pages + pos / page_size]; + long long page = (long long)page32; + int slot = pos % page_size; + long long k_row_base = page * k_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long v_row_base = page * v_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long cs_base = (long long)pos * 128; + float cosv[8]; + float sinv[8]; + { + 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 + (long long)d0) + (0))) : "memory"); + cosv[0 + 0] = __uint_as_float(_v4_12_0); + cosv[0 + 1] = __uint_as_float(_v4_12_1); + cosv[0 + 2] = __uint_as_float(_v4_12_2); + cosv[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 + (long long)d0 + 4) + (0))) : "memory"); + cosv[4 + 0] = __uint_as_float(_v4_13_0); + cosv[4 + 1] = __uint_as_float(_v4_13_1); + cosv[4 + 2] = __uint_as_float(_v4_13_2); + cosv[4 + 3] = __uint_as_float(_v4_13_3); + } + { + unsigned _v4_14_0; + unsigned _v4_14_1; + unsigned _v4_14_2; + unsigned _v4_14_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_14_0), "=r"(_v4_14_1), "=r"(_v4_14_2), "=r"(_v4_14_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64)) + (0))) : "memory"); + sinv[0 + 0] = __uint_as_float(_v4_14_0); + sinv[0 + 1] = __uint_as_float(_v4_14_1); + sinv[0 + 2] = __uint_as_float(_v4_14_2); + sinv[0 + 3] = __uint_as_float(_v4_14_3); + } + { + unsigned _v4_15_0; + unsigned _v4_15_1; + unsigned _v4_15_2; + unsigned _v4_15_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_15_0), "=r"(_v4_15_1), "=r"(_v4_15_2), "=r"(_v4_15_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64) + 4) + (0))) : "memory"); + sinv[4 + 0] = __uint_as_float(_v4_15_0); + sinv[4 + 1] = __uint_as_float(_v4_15_1); + sinv[4 + 2] = __uint_as_float(_v4_15_2); + sinv[4 + 3] = __uint_as_float(_v4_15_3); + } + float ss = 0.0f; + float inv = 0.0f; + float wj = 0.0f; + float x1 = 0.0f; + float x2 = 0.0f; + int head = 0; + int is_q = 0; + int kvh = 0; + long long q_dst = 0; + long long k_dst = 0; + #pragma unroll + for (int g_2 = 0; g_2 < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_2++) { + head = head_of[g_2]; + is_q = ((head < NUM_Q_HEADS) ? 1 : 0); + if (SKIP_MATH == 0 && qk_norm_policy == 2) { + ss = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ss += xs[g_2 * 2 * 8 + j_1] * xs[g_2 * 2 * 8 + j_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_0; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + wj = ((is_q != 0) ? qw[j_2] : kw[j_2]); + xs[g_2 * 2 * 8 + j_2] = xs[g_2 * 2 * 8 + j_2] * (inv * wj); + } + } + { + #pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x1 = xs[g_2 * 2 * 8 + j_3]; + x2 = xs[g_2 * 2 * 8 + 8 + j_3]; + float _fma_0 = __fmaf_rn(x1, cosv[j_3], -(x2 * sinv[j_3])); + xs[g_2 * 2 * 8 + j_3] = _fma_0; + float _fma_1 = __fmaf_rn(x2, cosv[j_3], x1 * sinv[j_3]); + xs[g_2 * 2 * 8 + 8 + j_3] = _fma_1; + } + } + if (SKIP_MATH == 0 && qk_norm_policy == 1) { + ss = 0.0f; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + ss += xs[g_2 * 2 * 8 + j_4] * xs[g_2 * 2 * 8 + j_4]; + } + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_1; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + wj = ((is_q != 0) ? qw[j_5] : kw[j_5]); + xs[g_2 * 2 * 8 + j_5] = xs[g_2 * 2 * 8 + j_5] * (inv * wj); + } + } + if (head < NUM_Q_HEADS) { + q_dst = (long long)row * (long long)(NUM_Q_HEADS * 128) + (long long)(head * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + q_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + (q_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else if (head < NUM_Q_HEADS + NUM_KV_HEADS) { + kvh = head - NUM_Q_HEADS; + if (use_out_k != 0) { + k_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + k_dst = k_row_base + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + int vh = 0; + long long v_dst = 0; + #pragma unroll + for (int g_3 = 0; g_3 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_3++) { + vh = vhead_of[g_3]; + if (vh < NUM_KV_HEADS) { + if (use_out_v != 0) { + v_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(out_v + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(out_v + (v_dst + 64) + 0) = _iv4; + } + } else { + v_dst = v_row_base + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(v_cache + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(v_cache + (v_dst + 64) + 0) = _iv4; + } + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_binding.cu b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_binding.cu new file mode 100644 index 00000000000..188fef2764f --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_binding.cu @@ -0,0 +1,160 @@ +/* + * 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_append_786a6a33470bb271f31a'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include + +extern "C" __global__ void kernel_cake_fused_qk_rope_append_786a6a33470bb271f31a(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page); + + +namespace cake_host_shim_7706ef7126a3238d { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_q_indptr, TensorView arg_seq_lens, TensorView arg_page_table, TensorView arg_k_cache, TensorView arg_v_cache, TensorView arg_out_q, TensorView arg_out_k, TensorView arg_out_v, TensorView arg_q_norm_weight, TensorView arg_k_norm_weight, int64_t arg_num_rows, int64_t arg_num_requests, int64_t arg_max_pages, int64_t arg_page_size, int64_t arg_k_page_stride, int64_t arg_v_page_stride, int64_t arg_num_row_ctas, int64_t arg_clear_units_per_request, int64_t arg_qk_norm_policy, double arg_eps, int64_t arg_use_out_k, int64_t arg_use_out_v, int64_t arg_clear_last_page, 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_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_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_page_table, "page_table"); + check_dtype(arg_page_table, DLDataType{kDLInt, 32, 1}, "page_table"); + check_contiguous(arg_page_table, "page_table"); + check_cuda_tensor(arg_k_cache, "k_cache"); + check_dtype(arg_k_cache, DLDataType{kDLBfloat, 16, 1}, "k_cache"); + check_contiguous(arg_k_cache, "k_cache"); + check_cuda_tensor(arg_v_cache, "v_cache"); + check_dtype(arg_v_cache, DLDataType{kDLBfloat, 16, 1}, "v_cache"); + check_contiguous(arg_v_cache, "v_cache"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLBfloat, 16, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLBfloat, 16, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLBfloat, 16, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + 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"); + 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_max_pages >= -2147483648LL && arg_max_pages <= 2147483647LL, ValueError) + << "scalar 'max_pages' value " << arg_max_pages + << " 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_num_row_ctas >= -2147483648LL && arg_num_row_ctas <= 2147483647LL, ValueError) + << "scalar 'num_row_ctas' value " << arg_num_row_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_units_per_request >= -2147483648LL && arg_clear_units_per_request <= 2147483647LL, ValueError) + << "scalar 'clear_units_per_request' value " << arg_clear_units_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_norm_policy >= -2147483648LL && arg_qk_norm_policy <= 2147483647LL, ValueError) + << "scalar 'qk_norm_policy' value " << arg_qk_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_k >= -2147483648LL && arg_use_out_k <= 2147483647LL, ValueError) + << "scalar 'use_out_k' value " << arg_use_out_k + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_v >= -2147483648LL && arg_use_out_v <= 2147483647LL, ValueError) + << "scalar 'use_out_v' value " << arg_use_out_v + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_last_page >= -2147483648LL && arg_clear_last_page <= 2147483647LL, ValueError) + << "scalar 'clear_last_page' value " << arg_clear_last_page + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_page_table, arg_qkv, "page_table", "qkv"); + check_same_device(arg_k_cache, arg_qkv, "k_cache", "qkv"); + check_same_device(arg_v_cache, arg_qkv, "v_cache", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "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_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "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_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_page_table = static_cast(arg_page_table.data_ptr()); + __nv_bfloat16* p_k_cache = static_cast<__nv_bfloat16*>(arg_k_cache.data_ptr()); + __nv_bfloat16* p_v_cache = static_cast<__nv_bfloat16*>(arg_v_cache.data_ptr()); + __nv_bfloat16* p_out_q = static_cast<__nv_bfloat16*>(arg_out_q.data_ptr()); + __nv_bfloat16* p_out_k = static_cast<__nv_bfloat16*>(arg_out_k.data_ptr()); + __nv_bfloat16* p_out_v = static_cast<__nv_bfloat16*>(arg_out_v.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()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_max_pages = (int32_t)arg_max_pages; + int32_t v_page_size = (int32_t)arg_page_size; + int64_t v_k_page_stride = (int64_t)arg_k_page_stride; + int64_t v_v_page_stride = (int64_t)arg_v_page_stride; + int32_t v_num_row_ctas = (int32_t)arg_num_row_ctas; + int32_t v_clear_units_per_request = (int32_t)arg_clear_units_per_request; + int32_t v_qk_norm_policy = (int32_t)arg_qk_norm_policy; + float v_eps = (float)arg_eps; + int32_t v_use_out_k = (int32_t)arg_use_out_k; + int32_t v_use_out_v = (int32_t)arg_use_out_v; + int32_t v_clear_last_page = (int32_t)arg_clear_last_page; + void* kargs[] = {&p_qkv, &p_cos_sin, &p_q_indptr, &p_seq_lens, &p_page_table, &p_k_cache, &p_v_cache, &p_out_q, &p_out_k, &p_out_v, &p_q_norm_weight, &p_k_norm_weight, &v_num_rows, &v_num_requests, &v_max_pages, &v_page_size, &v_k_page_stride, &v_v_page_stride, &v_num_row_ctas, &v_clear_units_per_request, &v_qk_norm_policy, &v_eps, &v_use_out_k, &v_use_out_v, &v_clear_last_page}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_append_786a6a33470bb271f31a), grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_append_786a6a33470bb271f31a failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_7706ef7126a3238d + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_7706ef7126a3238d::Run); diff --git a/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_kernel.cu b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_kernel.cu new file mode 100644 index 00000000000..51ec39da81e --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_786a6a33470bb271f31a_kernel.cu @@ -0,0 +1,608 @@ +/* + * 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 THREADS 128 +#define NUM_Q_HEADS 8 +#define NUM_KV_HEADS 1 +#define WARPS_PER_ROW 4 +#define SKIP_MATH 0 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_fused_qk_rope_append_786a6a33470bb271f31a(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // === Task calls (dependency order) === + int lane_0 = lane; + int warp_1 = warp; + int hsel = lane_0 / 8; + int d0 = lane_0 % 8 * 8; + if (bid >= num_row_ctas) { + int cid = (bid - num_row_ctas) * 4 + warp_1; + if (cid < num_requests * clear_units_per_request && clear_last_page != 0) { + int req_c = cid / clear_units_per_request; + int u = cid % clear_units_per_request; + int cache_sel = u % 2; + int unit_begin = u / 2 * 256; + int seq_len_c = seq_lens[req_c]; + if (seq_len_c > 0) { + int last_c = seq_len_c - 1; + int start_chunk = (last_c % page_size + 1) * (NUM_KV_HEADS * 128 / 8); + int end_chunk = page_size * (NUM_KV_HEADS * 128 / 8); + if (start_chunk < unit_begin + 256 && unit_begin < end_chunk) { + int page_c32 = page_table[req_c * max_pages + last_c / page_size]; + long long page_c = (long long)page_c32; + float zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0.0f; + } + int c = 0; + long long off_c = 0; + if (cache_sel == 0) { + long long base_k = page_c * k_page_stride; + #pragma unroll + for (int k = 0; k < 8; k++) { + c = unit_begin + lane_0 + 32 * k; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (base_k + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else { + long long base_v = page_c * v_page_stride; + #pragma unroll + for (int k_1 = 0; k_1 < 8; k_1++) { + c = unit_begin + lane_0 + 32 * k_1; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(v_cache + (base_v + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } else { + int row = bid / ((WARPS_PER_ROW + 3) / 4) * ((4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW) + warp_1 / WARPS_PER_ROW; + int sub = bid % ((WARPS_PER_ROW + 3) / 4) * 4 + warp_1 % WARPS_PER_ROW; + if (row < num_rows) { + long long row_base = (long long)row * (long long)((NUM_Q_HEADS + 2 * NUM_KV_HEADS) * 128); + float xs[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 2 * 8]; + int head_of[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_h = 0; + #pragma unroll + for (int g = 0; g < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g++) { + head_of[g] = 4 * (sub + WARPS_PER_ROW * g) + hsel; + src_h = ((head_of[g] < NUM_Q_HEADS + NUM_KV_HEADS) ? head_of[g] : NUM_Q_HEADS + NUM_KV_HEADS - 1); + { + const uint4* _vptr_0 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0)) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_0[_blk].x), "=r"(_vld_0[_blk].y), "=r"(_vld_0[_blk].z), "=r"(_vld_0[_blk].w) : "l"((const void*)(_vptr_0 + _blk)) : "memory"); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #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"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0 + 64)) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_1[_blk].x), "=r"(_vld_1[_blk].y), "=r"(_vld_1[_blk].z), "=r"(_vld_1[_blk].w) : "l"((const void*)(_vptr_1 + _blk)) : "memory"); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #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"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + int vbits[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 8]; + int vhead_of[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_v = 0; + #pragma unroll + for (int g_1 = 0; g_1 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_1++) { + vhead_of[g_1] = 4 * (sub + WARPS_PER_ROW * g_1) + hsel; + src_v = ((vhead_of[g_1] < NUM_KV_HEADS) ? vhead_of[g_1] : NUM_KV_HEADS - 1); + { + const int4* _ivptr_2 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0)) + 0); + int4 _ivld_2; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_2.x), "=r"(_ivld_2.y), "=r"(_ivld_2.z), "=r"(_ivld_2.w) : "l"((const void*)(_ivptr_2)) : "memory"); + vbits[g_1 * 8 + 0] = _ivld_2.x; + vbits[g_1 * 8 + 1] = _ivld_2.y; + vbits[g_1 * 8 + 2] = _ivld_2.z; + vbits[g_1 * 8 + 3] = _ivld_2.w; + } + { + const int4* _ivptr_3 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0 + 64)) + 0); + int4 _ivld_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_3.x), "=r"(_ivld_3.y), "=r"(_ivld_3.z), "=r"(_ivld_3.w) : "l"((const void*)(_ivptr_3)) : "memory"); + vbits[g_1 * 8 + 4 + 0] = _ivld_3.x; + vbits[g_1 * 8 + 4 + 1] = _ivld_3.y; + vbits[g_1 * 8 + 4 + 2] = _ivld_3.z; + vbits[g_1 * 8 + 4 + 3] = _ivld_3.w; + } + } + float qw[16]; + float kw[16]; + if (qk_norm_policy != 0) { + { + 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 + d0 + (0))) : "memory"); + qw[0 + 0] = __uint_as_float(_v4_4_0); + qw[0 + 1] = __uint_as_float(_v4_4_1); + qw[0 + 2] = __uint_as_float(_v4_4_2); + qw[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 + (d0 + 4) + (0))) : "memory"); + qw[4 + 0] = __uint_as_float(_v4_5_0); + qw[4 + 1] = __uint_as_float(_v4_5_1); + qw[4 + 2] = __uint_as_float(_v4_5_2); + qw[4 + 3] = __uint_as_float(_v4_5_3); + } + { + 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*)(q_norm_weight + (d0 + 64) + (0))) : "memory"); + qw[8 + 0] = __uint_as_float(_v4_6_0); + qw[8 + 1] = __uint_as_float(_v4_6_1); + qw[8 + 2] = __uint_as_float(_v4_6_2); + qw[8 + 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*)(q_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + qw[12 + 0] = __uint_as_float(_v4_7_0); + qw[12 + 1] = __uint_as_float(_v4_7_1); + qw[12 + 2] = __uint_as_float(_v4_7_2); + qw[12 + 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 + d0 + (0))) : "memory"); + kw[0 + 0] = __uint_as_float(_v4_8_0); + kw[0 + 1] = __uint_as_float(_v4_8_1); + kw[0 + 2] = __uint_as_float(_v4_8_2); + kw[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 + (d0 + 4) + (0))) : "memory"); + kw[4 + 0] = __uint_as_float(_v4_9_0); + kw[4 + 1] = __uint_as_float(_v4_9_1); + kw[4 + 2] = __uint_as_float(_v4_9_2); + kw[4 + 3] = __uint_as_float(_v4_9_3); + } + { + 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*)(k_norm_weight + (d0 + 64) + (0))) : "memory"); + kw[8 + 0] = __uint_as_float(_v4_10_0); + kw[8 + 1] = __uint_as_float(_v4_10_1); + kw[8 + 2] = __uint_as_float(_v4_10_2); + kw[8 + 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*)(k_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + kw[12 + 0] = __uint_as_float(_v4_11_0); + kw[12 + 1] = __uint_as_float(_v4_11_1); + kw[12 + 2] = __uint_as_float(_v4_11_2); + kw[12 + 3] = __uint_as_float(_v4_11_3); + } + } + int req = -1; + int pos = -1; + int lb[8]; + int le[8]; + int ll[8]; + int bk = 0; + #pragma unroll + for (int k_2 = 0; k_2 < 8; k_2++) { + lb[k_2] = 0; + le[k_2] = 0; + ll[k_2] = 0; + bk = lane_0 + 32 * k_2; + if (bk < num_requests) { + lb[k_2] = q_indptr[bk]; + le[k_2] = q_indptr[bk + 1]; + ll[k_2] = seq_lens[bk]; + } + } + int hit = 0; + unsigned int hit_mask = 0; + int src_lane = 0; + #pragma unroll + for (int k_3 = 0; k_3 < 8; k_3++) { + if (req < 0 && 32 * k_3 < num_requests) { + bk = lane_0 + 32 * k_3; + hit = 0; + if (bk < num_requests) { + if (lb[k_3] <= row && row < le[k_3]) { + hit = 1; + } + } + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_0; + if (hit_mask != 0) { + int _ffs_0 = __ffs(hit_mask); + src_lane = _ffs_0 - 1; + int _shfl_0 = __shfl_sync(0xFFFFFFFF, bk, src_lane); + req = _shfl_0; + int _shfl_1 = __shfl_sync(0xFFFFFFFF, row + ll[k_3] - le[k_3], src_lane); + pos = _shfl_1; + } + } + } + if (req < 0) { + #pragma unroll 1 + for (int base = 256; base < num_requests; base += 32) { + int b = base + lane_0; + hit = 0; + int b_begin = 0; + int b_end = 0; + int b_len = 0; + if (b < num_requests) { + b_begin = q_indptr[b]; + b_end = q_indptr[b + 1]; + b_len = seq_lens[b]; + if (b_begin <= row && row < b_end) { + hit = 1; + } + } + unsigned int _vote_1 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_1; + if (hit_mask != 0) { + int _ffs_1 = __ffs(hit_mask); + src_lane = _ffs_1 - 1; + int _shfl_2 = __shfl_sync(0xFFFFFFFF, b, src_lane); + req = _shfl_2; + int _shfl_3 = __shfl_sync(0xFFFFFFFF, row + b_len - b_end, src_lane); + pos = _shfl_3; + break; + } + } + } + if (req >= 0 && pos >= 0) { + int page32 = page_table[req * max_pages + pos / page_size]; + long long page = (long long)page32; + int slot = pos % page_size; + long long k_row_base = page * k_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long v_row_base = page * v_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long cs_base = (long long)pos * 128; + float cosv[8]; + float sinv[8]; + { + 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 + (long long)d0) + (0))) : "memory"); + cosv[0 + 0] = __uint_as_float(_v4_12_0); + cosv[0 + 1] = __uint_as_float(_v4_12_1); + cosv[0 + 2] = __uint_as_float(_v4_12_2); + cosv[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 + (long long)d0 + 4) + (0))) : "memory"); + cosv[4 + 0] = __uint_as_float(_v4_13_0); + cosv[4 + 1] = __uint_as_float(_v4_13_1); + cosv[4 + 2] = __uint_as_float(_v4_13_2); + cosv[4 + 3] = __uint_as_float(_v4_13_3); + } + { + unsigned _v4_14_0; + unsigned _v4_14_1; + unsigned _v4_14_2; + unsigned _v4_14_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_14_0), "=r"(_v4_14_1), "=r"(_v4_14_2), "=r"(_v4_14_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64)) + (0))) : "memory"); + sinv[0 + 0] = __uint_as_float(_v4_14_0); + sinv[0 + 1] = __uint_as_float(_v4_14_1); + sinv[0 + 2] = __uint_as_float(_v4_14_2); + sinv[0 + 3] = __uint_as_float(_v4_14_3); + } + { + unsigned _v4_15_0; + unsigned _v4_15_1; + unsigned _v4_15_2; + unsigned _v4_15_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_15_0), "=r"(_v4_15_1), "=r"(_v4_15_2), "=r"(_v4_15_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64) + 4) + (0))) : "memory"); + sinv[4 + 0] = __uint_as_float(_v4_15_0); + sinv[4 + 1] = __uint_as_float(_v4_15_1); + sinv[4 + 2] = __uint_as_float(_v4_15_2); + sinv[4 + 3] = __uint_as_float(_v4_15_3); + } + float ss = 0.0f; + float inv = 0.0f; + float wj = 0.0f; + float x1 = 0.0f; + float x2 = 0.0f; + int head = 0; + int is_q = 0; + int kvh = 0; + long long q_dst = 0; + long long k_dst = 0; + #pragma unroll + for (int g_2 = 0; g_2 < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_2++) { + head = head_of[g_2]; + is_q = ((head < NUM_Q_HEADS) ? 1 : 0); + if (SKIP_MATH == 0 && qk_norm_policy == 2) { + ss = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ss += xs[g_2 * 2 * 8 + j_1] * xs[g_2 * 2 * 8 + j_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_0; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + wj = ((is_q != 0) ? qw[j_2] : kw[j_2]); + xs[g_2 * 2 * 8 + j_2] = xs[g_2 * 2 * 8 + j_2] * (inv * wj); + } + } + { + #pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x1 = xs[g_2 * 2 * 8 + j_3]; + x2 = xs[g_2 * 2 * 8 + 8 + j_3]; + float _fma_0 = __fmaf_rn(x1, cosv[j_3], -(x2 * sinv[j_3])); + xs[g_2 * 2 * 8 + j_3] = _fma_0; + float _fma_1 = __fmaf_rn(x2, cosv[j_3], x1 * sinv[j_3]); + xs[g_2 * 2 * 8 + 8 + j_3] = _fma_1; + } + } + if (SKIP_MATH == 0 && qk_norm_policy == 1) { + ss = 0.0f; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + ss += xs[g_2 * 2 * 8 + j_4] * xs[g_2 * 2 * 8 + j_4]; + } + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_1; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + wj = ((is_q != 0) ? qw[j_5] : kw[j_5]); + xs[g_2 * 2 * 8 + j_5] = xs[g_2 * 2 * 8 + j_5] * (inv * wj); + } + } + if (head < NUM_Q_HEADS) { + q_dst = (long long)row * (long long)(NUM_Q_HEADS * 128) + (long long)(head * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + q_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + (q_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else if (head < NUM_Q_HEADS + NUM_KV_HEADS) { + kvh = head - NUM_Q_HEADS; + if (use_out_k != 0) { + k_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + k_dst = k_row_base + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + int vh = 0; + long long v_dst = 0; + #pragma unroll + for (int g_3 = 0; g_3 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_3++) { + vh = vhead_of[g_3]; + if (vh < NUM_KV_HEADS) { + if (use_out_v != 0) { + v_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(out_v + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(out_v + (v_dst + 64) + 0) = _iv4; + } + } else { + v_dst = v_row_base + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(v_cache + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(v_cache + (v_dst + 64) + 0) = _iv4; + } + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_binding.cu b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_binding.cu new file mode 100644 index 00000000000..4b6fe8e3d79 --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_binding.cu @@ -0,0 +1,160 @@ +/* + * 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_append_d46a979444ec64e1bda3'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include + +extern "C" __global__ void kernel_cake_fused_qk_rope_append_d46a979444ec64e1bda3(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page); + + +namespace cake_host_shim_c25c648e3fb07c7a { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_q_indptr, TensorView arg_seq_lens, TensorView arg_page_table, TensorView arg_k_cache, TensorView arg_v_cache, TensorView arg_out_q, TensorView arg_out_k, TensorView arg_out_v, TensorView arg_q_norm_weight, TensorView arg_k_norm_weight, int64_t arg_num_rows, int64_t arg_num_requests, int64_t arg_max_pages, int64_t arg_page_size, int64_t arg_k_page_stride, int64_t arg_v_page_stride, int64_t arg_num_row_ctas, int64_t arg_clear_units_per_request, int64_t arg_qk_norm_policy, double arg_eps, int64_t arg_use_out_k, int64_t arg_use_out_v, int64_t arg_clear_last_page, 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_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_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_page_table, "page_table"); + check_dtype(arg_page_table, DLDataType{kDLInt, 32, 1}, "page_table"); + check_contiguous(arg_page_table, "page_table"); + check_cuda_tensor(arg_k_cache, "k_cache"); + check_dtype(arg_k_cache, DLDataType{kDLBfloat, 16, 1}, "k_cache"); + check_contiguous(arg_k_cache, "k_cache"); + check_cuda_tensor(arg_v_cache, "v_cache"); + check_dtype(arg_v_cache, DLDataType{kDLBfloat, 16, 1}, "v_cache"); + check_contiguous(arg_v_cache, "v_cache"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLBfloat, 16, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLBfloat, 16, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLBfloat, 16, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + 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"); + 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_max_pages >= -2147483648LL && arg_max_pages <= 2147483647LL, ValueError) + << "scalar 'max_pages' value " << arg_max_pages + << " 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_num_row_ctas >= -2147483648LL && arg_num_row_ctas <= 2147483647LL, ValueError) + << "scalar 'num_row_ctas' value " << arg_num_row_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_units_per_request >= -2147483648LL && arg_clear_units_per_request <= 2147483647LL, ValueError) + << "scalar 'clear_units_per_request' value " << arg_clear_units_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_norm_policy >= -2147483648LL && arg_qk_norm_policy <= 2147483647LL, ValueError) + << "scalar 'qk_norm_policy' value " << arg_qk_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_k >= -2147483648LL && arg_use_out_k <= 2147483647LL, ValueError) + << "scalar 'use_out_k' value " << arg_use_out_k + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_v >= -2147483648LL && arg_use_out_v <= 2147483647LL, ValueError) + << "scalar 'use_out_v' value " << arg_use_out_v + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_last_page >= -2147483648LL && arg_clear_last_page <= 2147483647LL, ValueError) + << "scalar 'clear_last_page' value " << arg_clear_last_page + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_page_table, arg_qkv, "page_table", "qkv"); + check_same_device(arg_k_cache, arg_qkv, "k_cache", "qkv"); + check_same_device(arg_v_cache, arg_qkv, "v_cache", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "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_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "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_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_page_table = static_cast(arg_page_table.data_ptr()); + __nv_bfloat16* p_k_cache = static_cast<__nv_bfloat16*>(arg_k_cache.data_ptr()); + __nv_bfloat16* p_v_cache = static_cast<__nv_bfloat16*>(arg_v_cache.data_ptr()); + __nv_bfloat16* p_out_q = static_cast<__nv_bfloat16*>(arg_out_q.data_ptr()); + __nv_bfloat16* p_out_k = static_cast<__nv_bfloat16*>(arg_out_k.data_ptr()); + __nv_bfloat16* p_out_v = static_cast<__nv_bfloat16*>(arg_out_v.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()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_max_pages = (int32_t)arg_max_pages; + int32_t v_page_size = (int32_t)arg_page_size; + int64_t v_k_page_stride = (int64_t)arg_k_page_stride; + int64_t v_v_page_stride = (int64_t)arg_v_page_stride; + int32_t v_num_row_ctas = (int32_t)arg_num_row_ctas; + int32_t v_clear_units_per_request = (int32_t)arg_clear_units_per_request; + int32_t v_qk_norm_policy = (int32_t)arg_qk_norm_policy; + float v_eps = (float)arg_eps; + int32_t v_use_out_k = (int32_t)arg_use_out_k; + int32_t v_use_out_v = (int32_t)arg_use_out_v; + int32_t v_clear_last_page = (int32_t)arg_clear_last_page; + void* kargs[] = {&p_qkv, &p_cos_sin, &p_q_indptr, &p_seq_lens, &p_page_table, &p_k_cache, &p_v_cache, &p_out_q, &p_out_k, &p_out_v, &p_q_norm_weight, &p_k_norm_weight, &v_num_rows, &v_num_requests, &v_max_pages, &v_page_size, &v_k_page_stride, &v_v_page_stride, &v_num_row_ctas, &v_clear_units_per_request, &v_qk_norm_policy, &v_eps, &v_use_out_k, &v_use_out_v, &v_clear_last_page}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_append_d46a979444ec64e1bda3), grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_append_d46a979444ec64e1bda3 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_c25c648e3fb07c7a + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_c25c648e3fb07c7a::Run); diff --git a/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_kernel.cu b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_kernel.cu new file mode 100644 index 00000000000..d8d67f26ded --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_103a/cake_fused_qk_rope_append_d46a979444ec64e1bda3_kernel.cu @@ -0,0 +1,608 @@ +/* + * 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 THREADS 128 +#define NUM_Q_HEADS 64 +#define NUM_KV_HEADS 8 +#define WARPS_PER_ROW 8 +#define SKIP_MATH 0 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_fused_qk_rope_append_d46a979444ec64e1bda3(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // === Task calls (dependency order) === + int lane_0 = lane; + int warp_1 = warp; + int hsel = lane_0 / 8; + int d0 = lane_0 % 8 * 8; + if (bid >= num_row_ctas) { + int cid = (bid - num_row_ctas) * 4 + warp_1; + if (cid < num_requests * clear_units_per_request && clear_last_page != 0) { + int req_c = cid / clear_units_per_request; + int u = cid % clear_units_per_request; + int cache_sel = u % 2; + int unit_begin = u / 2 * 256; + int seq_len_c = seq_lens[req_c]; + if (seq_len_c > 0) { + int last_c = seq_len_c - 1; + int start_chunk = (last_c % page_size + 1) * (NUM_KV_HEADS * 128 / 8); + int end_chunk = page_size * (NUM_KV_HEADS * 128 / 8); + if (start_chunk < unit_begin + 256 && unit_begin < end_chunk) { + int page_c32 = page_table[req_c * max_pages + last_c / page_size]; + long long page_c = (long long)page_c32; + float zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0.0f; + } + int c = 0; + long long off_c = 0; + if (cache_sel == 0) { + long long base_k = page_c * k_page_stride; + #pragma unroll + for (int k = 0; k < 8; k++) { + c = unit_begin + lane_0 + 32 * k; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (base_k + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else { + long long base_v = page_c * v_page_stride; + #pragma unroll + for (int k_1 = 0; k_1 < 8; k_1++) { + c = unit_begin + lane_0 + 32 * k_1; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(v_cache + (base_v + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } else { + int row = bid / ((WARPS_PER_ROW + 3) / 4) * ((4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW) + warp_1 / WARPS_PER_ROW; + int sub = bid % ((WARPS_PER_ROW + 3) / 4) * 4 + warp_1 % WARPS_PER_ROW; + if (row < num_rows) { + long long row_base = (long long)row * (long long)((NUM_Q_HEADS + 2 * NUM_KV_HEADS) * 128); + float xs[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 2 * 8]; + int head_of[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_h = 0; + #pragma unroll + for (int g = 0; g < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g++) { + head_of[g] = 4 * (sub + WARPS_PER_ROW * g) + hsel; + src_h = ((head_of[g] < NUM_Q_HEADS + NUM_KV_HEADS) ? head_of[g] : NUM_Q_HEADS + NUM_KV_HEADS - 1); + { + const uint4* _vptr_0 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0)) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_0[_blk].x), "=r"(_vld_0[_blk].y), "=r"(_vld_0[_blk].z), "=r"(_vld_0[_blk].w) : "l"((const void*)(_vptr_0 + _blk)) : "memory"); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #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"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0 + 64)) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_1[_blk].x), "=r"(_vld_1[_blk].y), "=r"(_vld_1[_blk].z), "=r"(_vld_1[_blk].w) : "l"((const void*)(_vptr_1 + _blk)) : "memory"); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #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"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + int vbits[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 8]; + int vhead_of[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_v = 0; + #pragma unroll + for (int g_1 = 0; g_1 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_1++) { + vhead_of[g_1] = 4 * (sub + WARPS_PER_ROW * g_1) + hsel; + src_v = ((vhead_of[g_1] < NUM_KV_HEADS) ? vhead_of[g_1] : NUM_KV_HEADS - 1); + { + const int4* _ivptr_2 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0)) + 0); + int4 _ivld_2; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_2.x), "=r"(_ivld_2.y), "=r"(_ivld_2.z), "=r"(_ivld_2.w) : "l"((const void*)(_ivptr_2)) : "memory"); + vbits[g_1 * 8 + 0] = _ivld_2.x; + vbits[g_1 * 8 + 1] = _ivld_2.y; + vbits[g_1 * 8 + 2] = _ivld_2.z; + vbits[g_1 * 8 + 3] = _ivld_2.w; + } + { + const int4* _ivptr_3 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0 + 64)) + 0); + int4 _ivld_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_3.x), "=r"(_ivld_3.y), "=r"(_ivld_3.z), "=r"(_ivld_3.w) : "l"((const void*)(_ivptr_3)) : "memory"); + vbits[g_1 * 8 + 4 + 0] = _ivld_3.x; + vbits[g_1 * 8 + 4 + 1] = _ivld_3.y; + vbits[g_1 * 8 + 4 + 2] = _ivld_3.z; + vbits[g_1 * 8 + 4 + 3] = _ivld_3.w; + } + } + float qw[16]; + float kw[16]; + if (qk_norm_policy != 0) { + { + 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 + d0 + (0))) : "memory"); + qw[0 + 0] = __uint_as_float(_v4_4_0); + qw[0 + 1] = __uint_as_float(_v4_4_1); + qw[0 + 2] = __uint_as_float(_v4_4_2); + qw[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 + (d0 + 4) + (0))) : "memory"); + qw[4 + 0] = __uint_as_float(_v4_5_0); + qw[4 + 1] = __uint_as_float(_v4_5_1); + qw[4 + 2] = __uint_as_float(_v4_5_2); + qw[4 + 3] = __uint_as_float(_v4_5_3); + } + { + 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*)(q_norm_weight + (d0 + 64) + (0))) : "memory"); + qw[8 + 0] = __uint_as_float(_v4_6_0); + qw[8 + 1] = __uint_as_float(_v4_6_1); + qw[8 + 2] = __uint_as_float(_v4_6_2); + qw[8 + 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*)(q_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + qw[12 + 0] = __uint_as_float(_v4_7_0); + qw[12 + 1] = __uint_as_float(_v4_7_1); + qw[12 + 2] = __uint_as_float(_v4_7_2); + qw[12 + 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 + d0 + (0))) : "memory"); + kw[0 + 0] = __uint_as_float(_v4_8_0); + kw[0 + 1] = __uint_as_float(_v4_8_1); + kw[0 + 2] = __uint_as_float(_v4_8_2); + kw[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 + (d0 + 4) + (0))) : "memory"); + kw[4 + 0] = __uint_as_float(_v4_9_0); + kw[4 + 1] = __uint_as_float(_v4_9_1); + kw[4 + 2] = __uint_as_float(_v4_9_2); + kw[4 + 3] = __uint_as_float(_v4_9_3); + } + { + 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*)(k_norm_weight + (d0 + 64) + (0))) : "memory"); + kw[8 + 0] = __uint_as_float(_v4_10_0); + kw[8 + 1] = __uint_as_float(_v4_10_1); + kw[8 + 2] = __uint_as_float(_v4_10_2); + kw[8 + 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*)(k_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + kw[12 + 0] = __uint_as_float(_v4_11_0); + kw[12 + 1] = __uint_as_float(_v4_11_1); + kw[12 + 2] = __uint_as_float(_v4_11_2); + kw[12 + 3] = __uint_as_float(_v4_11_3); + } + } + int req = -1; + int pos = -1; + int lb[8]; + int le[8]; + int ll[8]; + int bk = 0; + #pragma unroll + for (int k_2 = 0; k_2 < 8; k_2++) { + lb[k_2] = 0; + le[k_2] = 0; + ll[k_2] = 0; + bk = lane_0 + 32 * k_2; + if (bk < num_requests) { + lb[k_2] = q_indptr[bk]; + le[k_2] = q_indptr[bk + 1]; + ll[k_2] = seq_lens[bk]; + } + } + int hit = 0; + unsigned int hit_mask = 0; + int src_lane = 0; + #pragma unroll + for (int k_3 = 0; k_3 < 8; k_3++) { + if (req < 0 && 32 * k_3 < num_requests) { + bk = lane_0 + 32 * k_3; + hit = 0; + if (bk < num_requests) { + if (lb[k_3] <= row && row < le[k_3]) { + hit = 1; + } + } + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_0; + if (hit_mask != 0) { + int _ffs_0 = __ffs(hit_mask); + src_lane = _ffs_0 - 1; + int _shfl_0 = __shfl_sync(0xFFFFFFFF, bk, src_lane); + req = _shfl_0; + int _shfl_1 = __shfl_sync(0xFFFFFFFF, row + ll[k_3] - le[k_3], src_lane); + pos = _shfl_1; + } + } + } + if (req < 0) { + #pragma unroll 1 + for (int base = 256; base < num_requests; base += 32) { + int b = base + lane_0; + hit = 0; + int b_begin = 0; + int b_end = 0; + int b_len = 0; + if (b < num_requests) { + b_begin = q_indptr[b]; + b_end = q_indptr[b + 1]; + b_len = seq_lens[b]; + if (b_begin <= row && row < b_end) { + hit = 1; + } + } + unsigned int _vote_1 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_1; + if (hit_mask != 0) { + int _ffs_1 = __ffs(hit_mask); + src_lane = _ffs_1 - 1; + int _shfl_2 = __shfl_sync(0xFFFFFFFF, b, src_lane); + req = _shfl_2; + int _shfl_3 = __shfl_sync(0xFFFFFFFF, row + b_len - b_end, src_lane); + pos = _shfl_3; + break; + } + } + } + if (req >= 0 && pos >= 0) { + int page32 = page_table[req * max_pages + pos / page_size]; + long long page = (long long)page32; + int slot = pos % page_size; + long long k_row_base = page * k_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long v_row_base = page * v_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long cs_base = (long long)pos * 128; + float cosv[8]; + float sinv[8]; + { + 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 + (long long)d0) + (0))) : "memory"); + cosv[0 + 0] = __uint_as_float(_v4_12_0); + cosv[0 + 1] = __uint_as_float(_v4_12_1); + cosv[0 + 2] = __uint_as_float(_v4_12_2); + cosv[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 + (long long)d0 + 4) + (0))) : "memory"); + cosv[4 + 0] = __uint_as_float(_v4_13_0); + cosv[4 + 1] = __uint_as_float(_v4_13_1); + cosv[4 + 2] = __uint_as_float(_v4_13_2); + cosv[4 + 3] = __uint_as_float(_v4_13_3); + } + { + unsigned _v4_14_0; + unsigned _v4_14_1; + unsigned _v4_14_2; + unsigned _v4_14_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_14_0), "=r"(_v4_14_1), "=r"(_v4_14_2), "=r"(_v4_14_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64)) + (0))) : "memory"); + sinv[0 + 0] = __uint_as_float(_v4_14_0); + sinv[0 + 1] = __uint_as_float(_v4_14_1); + sinv[0 + 2] = __uint_as_float(_v4_14_2); + sinv[0 + 3] = __uint_as_float(_v4_14_3); + } + { + unsigned _v4_15_0; + unsigned _v4_15_1; + unsigned _v4_15_2; + unsigned _v4_15_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_15_0), "=r"(_v4_15_1), "=r"(_v4_15_2), "=r"(_v4_15_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64) + 4) + (0))) : "memory"); + sinv[4 + 0] = __uint_as_float(_v4_15_0); + sinv[4 + 1] = __uint_as_float(_v4_15_1); + sinv[4 + 2] = __uint_as_float(_v4_15_2); + sinv[4 + 3] = __uint_as_float(_v4_15_3); + } + float ss = 0.0f; + float inv = 0.0f; + float wj = 0.0f; + float x1 = 0.0f; + float x2 = 0.0f; + int head = 0; + int is_q = 0; + int kvh = 0; + long long q_dst = 0; + long long k_dst = 0; + #pragma unroll + for (int g_2 = 0; g_2 < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_2++) { + head = head_of[g_2]; + is_q = ((head < NUM_Q_HEADS) ? 1 : 0); + if (SKIP_MATH == 0 && qk_norm_policy == 2) { + ss = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ss += xs[g_2 * 2 * 8 + j_1] * xs[g_2 * 2 * 8 + j_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_0; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + wj = ((is_q != 0) ? qw[j_2] : kw[j_2]); + xs[g_2 * 2 * 8 + j_2] = xs[g_2 * 2 * 8 + j_2] * (inv * wj); + } + } + { + #pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x1 = xs[g_2 * 2 * 8 + j_3]; + x2 = xs[g_2 * 2 * 8 + 8 + j_3]; + float _fma_0 = __fmaf_rn(x1, cosv[j_3], -(x2 * sinv[j_3])); + xs[g_2 * 2 * 8 + j_3] = _fma_0; + float _fma_1 = __fmaf_rn(x2, cosv[j_3], x1 * sinv[j_3]); + xs[g_2 * 2 * 8 + 8 + j_3] = _fma_1; + } + } + if (SKIP_MATH == 0 && qk_norm_policy == 1) { + ss = 0.0f; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + ss += xs[g_2 * 2 * 8 + j_4] * xs[g_2 * 2 * 8 + j_4]; + } + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_1; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + wj = ((is_q != 0) ? qw[j_5] : kw[j_5]); + xs[g_2 * 2 * 8 + j_5] = xs[g_2 * 2 * 8 + j_5] * (inv * wj); + } + } + if (head < NUM_Q_HEADS) { + q_dst = (long long)row * (long long)(NUM_Q_HEADS * 128) + (long long)(head * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + q_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + (q_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else if (head < NUM_Q_HEADS + NUM_KV_HEADS) { + kvh = head - NUM_Q_HEADS; + if (use_out_k != 0) { + k_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + k_dst = k_row_base + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + int vh = 0; + long long v_dst = 0; + #pragma unroll + for (int g_3 = 0; g_3 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_3++) { + vh = vhead_of[g_3]; + if (vh < NUM_KV_HEADS) { + if (use_out_v != 0) { + v_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(out_v + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(out_v + (v_dst + 64) + 0) = _iv4; + } + } else { + v_dst = v_row_base + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(v_cache + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(v_cache + (v_dst + 64) + 0) = _iv4; + } + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_binding.cu b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_binding.cu new file mode 100644 index 00000000000..3cf08678145 --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_binding.cu @@ -0,0 +1,160 @@ +/* + * 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_append_81fc2fe5ac8bb14706e9'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include + +extern "C" __global__ void kernel_cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page); + + +namespace cake_host_shim_dde02df85b9cbdfb { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_q_indptr, TensorView arg_seq_lens, TensorView arg_page_table, TensorView arg_k_cache, TensorView arg_v_cache, TensorView arg_out_q, TensorView arg_out_k, TensorView arg_out_v, TensorView arg_q_norm_weight, TensorView arg_k_norm_weight, int64_t arg_num_rows, int64_t arg_num_requests, int64_t arg_max_pages, int64_t arg_page_size, int64_t arg_k_page_stride, int64_t arg_v_page_stride, int64_t arg_num_row_ctas, int64_t arg_clear_units_per_request, int64_t arg_qk_norm_policy, double arg_eps, int64_t arg_use_out_k, int64_t arg_use_out_v, int64_t arg_clear_last_page, 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_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_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_page_table, "page_table"); + check_dtype(arg_page_table, DLDataType{kDLInt, 32, 1}, "page_table"); + check_contiguous(arg_page_table, "page_table"); + check_cuda_tensor(arg_k_cache, "k_cache"); + check_dtype(arg_k_cache, DLDataType{kDLBfloat, 16, 1}, "k_cache"); + check_contiguous(arg_k_cache, "k_cache"); + check_cuda_tensor(arg_v_cache, "v_cache"); + check_dtype(arg_v_cache, DLDataType{kDLBfloat, 16, 1}, "v_cache"); + check_contiguous(arg_v_cache, "v_cache"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLBfloat, 16, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLBfloat, 16, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLBfloat, 16, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + 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"); + 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_max_pages >= -2147483648LL && arg_max_pages <= 2147483647LL, ValueError) + << "scalar 'max_pages' value " << arg_max_pages + << " 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_num_row_ctas >= -2147483648LL && arg_num_row_ctas <= 2147483647LL, ValueError) + << "scalar 'num_row_ctas' value " << arg_num_row_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_units_per_request >= -2147483648LL && arg_clear_units_per_request <= 2147483647LL, ValueError) + << "scalar 'clear_units_per_request' value " << arg_clear_units_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_norm_policy >= -2147483648LL && arg_qk_norm_policy <= 2147483647LL, ValueError) + << "scalar 'qk_norm_policy' value " << arg_qk_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_k >= -2147483648LL && arg_use_out_k <= 2147483647LL, ValueError) + << "scalar 'use_out_k' value " << arg_use_out_k + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_v >= -2147483648LL && arg_use_out_v <= 2147483647LL, ValueError) + << "scalar 'use_out_v' value " << arg_use_out_v + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_last_page >= -2147483648LL && arg_clear_last_page <= 2147483647LL, ValueError) + << "scalar 'clear_last_page' value " << arg_clear_last_page + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_page_table, arg_qkv, "page_table", "qkv"); + check_same_device(arg_k_cache, arg_qkv, "k_cache", "qkv"); + check_same_device(arg_v_cache, arg_qkv, "v_cache", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "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_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "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_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_page_table = static_cast(arg_page_table.data_ptr()); + __nv_bfloat16* p_k_cache = static_cast<__nv_bfloat16*>(arg_k_cache.data_ptr()); + __nv_bfloat16* p_v_cache = static_cast<__nv_bfloat16*>(arg_v_cache.data_ptr()); + __nv_bfloat16* p_out_q = static_cast<__nv_bfloat16*>(arg_out_q.data_ptr()); + __nv_bfloat16* p_out_k = static_cast<__nv_bfloat16*>(arg_out_k.data_ptr()); + __nv_bfloat16* p_out_v = static_cast<__nv_bfloat16*>(arg_out_v.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()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_max_pages = (int32_t)arg_max_pages; + int32_t v_page_size = (int32_t)arg_page_size; + int64_t v_k_page_stride = (int64_t)arg_k_page_stride; + int64_t v_v_page_stride = (int64_t)arg_v_page_stride; + int32_t v_num_row_ctas = (int32_t)arg_num_row_ctas; + int32_t v_clear_units_per_request = (int32_t)arg_clear_units_per_request; + int32_t v_qk_norm_policy = (int32_t)arg_qk_norm_policy; + float v_eps = (float)arg_eps; + int32_t v_use_out_k = (int32_t)arg_use_out_k; + int32_t v_use_out_v = (int32_t)arg_use_out_v; + int32_t v_clear_last_page = (int32_t)arg_clear_last_page; + void* kargs[] = {&p_qkv, &p_cos_sin, &p_q_indptr, &p_seq_lens, &p_page_table, &p_k_cache, &p_v_cache, &p_out_q, &p_out_k, &p_out_v, &p_q_norm_weight, &p_k_norm_weight, &v_num_rows, &v_num_requests, &v_max_pages, &v_page_size, &v_k_page_stride, &v_v_page_stride, &v_num_row_ctas, &v_clear_units_per_request, &v_qk_norm_policy, &v_eps, &v_use_out_k, &v_use_out_v, &v_clear_last_page}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9), grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_dde02df85b9cbdfb + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_dde02df85b9cbdfb::Run); diff --git a/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_kernel.cu b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_kernel.cu new file mode 100644 index 00000000000..6110bbf5713 --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9_kernel.cu @@ -0,0 +1,608 @@ +/* + * 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 THREADS 128 +#define NUM_Q_HEADS 64 +#define NUM_KV_HEADS 8 +#define WARPS_PER_ROW 8 +#define SKIP_MATH 0 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_fused_qk_rope_append_81fc2fe5ac8bb14706e9(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // === Task calls (dependency order) === + int lane_0 = lane; + int warp_1 = warp; + int hsel = lane_0 / 8; + int d0 = lane_0 % 8 * 8; + if (bid >= num_row_ctas) { + int cid = (bid - num_row_ctas) * 4 + warp_1; + if (cid < num_requests * clear_units_per_request && clear_last_page != 0) { + int req_c = cid / clear_units_per_request; + int u = cid % clear_units_per_request; + int cache_sel = u % 2; + int unit_begin = u / 2 * 256; + int seq_len_c = seq_lens[req_c]; + if (seq_len_c > 0) { + int last_c = seq_len_c - 1; + int start_chunk = (last_c % page_size + 1) * (NUM_KV_HEADS * 128 / 8); + int end_chunk = page_size * (NUM_KV_HEADS * 128 / 8); + if (start_chunk < unit_begin + 256 && unit_begin < end_chunk) { + int page_c32 = page_table[req_c * max_pages + last_c / page_size]; + long long page_c = (long long)page_c32; + float zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0.0f; + } + int c = 0; + long long off_c = 0; + if (cache_sel == 0) { + long long base_k = page_c * k_page_stride; + #pragma unroll + for (int k = 0; k < 8; k++) { + c = unit_begin + lane_0 + 32 * k; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (base_k + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else { + long long base_v = page_c * v_page_stride; + #pragma unroll + for (int k_1 = 0; k_1 < 8; k_1++) { + c = unit_begin + lane_0 + 32 * k_1; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(v_cache + (base_v + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } else { + int row = bid / ((WARPS_PER_ROW + 3) / 4) * ((4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW) + warp_1 / WARPS_PER_ROW; + int sub = bid % ((WARPS_PER_ROW + 3) / 4) * 4 + warp_1 % WARPS_PER_ROW; + if (row < num_rows) { + long long row_base = (long long)row * (long long)((NUM_Q_HEADS + 2 * NUM_KV_HEADS) * 128); + float xs[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 2 * 8]; + int head_of[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_h = 0; + #pragma unroll + for (int g = 0; g < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g++) { + head_of[g] = 4 * (sub + WARPS_PER_ROW * g) + hsel; + src_h = ((head_of[g] < NUM_Q_HEADS + NUM_KV_HEADS) ? head_of[g] : NUM_Q_HEADS + NUM_KV_HEADS - 1); + { + const uint4* _vptr_0 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0)) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_0[_blk].x), "=r"(_vld_0[_blk].y), "=r"(_vld_0[_blk].z), "=r"(_vld_0[_blk].w) : "l"((const void*)(_vptr_0 + _blk)) : "memory"); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #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"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0 + 64)) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_1[_blk].x), "=r"(_vld_1[_blk].y), "=r"(_vld_1[_blk].z), "=r"(_vld_1[_blk].w) : "l"((const void*)(_vptr_1 + _blk)) : "memory"); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #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"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + int vbits[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 8]; + int vhead_of[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_v = 0; + #pragma unroll + for (int g_1 = 0; g_1 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_1++) { + vhead_of[g_1] = 4 * (sub + WARPS_PER_ROW * g_1) + hsel; + src_v = ((vhead_of[g_1] < NUM_KV_HEADS) ? vhead_of[g_1] : NUM_KV_HEADS - 1); + { + const int4* _ivptr_2 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0)) + 0); + int4 _ivld_2; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_2.x), "=r"(_ivld_2.y), "=r"(_ivld_2.z), "=r"(_ivld_2.w) : "l"((const void*)(_ivptr_2)) : "memory"); + vbits[g_1 * 8 + 0] = _ivld_2.x; + vbits[g_1 * 8 + 1] = _ivld_2.y; + vbits[g_1 * 8 + 2] = _ivld_2.z; + vbits[g_1 * 8 + 3] = _ivld_2.w; + } + { + const int4* _ivptr_3 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0 + 64)) + 0); + int4 _ivld_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_3.x), "=r"(_ivld_3.y), "=r"(_ivld_3.z), "=r"(_ivld_3.w) : "l"((const void*)(_ivptr_3)) : "memory"); + vbits[g_1 * 8 + 4 + 0] = _ivld_3.x; + vbits[g_1 * 8 + 4 + 1] = _ivld_3.y; + vbits[g_1 * 8 + 4 + 2] = _ivld_3.z; + vbits[g_1 * 8 + 4 + 3] = _ivld_3.w; + } + } + float qw[16]; + float kw[16]; + if (qk_norm_policy != 0) { + { + 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 + d0 + (0))) : "memory"); + qw[0 + 0] = __uint_as_float(_v4_4_0); + qw[0 + 1] = __uint_as_float(_v4_4_1); + qw[0 + 2] = __uint_as_float(_v4_4_2); + qw[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 + (d0 + 4) + (0))) : "memory"); + qw[4 + 0] = __uint_as_float(_v4_5_0); + qw[4 + 1] = __uint_as_float(_v4_5_1); + qw[4 + 2] = __uint_as_float(_v4_5_2); + qw[4 + 3] = __uint_as_float(_v4_5_3); + } + { + 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*)(q_norm_weight + (d0 + 64) + (0))) : "memory"); + qw[8 + 0] = __uint_as_float(_v4_6_0); + qw[8 + 1] = __uint_as_float(_v4_6_1); + qw[8 + 2] = __uint_as_float(_v4_6_2); + qw[8 + 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*)(q_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + qw[12 + 0] = __uint_as_float(_v4_7_0); + qw[12 + 1] = __uint_as_float(_v4_7_1); + qw[12 + 2] = __uint_as_float(_v4_7_2); + qw[12 + 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 + d0 + (0))) : "memory"); + kw[0 + 0] = __uint_as_float(_v4_8_0); + kw[0 + 1] = __uint_as_float(_v4_8_1); + kw[0 + 2] = __uint_as_float(_v4_8_2); + kw[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 + (d0 + 4) + (0))) : "memory"); + kw[4 + 0] = __uint_as_float(_v4_9_0); + kw[4 + 1] = __uint_as_float(_v4_9_1); + kw[4 + 2] = __uint_as_float(_v4_9_2); + kw[4 + 3] = __uint_as_float(_v4_9_3); + } + { + 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*)(k_norm_weight + (d0 + 64) + (0))) : "memory"); + kw[8 + 0] = __uint_as_float(_v4_10_0); + kw[8 + 1] = __uint_as_float(_v4_10_1); + kw[8 + 2] = __uint_as_float(_v4_10_2); + kw[8 + 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*)(k_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + kw[12 + 0] = __uint_as_float(_v4_11_0); + kw[12 + 1] = __uint_as_float(_v4_11_1); + kw[12 + 2] = __uint_as_float(_v4_11_2); + kw[12 + 3] = __uint_as_float(_v4_11_3); + } + } + int req = -1; + int pos = -1; + int lb[8]; + int le[8]; + int ll[8]; + int bk = 0; + #pragma unroll + for (int k_2 = 0; k_2 < 8; k_2++) { + lb[k_2] = 0; + le[k_2] = 0; + ll[k_2] = 0; + bk = lane_0 + 32 * k_2; + if (bk < num_requests) { + lb[k_2] = q_indptr[bk]; + le[k_2] = q_indptr[bk + 1]; + ll[k_2] = seq_lens[bk]; + } + } + int hit = 0; + unsigned int hit_mask = 0; + int src_lane = 0; + #pragma unroll + for (int k_3 = 0; k_3 < 8; k_3++) { + if (req < 0 && 32 * k_3 < num_requests) { + bk = lane_0 + 32 * k_3; + hit = 0; + if (bk < num_requests) { + if (lb[k_3] <= row && row < le[k_3]) { + hit = 1; + } + } + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_0; + if (hit_mask != 0) { + int _ffs_0 = __ffs(hit_mask); + src_lane = _ffs_0 - 1; + int _shfl_0 = __shfl_sync(0xFFFFFFFF, bk, src_lane); + req = _shfl_0; + int _shfl_1 = __shfl_sync(0xFFFFFFFF, row + ll[k_3] - le[k_3], src_lane); + pos = _shfl_1; + } + } + } + if (req < 0) { + #pragma unroll 1 + for (int base = 256; base < num_requests; base += 32) { + int b = base + lane_0; + hit = 0; + int b_begin = 0; + int b_end = 0; + int b_len = 0; + if (b < num_requests) { + b_begin = q_indptr[b]; + b_end = q_indptr[b + 1]; + b_len = seq_lens[b]; + if (b_begin <= row && row < b_end) { + hit = 1; + } + } + unsigned int _vote_1 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_1; + if (hit_mask != 0) { + int _ffs_1 = __ffs(hit_mask); + src_lane = _ffs_1 - 1; + int _shfl_2 = __shfl_sync(0xFFFFFFFF, b, src_lane); + req = _shfl_2; + int _shfl_3 = __shfl_sync(0xFFFFFFFF, row + b_len - b_end, src_lane); + pos = _shfl_3; + break; + } + } + } + if (req >= 0 && pos >= 0) { + int page32 = page_table[req * max_pages + pos / page_size]; + long long page = (long long)page32; + int slot = pos % page_size; + long long k_row_base = page * k_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long v_row_base = page * v_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long cs_base = (long long)pos * 128; + float cosv[8]; + float sinv[8]; + { + 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 + (long long)d0) + (0))) : "memory"); + cosv[0 + 0] = __uint_as_float(_v4_12_0); + cosv[0 + 1] = __uint_as_float(_v4_12_1); + cosv[0 + 2] = __uint_as_float(_v4_12_2); + cosv[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 + (long long)d0 + 4) + (0))) : "memory"); + cosv[4 + 0] = __uint_as_float(_v4_13_0); + cosv[4 + 1] = __uint_as_float(_v4_13_1); + cosv[4 + 2] = __uint_as_float(_v4_13_2); + cosv[4 + 3] = __uint_as_float(_v4_13_3); + } + { + unsigned _v4_14_0; + unsigned _v4_14_1; + unsigned _v4_14_2; + unsigned _v4_14_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_14_0), "=r"(_v4_14_1), "=r"(_v4_14_2), "=r"(_v4_14_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64)) + (0))) : "memory"); + sinv[0 + 0] = __uint_as_float(_v4_14_0); + sinv[0 + 1] = __uint_as_float(_v4_14_1); + sinv[0 + 2] = __uint_as_float(_v4_14_2); + sinv[0 + 3] = __uint_as_float(_v4_14_3); + } + { + unsigned _v4_15_0; + unsigned _v4_15_1; + unsigned _v4_15_2; + unsigned _v4_15_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_15_0), "=r"(_v4_15_1), "=r"(_v4_15_2), "=r"(_v4_15_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64) + 4) + (0))) : "memory"); + sinv[4 + 0] = __uint_as_float(_v4_15_0); + sinv[4 + 1] = __uint_as_float(_v4_15_1); + sinv[4 + 2] = __uint_as_float(_v4_15_2); + sinv[4 + 3] = __uint_as_float(_v4_15_3); + } + float ss = 0.0f; + float inv = 0.0f; + float wj = 0.0f; + float x1 = 0.0f; + float x2 = 0.0f; + int head = 0; + int is_q = 0; + int kvh = 0; + long long q_dst = 0; + long long k_dst = 0; + #pragma unroll + for (int g_2 = 0; g_2 < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_2++) { + head = head_of[g_2]; + is_q = ((head < NUM_Q_HEADS) ? 1 : 0); + if (SKIP_MATH == 0 && qk_norm_policy == 2) { + ss = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ss += xs[g_2 * 2 * 8 + j_1] * xs[g_2 * 2 * 8 + j_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_0; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + wj = ((is_q != 0) ? qw[j_2] : kw[j_2]); + xs[g_2 * 2 * 8 + j_2] = xs[g_2 * 2 * 8 + j_2] * (inv * wj); + } + } + { + #pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x1 = xs[g_2 * 2 * 8 + j_3]; + x2 = xs[g_2 * 2 * 8 + 8 + j_3]; + float _fma_0 = __fmaf_rn(x1, cosv[j_3], -(x2 * sinv[j_3])); + xs[g_2 * 2 * 8 + j_3] = _fma_0; + float _fma_1 = __fmaf_rn(x2, cosv[j_3], x1 * sinv[j_3]); + xs[g_2 * 2 * 8 + 8 + j_3] = _fma_1; + } + } + if (SKIP_MATH == 0 && qk_norm_policy == 1) { + ss = 0.0f; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + ss += xs[g_2 * 2 * 8 + j_4] * xs[g_2 * 2 * 8 + j_4]; + } + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_1; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + wj = ((is_q != 0) ? qw[j_5] : kw[j_5]); + xs[g_2 * 2 * 8 + j_5] = xs[g_2 * 2 * 8 + j_5] * (inv * wj); + } + } + if (head < NUM_Q_HEADS) { + q_dst = (long long)row * (long long)(NUM_Q_HEADS * 128) + (long long)(head * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + q_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + (q_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else if (head < NUM_Q_HEADS + NUM_KV_HEADS) { + kvh = head - NUM_Q_HEADS; + if (use_out_k != 0) { + k_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + k_dst = k_row_base + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + int vh = 0; + long long v_dst = 0; + #pragma unroll + for (int g_3 = 0; g_3 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_3++) { + vh = vhead_of[g_3]; + if (vh < NUM_KV_HEADS) { + if (use_out_v != 0) { + v_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(out_v + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(out_v + (v_dst + 64) + 0) = _iv4; + } + } else { + v_dst = v_row_base + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(v_cache + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(v_cache + (v_dst + 64) + 0) = _iv4; + } + } + } + } + } + } + } +} + +} // extern "C" diff --git a/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_binding.cu b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_binding.cu new file mode 100644 index 00000000000..f48129fc894 --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_binding.cu @@ -0,0 +1,160 @@ +/* + * 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_append_ef57b221493533ee8018'. +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +#include + +extern "C" __global__ void kernel_cake_fused_qk_rope_append_ef57b221493533ee8018(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page); + + +namespace cake_host_shim_187482d418020437 { + +using tvm::ffi::TensorView; + +void Run(TensorView arg_qkv, TensorView arg_cos_sin, TensorView arg_q_indptr, TensorView arg_seq_lens, TensorView arg_page_table, TensorView arg_k_cache, TensorView arg_v_cache, TensorView arg_out_q, TensorView arg_out_k, TensorView arg_out_v, TensorView arg_q_norm_weight, TensorView arg_k_norm_weight, int64_t arg_num_rows, int64_t arg_num_requests, int64_t arg_max_pages, int64_t arg_page_size, int64_t arg_k_page_stride, int64_t arg_v_page_stride, int64_t arg_num_row_ctas, int64_t arg_clear_units_per_request, int64_t arg_qk_norm_policy, double arg_eps, int64_t arg_use_out_k, int64_t arg_use_out_v, int64_t arg_clear_last_page, 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_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_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_page_table, "page_table"); + check_dtype(arg_page_table, DLDataType{kDLInt, 32, 1}, "page_table"); + check_contiguous(arg_page_table, "page_table"); + check_cuda_tensor(arg_k_cache, "k_cache"); + check_dtype(arg_k_cache, DLDataType{kDLBfloat, 16, 1}, "k_cache"); + check_contiguous(arg_k_cache, "k_cache"); + check_cuda_tensor(arg_v_cache, "v_cache"); + check_dtype(arg_v_cache, DLDataType{kDLBfloat, 16, 1}, "v_cache"); + check_contiguous(arg_v_cache, "v_cache"); + check_cuda_tensor(arg_out_q, "out_q"); + check_dtype(arg_out_q, DLDataType{kDLBfloat, 16, 1}, "out_q"); + check_contiguous(arg_out_q, "out_q"); + check_cuda_tensor(arg_out_k, "out_k"); + check_dtype(arg_out_k, DLDataType{kDLBfloat, 16, 1}, "out_k"); + check_contiguous(arg_out_k, "out_k"); + check_cuda_tensor(arg_out_v, "out_v"); + check_dtype(arg_out_v, DLDataType{kDLBfloat, 16, 1}, "out_v"); + check_contiguous(arg_out_v, "out_v"); + 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"); + 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_max_pages >= -2147483648LL && arg_max_pages <= 2147483647LL, ValueError) + << "scalar 'max_pages' value " << arg_max_pages + << " 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_num_row_ctas >= -2147483648LL && arg_num_row_ctas <= 2147483647LL, ValueError) + << "scalar 'num_row_ctas' value " << arg_num_row_ctas + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_units_per_request >= -2147483648LL && arg_clear_units_per_request <= 2147483647LL, ValueError) + << "scalar 'clear_units_per_request' value " << arg_clear_units_per_request + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_qk_norm_policy >= -2147483648LL && arg_qk_norm_policy <= 2147483647LL, ValueError) + << "scalar 'qk_norm_policy' value " << arg_qk_norm_policy + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_k >= -2147483648LL && arg_use_out_k <= 2147483647LL, ValueError) + << "scalar 'use_out_k' value " << arg_use_out_k + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_use_out_v >= -2147483648LL && arg_use_out_v <= 2147483647LL, ValueError) + << "scalar 'use_out_v' value " << arg_use_out_v + << " is outside i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(arg_clear_last_page >= -2147483648LL && arg_clear_last_page <= 2147483647LL, ValueError) + << "scalar 'clear_last_page' value " << arg_clear_last_page + << " is outside i32 range [-2147483648, 2147483647]"; + check_same_device(arg_cos_sin, arg_qkv, "cos_sin", "qkv"); + check_same_device(arg_q_indptr, arg_qkv, "q_indptr", "qkv"); + check_same_device(arg_seq_lens, arg_qkv, "seq_lens", "qkv"); + check_same_device(arg_page_table, arg_qkv, "page_table", "qkv"); + check_same_device(arg_k_cache, arg_qkv, "k_cache", "qkv"); + check_same_device(arg_v_cache, arg_qkv, "v_cache", "qkv"); + check_same_device(arg_out_q, arg_qkv, "out_q", "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_norm_weight, arg_qkv, "q_norm_weight", "qkv"); + check_same_device(arg_k_norm_weight, arg_qkv, "k_norm_weight", "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_q_indptr = static_cast(arg_q_indptr.data_ptr()); + int* p_seq_lens = static_cast(arg_seq_lens.data_ptr()); + int* p_page_table = static_cast(arg_page_table.data_ptr()); + __nv_bfloat16* p_k_cache = static_cast<__nv_bfloat16*>(arg_k_cache.data_ptr()); + __nv_bfloat16* p_v_cache = static_cast<__nv_bfloat16*>(arg_v_cache.data_ptr()); + __nv_bfloat16* p_out_q = static_cast<__nv_bfloat16*>(arg_out_q.data_ptr()); + __nv_bfloat16* p_out_k = static_cast<__nv_bfloat16*>(arg_out_k.data_ptr()); + __nv_bfloat16* p_out_v = static_cast<__nv_bfloat16*>(arg_out_v.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()); + int32_t v_num_rows = (int32_t)arg_num_rows; + int32_t v_num_requests = (int32_t)arg_num_requests; + int32_t v_max_pages = (int32_t)arg_max_pages; + int32_t v_page_size = (int32_t)arg_page_size; + int64_t v_k_page_stride = (int64_t)arg_k_page_stride; + int64_t v_v_page_stride = (int64_t)arg_v_page_stride; + int32_t v_num_row_ctas = (int32_t)arg_num_row_ctas; + int32_t v_clear_units_per_request = (int32_t)arg_clear_units_per_request; + int32_t v_qk_norm_policy = (int32_t)arg_qk_norm_policy; + float v_eps = (float)arg_eps; + int32_t v_use_out_k = (int32_t)arg_use_out_k; + int32_t v_use_out_v = (int32_t)arg_use_out_v; + int32_t v_clear_last_page = (int32_t)arg_clear_last_page; + void* kargs[] = {&p_qkv, &p_cos_sin, &p_q_indptr, &p_seq_lens, &p_page_table, &p_k_cache, &p_v_cache, &p_out_q, &p_out_k, &p_out_v, &p_q_norm_weight, &p_k_norm_weight, &v_num_rows, &v_num_requests, &v_max_pages, &v_page_size, &v_k_page_stride, &v_v_page_stride, &v_num_row_ctas, &v_clear_units_per_request, &v_qk_norm_policy, &v_eps, &v_use_out_k, &v_use_out_v, &v_clear_last_page}; + + + dim3 grid((uint32_t)grid_x, (uint32_t)grid_y, (uint32_t)grid_z); + dim3 block(128u, 1u, 1u); + + cudaError_t launch_status = cudaLaunchKernel( + reinterpret_cast(kernel_cake_fused_qk_rope_append_ef57b221493533ee8018), grid, block, kargs, 0u, stream); + TVM_FFI_CHECK(launch_status == cudaSuccess, RuntimeError) + << "cudaLaunchKernel for kernel_cake_fused_qk_rope_append_ef57b221493533ee8018 failed: " + << cudaGetErrorString(launch_status); +} + +} // namespace cake_host_shim_187482d418020437 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_host_shim_187482d418020437::Run); diff --git a/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_kernel.cu b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_kernel.cu new file mode 100644 index 00000000000..fa63e5e5d23 --- /dev/null +++ b/csrc/cake_fused_qk_rope_append/sm_90a/cake_fused_qk_rope_append_ef57b221493533ee8018_kernel.cu @@ -0,0 +1,608 @@ +/* + * 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 THREADS 128 +#define NUM_Q_HEADS 8 +#define NUM_KV_HEADS 1 +#define WARPS_PER_ROW 4 +#define SKIP_MATH 0 + +#include + +extern "C" { + +__global__ __launch_bounds__(128) void +kernel_cake_fused_qk_rope_append_ef57b221493533ee8018(__nv_bfloat16* __restrict__ qkv, float* __restrict__ cos_sin, int* __restrict__ q_indptr, int* __restrict__ seq_lens, int* __restrict__ page_table, __nv_bfloat16* __restrict__ k_cache, __nv_bfloat16* __restrict__ v_cache, __nv_bfloat16* __restrict__ out_q, __nv_bfloat16* __restrict__ out_k, __nv_bfloat16* __restrict__ out_v, float* __restrict__ q_norm_weight, float* __restrict__ k_norm_weight, int num_rows, int num_requests, int max_pages, int page_size, long long k_page_stride, long long v_page_stride, int num_row_ctas, int clear_units_per_request, int qk_norm_policy, float eps, int use_out_k, int use_out_v, int clear_last_page) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + const int cta_rank = 0; + + // === Task calls (dependency order) === + int lane_0 = lane; + int warp_1 = warp; + int hsel = lane_0 / 8; + int d0 = lane_0 % 8 * 8; + if (bid >= num_row_ctas) { + int cid = (bid - num_row_ctas) * 4 + warp_1; + if (cid < num_requests * clear_units_per_request && clear_last_page != 0) { + int req_c = cid / clear_units_per_request; + int u = cid % clear_units_per_request; + int cache_sel = u % 2; + int unit_begin = u / 2 * 256; + int seq_len_c = seq_lens[req_c]; + if (seq_len_c > 0) { + int last_c = seq_len_c - 1; + int start_chunk = (last_c % page_size + 1) * (NUM_KV_HEADS * 128 / 8); + int end_chunk = page_size * (NUM_KV_HEADS * 128 / 8); + if (start_chunk < unit_begin + 256 && unit_begin < end_chunk) { + int page_c32 = page_table[req_c * max_pages + last_c / page_size]; + long long page_c = (long long)page_c32; + float zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0.0f; + } + int c = 0; + long long off_c = 0; + if (cache_sel == 0) { + long long base_k = page_c * k_page_stride; + #pragma unroll + for (int k = 0; k < 8; k++) { + c = unit_begin + lane_0 + 32 * k; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (base_k + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } else { + long long base_v = page_c * v_page_stride; + #pragma unroll + for (int k_1 = 0; k_1 < 8; k_1++) { + c = unit_begin + lane_0 + 32 * k_1; + if (c >= start_chunk && c < end_chunk) { + off_c = (long long)c * 8; + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(zeros[0 + 0], zeros[0 + 1]); + _pk[1] = __floats2bfloat162_rn(zeros[0 + 2], zeros[0 + 3]); + _pk[2] = __floats2bfloat162_rn(zeros[0 + 4], zeros[0 + 5]); + _pk[3] = __floats2bfloat162_rn(zeros[0 + 6], zeros[0 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(v_cache + (base_v + off_c)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + } + } + } + } else { + int row = bid / ((WARPS_PER_ROW + 3) / 4) * ((4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW) + warp_1 / WARPS_PER_ROW; + int sub = bid % ((WARPS_PER_ROW + 3) / 4) * 4 + warp_1 % WARPS_PER_ROW; + if (row < num_rows) { + long long row_base = (long long)row * (long long)((NUM_Q_HEADS + 2 * NUM_KV_HEADS) * 128); + float xs[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 2 * 8]; + int head_of[((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_h = 0; + #pragma unroll + for (int g = 0; g < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g++) { + head_of[g] = 4 * (sub + WARPS_PER_ROW * g) + hsel; + src_h = ((head_of[g] < NUM_Q_HEADS + NUM_KV_HEADS) ? head_of[g] : NUM_Q_HEADS + NUM_KV_HEADS - 1); + { + const uint4* _vptr_0 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0)) + 0); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_0[_blk].x), "=r"(_vld_0[_blk].y), "=r"(_vld_0[_blk].z), "=r"(_vld_0[_blk].w) : "l"((const void*)(_vptr_0 + _blk)) : "memory"); + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #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"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(qkv + (row_base + (long long)(src_h * 128 + d0 + 64)) + 0); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_vld_1[_blk].x), "=r"(_vld_1[_blk].y), "=r"(_vld_1[_blk].z), "=r"(_vld_1[_blk].w) : "l"((const void*)(_vptr_1 + _blk)) : "memory"); + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #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"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[0]), "=f"((&xs[g * 2 * 8 + 8 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + int vbits[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW * 8]; + int vhead_of[((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW]; + int src_v = 0; + #pragma unroll + for (int g_1 = 0; g_1 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_1++) { + vhead_of[g_1] = 4 * (sub + WARPS_PER_ROW * g_1) + hsel; + src_v = ((vhead_of[g_1] < NUM_KV_HEADS) ? vhead_of[g_1] : NUM_KV_HEADS - 1); + { + const int4* _ivptr_2 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0)) + 0); + int4 _ivld_2; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_2.x), "=r"(_ivld_2.y), "=r"(_ivld_2.z), "=r"(_ivld_2.w) : "l"((const void*)(_ivptr_2)) : "memory"); + vbits[g_1 * 8 + 0] = _ivld_2.x; + vbits[g_1 * 8 + 1] = _ivld_2.y; + vbits[g_1 * 8 + 2] = _ivld_2.z; + vbits[g_1 * 8 + 3] = _ivld_2.w; + } + { + const int4* _ivptr_3 = reinterpret_cast(qkv + (row_base + (long long)((NUM_Q_HEADS + NUM_KV_HEADS + src_v) * 128 + d0 + 64)) + 0); + int4 _ivld_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_ivld_3.x), "=r"(_ivld_3.y), "=r"(_ivld_3.z), "=r"(_ivld_3.w) : "l"((const void*)(_ivptr_3)) : "memory"); + vbits[g_1 * 8 + 4 + 0] = _ivld_3.x; + vbits[g_1 * 8 + 4 + 1] = _ivld_3.y; + vbits[g_1 * 8 + 4 + 2] = _ivld_3.z; + vbits[g_1 * 8 + 4 + 3] = _ivld_3.w; + } + } + float qw[16]; + float kw[16]; + if (qk_norm_policy != 0) { + { + 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 + d0 + (0))) : "memory"); + qw[0 + 0] = __uint_as_float(_v4_4_0); + qw[0 + 1] = __uint_as_float(_v4_4_1); + qw[0 + 2] = __uint_as_float(_v4_4_2); + qw[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 + (d0 + 4) + (0))) : "memory"); + qw[4 + 0] = __uint_as_float(_v4_5_0); + qw[4 + 1] = __uint_as_float(_v4_5_1); + qw[4 + 2] = __uint_as_float(_v4_5_2); + qw[4 + 3] = __uint_as_float(_v4_5_3); + } + { + 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*)(q_norm_weight + (d0 + 64) + (0))) : "memory"); + qw[8 + 0] = __uint_as_float(_v4_6_0); + qw[8 + 1] = __uint_as_float(_v4_6_1); + qw[8 + 2] = __uint_as_float(_v4_6_2); + qw[8 + 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*)(q_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + qw[12 + 0] = __uint_as_float(_v4_7_0); + qw[12 + 1] = __uint_as_float(_v4_7_1); + qw[12 + 2] = __uint_as_float(_v4_7_2); + qw[12 + 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 + d0 + (0))) : "memory"); + kw[0 + 0] = __uint_as_float(_v4_8_0); + kw[0 + 1] = __uint_as_float(_v4_8_1); + kw[0 + 2] = __uint_as_float(_v4_8_2); + kw[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 + (d0 + 4) + (0))) : "memory"); + kw[4 + 0] = __uint_as_float(_v4_9_0); + kw[4 + 1] = __uint_as_float(_v4_9_1); + kw[4 + 2] = __uint_as_float(_v4_9_2); + kw[4 + 3] = __uint_as_float(_v4_9_3); + } + { + 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*)(k_norm_weight + (d0 + 64) + (0))) : "memory"); + kw[8 + 0] = __uint_as_float(_v4_10_0); + kw[8 + 1] = __uint_as_float(_v4_10_1); + kw[8 + 2] = __uint_as_float(_v4_10_2); + kw[8 + 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*)(k_norm_weight + (d0 + 64 + 4) + (0))) : "memory"); + kw[12 + 0] = __uint_as_float(_v4_11_0); + kw[12 + 1] = __uint_as_float(_v4_11_1); + kw[12 + 2] = __uint_as_float(_v4_11_2); + kw[12 + 3] = __uint_as_float(_v4_11_3); + } + } + int req = -1; + int pos = -1; + int lb[8]; + int le[8]; + int ll[8]; + int bk = 0; + #pragma unroll + for (int k_2 = 0; k_2 < 8; k_2++) { + lb[k_2] = 0; + le[k_2] = 0; + ll[k_2] = 0; + bk = lane_0 + 32 * k_2; + if (bk < num_requests) { + lb[k_2] = q_indptr[bk]; + le[k_2] = q_indptr[bk + 1]; + ll[k_2] = seq_lens[bk]; + } + } + int hit = 0; + unsigned int hit_mask = 0; + int src_lane = 0; + #pragma unroll + for (int k_3 = 0; k_3 < 8; k_3++) { + if (req < 0 && 32 * k_3 < num_requests) { + bk = lane_0 + 32 * k_3; + hit = 0; + if (bk < num_requests) { + if (lb[k_3] <= row && row < le[k_3]) { + hit = 1; + } + } + unsigned int _vote_0 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_0; + if (hit_mask != 0) { + int _ffs_0 = __ffs(hit_mask); + src_lane = _ffs_0 - 1; + int _shfl_0 = __shfl_sync(0xFFFFFFFF, bk, src_lane); + req = _shfl_0; + int _shfl_1 = __shfl_sync(0xFFFFFFFF, row + ll[k_3] - le[k_3], src_lane); + pos = _shfl_1; + } + } + } + if (req < 0) { + #pragma unroll 1 + for (int base = 256; base < num_requests; base += 32) { + int b = base + lane_0; + hit = 0; + int b_begin = 0; + int b_end = 0; + int b_len = 0; + if (b < num_requests) { + b_begin = q_indptr[b]; + b_end = q_indptr[b + 1]; + b_len = seq_lens[b]; + if (b_begin <= row && row < b_end) { + hit = 1; + } + } + unsigned int _vote_1 = __ballot_sync(0xFFFFFFFF, hit != 0); + hit_mask = _vote_1; + if (hit_mask != 0) { + int _ffs_1 = __ffs(hit_mask); + src_lane = _ffs_1 - 1; + int _shfl_2 = __shfl_sync(0xFFFFFFFF, b, src_lane); + req = _shfl_2; + int _shfl_3 = __shfl_sync(0xFFFFFFFF, row + b_len - b_end, src_lane); + pos = _shfl_3; + break; + } + } + } + if (req >= 0 && pos >= 0) { + int page32 = page_table[req * max_pages + pos / page_size]; + long long page = (long long)page32; + int slot = pos % page_size; + long long k_row_base = page * k_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long v_row_base = page * v_page_stride + (long long)(slot * (NUM_KV_HEADS * 128)); + long long cs_base = (long long)pos * 128; + float cosv[8]; + float sinv[8]; + { + 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 + (long long)d0) + (0))) : "memory"); + cosv[0 + 0] = __uint_as_float(_v4_12_0); + cosv[0 + 1] = __uint_as_float(_v4_12_1); + cosv[0 + 2] = __uint_as_float(_v4_12_2); + cosv[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 + (long long)d0 + 4) + (0))) : "memory"); + cosv[4 + 0] = __uint_as_float(_v4_13_0); + cosv[4 + 1] = __uint_as_float(_v4_13_1); + cosv[4 + 2] = __uint_as_float(_v4_13_2); + cosv[4 + 3] = __uint_as_float(_v4_13_3); + } + { + unsigned _v4_14_0; + unsigned _v4_14_1; + unsigned _v4_14_2; + unsigned _v4_14_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_14_0), "=r"(_v4_14_1), "=r"(_v4_14_2), "=r"(_v4_14_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64)) + (0))) : "memory"); + sinv[0 + 0] = __uint_as_float(_v4_14_0); + sinv[0 + 1] = __uint_as_float(_v4_14_1); + sinv[0 + 2] = __uint_as_float(_v4_14_2); + sinv[0 + 3] = __uint_as_float(_v4_14_3); + } + { + unsigned _v4_15_0; + unsigned _v4_15_1; + unsigned _v4_15_2; + unsigned _v4_15_3; + asm volatile("ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];" + : "=r"(_v4_15_0), "=r"(_v4_15_1), "=r"(_v4_15_2), "=r"(_v4_15_3) : "l"((const void*)(cos_sin + (cs_base + (long long)(d0 + 64) + 4) + (0))) : "memory"); + sinv[4 + 0] = __uint_as_float(_v4_15_0); + sinv[4 + 1] = __uint_as_float(_v4_15_1); + sinv[4 + 2] = __uint_as_float(_v4_15_2); + sinv[4 + 3] = __uint_as_float(_v4_15_3); + } + float ss = 0.0f; + float inv = 0.0f; + float wj = 0.0f; + float x1 = 0.0f; + float x2 = 0.0f; + int head = 0; + int is_q = 0; + int kvh = 0; + long long q_dst = 0; + long long k_dst = 0; + #pragma unroll + for (int g_2 = 0; g_2 < ((NUM_Q_HEADS + NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_2++) { + head = head_of[g_2]; + is_q = ((head < NUM_Q_HEADS) ? 1 : 0); + if (SKIP_MATH == 0 && qk_norm_policy == 2) { + ss = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 16; j_1++) { + ss += xs[g_2 * 2 * 8 + j_1] * xs[g_2 * 2 * 8 + j_1]; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_2; + float _rsqrt_0 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_0; + #pragma unroll + for (int j_2 = 0; j_2 < 16; j_2++) { + wj = ((is_q != 0) ? qw[j_2] : kw[j_2]); + xs[g_2 * 2 * 8 + j_2] = xs[g_2 * 2 * 8 + j_2] * (inv * wj); + } + } + { + #pragma unroll + for (int j_3 = 0; j_3 < 8; j_3++) { + x1 = xs[g_2 * 2 * 8 + j_3]; + x2 = xs[g_2 * 2 * 8 + 8 + j_3]; + float _fma_0 = __fmaf_rn(x1, cosv[j_3], -(x2 * sinv[j_3])); + xs[g_2 * 2 * 8 + j_3] = _fma_0; + float _fma_1 = __fmaf_rn(x2, cosv[j_3], x1 * sinv[j_3]); + xs[g_2 * 2 * 8 + 8 + j_3] = _fma_1; + } + } + if (SKIP_MATH == 0 && qk_norm_policy == 1) { + ss = 0.0f; + #pragma unroll + for (int j_4 = 0; j_4 < 16; j_4++) { + ss += xs[g_2 * 2 * 8 + j_4] * xs[g_2 * 2 * 8 + j_4]; + } + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, ss, 4); + ss += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, ss, 2); + ss += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, ss, 1); + ss += _shfl_xor_5; + float _rsqrt_1 = rsqrtf(ss * 0.0078125f + eps); + inv = _rsqrt_1; + #pragma unroll + for (int j_5 = 0; j_5 < 16; j_5++) { + wj = ((is_q != 0) ? qw[j_5] : kw[j_5]); + xs[g_2 * 2 * 8 + j_5] = xs[g_2 * 2 * 8 + j_5] * (inv * wj); + } + } + if (head < NUM_Q_HEADS) { + q_dst = (long long)row * (long long)(NUM_Q_HEADS * 128) + (long long)(head * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + q_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_q + (q_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else if (head < NUM_Q_HEADS + NUM_KV_HEADS) { + kvh = head - NUM_Q_HEADS; + if (use_out_k != 0) { + k_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(out_k + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } else { + k_dst = k_row_base + (long long)(kvh * 128 + d0); + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 0], xs[g_2 * 2 * 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 2], xs[g_2 * 2 * 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 4], xs[g_2 * 2 * 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 6], xs[g_2 * 2 * 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + k_dst))[0]) = *reinterpret_cast(&_pk[0]); + } + { + __nv_bfloat162 _pk[4]; + _pk[0] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 0], xs[g_2 * 2 * 8 + 8 + 1]); + _pk[1] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 2], xs[g_2 * 2 * 8 + 8 + 3]); + _pk[2] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 4], xs[g_2 * 2 * 8 + 8 + 5]); + _pk[3] = __floats2bfloat162_rn(xs[g_2 * 2 * 8 + 8 + 6], xs[g_2 * 2 * 8 + 8 + 7]); + *reinterpret_cast(&((__nv_bfloat16*)(k_cache + (k_dst + 64)))[0]) = *reinterpret_cast(&_pk[0]); + } + } + } + } + int vh = 0; + long long v_dst = 0; + #pragma unroll + for (int g_3 = 0; g_3 < ((NUM_KV_HEADS + 4 - 1) / 4 + WARPS_PER_ROW - 1) / WARPS_PER_ROW; g_3++) { + vh = vhead_of[g_3]; + if (vh < NUM_KV_HEADS) { + if (use_out_v != 0) { + v_dst = (long long)row * (long long)(NUM_KV_HEADS * 128) + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(out_v + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(out_v + (v_dst + 64) + 0) = _iv4; + } + } else { + v_dst = v_row_base + (long long)(vh * 128 + d0); + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 0], vbits[g_3 * 8 + 1], vbits[g_3 * 8 + 2], vbits[g_3 * 8 + 3]); + *reinterpret_cast(v_cache + v_dst + 0) = _iv4; + } + { + int4 _iv4 = make_int4(vbits[g_3 * 8 + 4 + 0], vbits[g_3 * 8 + 4 + 1], vbits[g_3 * 8 + 4 + 2], vbits[g_3 * 8 + 4 + 3]); + *reinterpret_cast(v_cache + (v_dst + 64) + 0) = _iv4; + } + } + } + } + } + } + } +} + +} // extern "C" diff --git a/flashinfer/cake_fused_qk_rope_append.py b/flashinfer/cake_fused_qk_rope_append.py new file mode 100644 index 00000000000..3f3af9c21ca --- /dev/null +++ b/flashinfer/cake_fused_qk_rope_append.py @@ -0,0 +1,295 @@ +# 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 + BF16 Q output + paged KV append (one launch). + +:func:`cake_fused_qk_rmsnorm_rope_append_paged_kv_cache` takes the packed BF16 +``qkv`` projection of a decode or small-chunk 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`` default ``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`` table indexed by absolute position, +4. writes Q as BF16 ``[T, Hq, 128]`` and appends K/V to an NHD paged cache + ``[num_pages, page_size, Hkv, 128]`` (or into caller-owned ``out_k`` / + ``out_v`` ``[T, Hkv, 128]``), +5. zeroes the unused rows of every request's last page in both caches + (``clear_unused_last_page_rows=True``), including when K/V go to + caller-owned buffers. + +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). Rows outside every +request or with a negative position produce no output. 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 in/out, SM90 +(H100/H200), SM100 (B200) and SM103 (B300). The kernel is generated from the +Cake Weave schedule ``fused_qk_rmsnorm_rope_paged_kv_append_bf16``; sources and +manifest live under ``csrc/cake_fused_qk_rope_append`` (JIT-only). The call is +CUDA-graph capturable: every launch argument is a by-value scalar or a tensor +pointer, no host synchronisation or workspace is involved. +""" + +from __future__ import annotations + +from typing import Any, Optional + +import torch + +from .jit.cake_fused_qk_rope_append import ( + arch_for, + load_cake_fused_qk_rope_append_module, + stage_for, +) +from .utils import get_compute_capability + +HEAD_DIM = 128 +_VEC = 8 +_CLEAR_UNIT_CHUNKS = 256 +_THREADS = 128 + + +def _clear_units_per_request(page_size: int, num_kv_heads: int) -> int: + chunks = page_size * num_kv_heads * HEAD_DIM // _VEC + return 2 * ((chunks + _CLEAR_UNIT_CHUNKS - 1) // _CLEAR_UNIT_CHUNKS) + + +def launch_plan( + num_rows: int, + num_requests: int, + page_size: int, + num_kv_heads: int, + warps_per_row: int, +) -> dict[str, Any]: + """Grid and host-derived scalars (mirrors the Cake kernel's ``launch_plan``).""" + rows_per_cta = max(1, 4 // warps_per_row) + ctas_per_row = max(1, warps_per_row // 4) + num_row_ctas = ((num_rows + rows_per_cta - 1) // rows_per_cta) * ctas_per_row + units = _clear_units_per_request(page_size, num_kv_heads) + num_clear_ctas = (num_requests * units + 3) // 4 + return { + "num_row_ctas": num_row_ctas, + "clear_units_per_request": units, + "grid": (num_row_ctas + num_clear_ctas, 1, 1), + } + + +def _check(cond: bool, message: str) -> None: + if not cond: + raise ValueError(message) + + +def cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + qkv: torch.Tensor, + cos_sin: torch.Tensor, + seq_lens: torch.Tensor, + q_indptr: torch.Tensor, + page_indices: torch.Tensor, + key_cache: torch.Tensor, + value_cache: torch.Tensor, + *, + num_q_heads: int, + num_kv_heads: int, + qk_norm_policy: int = 0, + q_norm_weight: Optional[torch.Tensor] = None, + k_norm_weight: Optional[torch.Tensor] = None, + eps: float = 1e-6, + out_q: Optional[torch.Tensor] = None, + out_k: Optional[torch.Tensor] = None, + out_v: Optional[torch.Tensor] = None, + clear_unused_last_page_rows: bool = True, +) -> torch.Tensor: + r"""Fused Q/K RMSNorm + NeoX RoPE + BF16 Q output + 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. + cos_sin : 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). + key_cache, value_cache : torch.Tensor + BF16 NHD caches ``[num_pages, page_size, num_kv_heads, 128]`` (written in place). + num_q_heads, num_kv_heads : int + Head configuration; ``(8, 1)`` and ``(64, 8)`` are supported. + qk_norm_policy : int + 0 no norm, 1 RoPE then RMSNorm, 2 RMSNorm then RoPE. + q_norm_weight, k_norm_weight : torch.Tensor, optional + F32 ``[128]`` per-head RMSNorm weights (required unless ``qk_norm_policy == 0``). + eps : float + RMSNorm epsilon (default ``1e-6``). + out_q : torch.Tensor, optional + BF16 ``[T, num_q_heads, 128]`` output buffer (allocated when ``None``). + out_k, out_v : torch.Tensor, optional + BF16 ``[T, num_kv_heads, 128]``; when given, K / V are written there instead + of the caches (the last-page clear still runs). + clear_unused_last_page_rows : bool + Zero rows ``[(seq_lens[b] - 1) % page_size + 1, page_size)`` of every request's + last page in both caches (all kv heads). + + Returns + ------- + out_q : torch.Tensor + BF16 ``[T, num_q_heads, 128]`` normalised and rotated queries. + """ + _check(qk_norm_policy in (0, 1, 2), "qk_norm_policy must be 0, 1 or 2") + _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_rows = int(qkv.shape[0]) + width = (num_q_heads + 2 * num_kv_heads) * HEAD_DIM + _check( + int(qkv.shape[1]) == width, + f"qkv must have {width} columns for this head configuration", + ) + _check( + cos_sin.dtype == torch.float32 + and cos_sin.dim() == 2 + and int(cos_sin.shape[1]) == HEAD_DIM, + "cos_sin 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.bfloat16 + and t.dim() == 4 + and int(t.shape[2]) == num_kv_heads + and int(t.shape[3]) == HEAD_DIM, + f"{name} must be bf16 [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]) + if out_q is None: + out_q = torch.empty( + (num_rows, num_q_heads, HEAD_DIM), dtype=torch.bfloat16, device=device + ) + _check( + out_q.dtype == torch.bfloat16 + and tuple(out_q.shape) == (num_rows, num_q_heads, HEAD_DIM), + "out_q must be bf16 [T, num_q_heads, 128]", + ) + for name, t in (("out_k", out_k), ("out_v", out_v)): + if t is not None: + _check( + t.dtype == torch.bfloat16 + and tuple(t.shape) == (num_rows, num_kv_heads, HEAD_DIM), + f"{name} must be bf16 [T, num_kv_heads, 128]", + ) + 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 (no extra kernel, graph-safe). + placeholder = cos_sin[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]", + ) + tensors = { + "qkv": qkv, + "cos_sin": cos_sin, + "q_indptr": q_indptr, + "seq_lens": seq_lens, + "page_table": page_indices, + "k_cache": key_cache, + "v_cache": value_cache, + "out_q": out_q, + "out_k": out_k if out_k is not None else key_cache, + "out_v": out_v if out_v is not None else value_cache, + "q_norm_weight": q_norm_weight, + "k_norm_weight": k_norm_weight, + } + 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}") + + stage = stage_for(num_q_heads, num_kv_heads) + arch = arch_for(get_compute_capability(device)) + module, record = load_cake_fused_qk_rope_append_module(stage, arch) + warps_per_row = int( + dict(record["route"]).get("specialization", {}).get("WARPS_PER_ROW", 1) + ) + plan = launch_plan(num_rows, num_requests, page_size, num_kv_heads, warps_per_row) + scalars = { + "num_rows": num_rows, + "num_requests": num_requests, + "max_pages": int(page_indices.shape[1]), + "page_size": page_size, + "k_page_stride": int(key_cache.stride(0)), + "v_page_stride": int(value_cache.stride(0)), + "num_row_ctas": plan["num_row_ctas"], + "clear_units_per_request": plan["clear_units_per_request"], + "qk_norm_policy": int(qk_norm_policy), + "eps": float(eps), + "use_out_k": 0 if out_k is None else 1, + "use_out_v": 0 if out_v is None else 1, + "clear_last_page": 1 if clear_unused_last_page_rows else 0, + "grid_x": plan["grid"][0], + "grid_y": 1, + "grid_z": 1, + } + 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 + + +__all__ = ["cake_fused_qk_rmsnorm_rope_append_paged_kv_cache", "launch_plan"] diff --git a/flashinfer/jit/cake_fused_qk_rope_append.py b/flashinfer/jit/cake_fused_qk_rope_append.py new file mode 100644 index 00000000000..5e111bdcc5a --- /dev/null +++ b/flashinfer/jit/cake_fused_qk_rope_append.py @@ -0,0 +1,193 @@ +# 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 + paged KV append (BF16) kernel. + +One source-only manifest lives under ``csrc/cake_fused_qk_rope_append`` +(``cake_fused_qk_rope_append_manifest.json``, schema ``cake.library_export.v5``): +two head-configuration stages (``hq8_hkv1`` for ``(Hq, Hkv) = (8, 1)`` and +``hq64_hkv8`` for ``(64, 8)``) 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_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_append" +_GENERATED_ROOT = f"csrc/{_PACKAGE}" +_MANIFEST_NAME = f"{_PACKAGE}_manifest.json" + +STAGES = ("hq8_hkv1", "hq64_hkv8") +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 stage_for(num_q_heads: int, num_kv_heads: int) -> str: + if (num_q_heads, num_kv_heads) == (8, 1): + return "hq8_hkv1" + if (num_q_heads, num_kv_heads) == (64, 8): + return "hq64_hkv8" + raise ValueError( + f"cake_fused_qk_rope_append supports (num_q_heads, num_kv_heads) in {{(8, 1), (64, 8)}}, " + f"got ({num_q_heads}, {num_kv_heads})" + ) + + +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_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_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_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_append_module(stage: str, arch: str): + """Build and load one module; return ``(module, manifest record)``.""" + spec = gen_cake_fused_qk_rope_append_module(stage, arch) + module = spec.build_and_load() + return module, _record(stage, arch) + + +def build_all_cake_fused_qk_rope_append_modules(arch: str) -> dict[str, str]: + built = {} + for stage in STAGES: + spec = gen_cake_fused_qk_rope_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 3e9f9526f38..b43047f6108 100644 --- a/flashinfer/rope.py +++ b/flashinfer/rope.py @@ -20,6 +20,9 @@ import torch from .api_logging import flashinfer_api +from .cake_fused_qk_rope_append import ( # noqa: F401 (re-exported Cake fused route) + cake_fused_qk_rmsnorm_rope_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_append.py b/tests/attention/test_cake_fused_qk_rope_append.py new file mode 100644 index 00000000000..7449406e7c9 --- /dev/null +++ b/tests/attention/test_cake_fused_qk_rope_append.py @@ -0,0 +1,480 @@ +"""Tests for the Cake fused QK RMSNorm + NeoX RoPE + paged KV append (BF16) kernel. + +Independent torch reference (per-head RMSNorm in f32, NeoX rotation from the cos/sin table, NHD paged +append, last-page clearing). Tolerances follow the Cake BF16 standard: Q/K ``atol = rtol = 1e-2``, +V bit-exact, every cache / output element outside the appended rows and the cleared range untouched. +""" + +import pytest +import torch + +from flashinfer.cake_fused_qk_rope_append import ( + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache, +) +from flashinfer.utils import get_compute_capability + +HEAD_DIM = 128 +POISON = 1234.0 + + +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_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 reference( + qkv, + cos_sin, + seq_lens, + q_indptr, + page_indices, + key_cache, + value_cache, + *, + num_q_heads, + num_kv_heads, + qk_norm_policy, + q_norm_weight, + k_norm_weight, + eps=1e-6, + out_q=None, + out_k=None, + out_v=None, + clear_last_page=True, +): + T = qkv.shape[0] + hq, hkv = num_q_heads, num_kv_heads + 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 = qkv[:, (hq + hkv) * HEAD_DIM :].reshape( + T, hkv, HEAD_DIM + ) # bf16, copied bit-exactly + B = seq_lens.shape[0] + qi = q_indptr.tolist() + sl = seq_lens.tolist() + page_size = key_cache.shape[1] + if out_q is None: + out_q = torch.empty((T, hq, HEAD_DIM), dtype=torch.bfloat16, device=qkv.device) + + def rms(t, w): + var = (t * t).mean(dim=-1, keepdim=True) + return t * torch.rsqrt(var + eps) * w + + def rope(t, cs): + half = HEAD_DIM // 2 + c = cs[:, None, :half] + s = cs[:, None, half:] + x1, x2 = t[..., :half], t[..., half:] + return torch.cat([x1 * c - x2 * s, x2 * c + x1 * s], dim=-1) + + if clear_last_page: + for b in range(B): + if sl[b] <= 0: + continue + last = sl[b] - 1 + page = int(page_indices[b, last // page_size]) + first_unused = last % page_size + 1 + key_cache[page, first_unused:] = 0 + value_cache[page, first_unused:] = 0 + for b in range(B): + rows = list(range(qi[b], qi[b + 1])) + if not rows: + continue + positions = torch.tensor( + [r + sl[b] - qi[b + 1] for r in rows], device=qkv.device + ) + valid = positions >= 0 + rows_t = torch.tensor(rows, device=qkv.device)[valid] + positions = positions[valid] + if rows_t.numel() == 0: + continue + cs = cos_sin[positions] + qb, kb = q[rows_t], k[rows_t] + if qk_norm_policy == 2: + qb, kb = rms(qb, q_norm_weight), rms(kb, k_norm_weight) + qb, kb = rope(qb, cs), rope(kb, cs) + if qk_norm_policy == 1: + qb, kb = rms(qb, q_norm_weight), rms(kb, k_norm_weight) + out_q[rows_t] = qb.to(torch.bfloat16) + kb16 = kb.to(torch.bfloat16) + vb = v[rows_t] + if out_k is not None: + out_k[rows_t] = kb16 + if out_v is not None: + out_v[rows_t] = vb + if out_k is None or out_v is None: + pages = page_indices[b][positions // page_size].long() + slots = positions % page_size + if out_k is None: + key_cache[pages, slots] = kb16 + if out_v is None: + value_cache[pages, slots] = vb + return out_q + + +def make_case( + *, + q_lens, + ctx_lens, + num_q_heads, + num_kv_heads, + page_size, + policy, + seed, + 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) + ) + cos_sin = rotary_cos_sin_table(max(seq_lens) + 1, device) + q_w = ( + (torch.rand(HEAD_DIM, generator=g) + 0.5).to(torch.bfloat16).float().to(device) + ) + k_w = ( + (torch.rand(HEAD_DIM, generator=g) + 0.5).to(torch.bfloat16).float().to(device) + ) + key_cache = torch.full( + (total_pages, page_size, num_kv_heads, HEAD_DIM), + POISON, + dtype=torch.bfloat16, + device=device, + ) + value_cache = torch.full_like(key_cache, POISON) + out_q = torch.full( + (T, num_q_heads, HEAD_DIM), POISON, dtype=torch.bfloat16, device=device + ) + return dict( + qkv=qkv, + cos_sin=cos_sin, + seq_lens=torch.tensor(seq_lens, dtype=torch.int32, device=device), + q_indptr=q_indptr.to(device), + page_indices=page_indices.to(device), + key_cache=key_cache, + value_cache=value_cache, + out_q=out_q, + q_norm_weight=q_w, + k_norm_weight=k_w, + policy=policy, + num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, + ) + + +def run_and_compare(case, *, caller_owned_kv=False, clear=True): + T = case["qkv"].shape[0] + hkv = case["num_kv_heads"] + dev = case["qkv"].device + kw = dict( + num_q_heads=case["num_q_heads"], + num_kv_heads=hkv, + qk_norm_policy=case["policy"], + q_norm_weight=case["q_norm_weight"] if case["policy"] else None, + k_norm_weight=case["k_norm_weight"] if case["policy"] else None, + ) + ok = ov = rok = rov = None + if caller_owned_kv: + ok = torch.full((T, hkv, HEAD_DIM), POISON, dtype=torch.bfloat16, device=dev) + ov = torch.full_like(ok, POISON) + rok, rov = ok.clone(), ov.clone() + kc, vc, oq = ( + case["key_cache"].clone(), + case["value_cache"].clone(), + case["out_q"].clone(), + ) + ret = cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + kc, + vc, + out_q=oq, + out_k=ok, + out_v=ov, + clear_unused_last_page_rows=clear, + **kw, + ) + assert ret is oq + torch.cuda.synchronize() + rkc, rvc, roq = ( + case["key_cache"].clone(), + case["value_cache"].clone(), + case["out_q"].clone(), + ) + reference( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + rkc, + rvc, + num_q_heads=case["num_q_heads"], + num_kv_heads=hkv, + qk_norm_policy=case["policy"], + q_norm_weight=case["q_norm_weight"], + k_norm_weight=case["k_norm_weight"], + out_q=roq, + out_k=rok, + out_v=rov, + clear_last_page=clear, + ) + torch.testing.assert_close(oq.float(), roq.float(), atol=1e-2, rtol=1e-2) + torch.testing.assert_close(kc.float(), rkc.float(), atol=1e-2, rtol=1e-2) + assert torch.equal(vc.view(torch.int16), rvc.view(torch.int16)), ( + "V must be copied bit-exactly" + ) + if caller_owned_kv: + torch.testing.assert_close(ok.float(), rok.float(), atol=1e-2, rtol=1e-2) + assert torch.equal(ov.view(torch.int16), rov.view(torch.int16)) + # untouched regions: everything the reference left at POISON must still be POISON + for mine, ref in ((kc, rkc), (vc, rvc), (oq, roq)): + mask = ref == POISON + assert torch.equal(mine[mask], ref[mask]), ( + "kernel wrote outside the appended / cleared range" + ) + + +@pytest.mark.parametrize("heads", [(8, 1), (64, 8)]) +@pytest.mark.parametrize("policy", [0, 1, 2]) +@pytest.mark.parametrize("page_size", [16, 64]) +def test_decode_batch(heads, policy, page_size): + _skip_unless_supported() + case = make_case( + q_lens=[1] * 32, + ctx_lens=[2048 + 7 * i for i in range(32)], + num_q_heads=heads[0], + num_kv_heads=heads[1], + page_size=page_size, + policy=policy, + seed=892, + ) + run_and_compare(case) + + +@pytest.mark.parametrize("heads", [(8, 1), (64, 8)]) +@pytest.mark.parametrize("policy", [1, 2]) +def test_prefill_chunks_cross_pages(heads, policy): + _skip_unless_supported() + case = make_case( + q_lens=[16, 64, 128, 3], + ctx_lens=[2047, 0, 2000, 63], + num_q_heads=heads[0], + num_kv_heads=heads[1], + page_size=64, + policy=policy, + seed=893, + ) + run_and_compare(case) + + +def test_ragged_with_empty_requests_and_caller_owned_kv(): + _skip_unless_supported() + case = make_case( + q_lens=[5, 0, 1, 33, 0, 2], + ctx_lens=[100, 64, 4095, 0, 1, 127], + num_q_heads=8, + num_kv_heads=1, + page_size=64, + policy=2, + seed=894, + ) + run_and_compare(case) + run_and_compare(case, caller_owned_kv=True) + + +def test_no_clear_leaves_last_page_tail(): + _skip_unless_supported() + case = make_case( + q_lens=[1] * 4, + ctx_lens=[10, 70, 128, 5], + num_q_heads=8, + num_kv_heads=1, + page_size=64, + policy=2, + seed=895, + ) + run_and_compare(case, clear=False) + + +def test_large_batch_lookup_fallback(): + _skip_unless_supported() + case = make_case( + q_lens=[1] * 300, + ctx_lens=[64 + (i % 200) for i in range(300)], + num_q_heads=8, + num_kv_heads=1, + page_size=64, + policy=2, + seed=896, + ) + run_and_compare(case) + + +def test_cuda_graph_replay_matches_eager(): + _skip_unless_supported() + case = make_case( + q_lens=[1] * 8, + ctx_lens=[2048] * 8, + num_q_heads=8, + num_kv_heads=1, + page_size=64, + policy=2, + seed=897, + ) + kw = dict( + num_q_heads=8, + num_kv_heads=1, + qk_norm_policy=2, + q_norm_weight=case["q_norm_weight"], + k_norm_weight=case["k_norm_weight"], + ) + kc, vc, oq = ( + case["key_cache"].clone(), + case["value_cache"].clone(), + case["out_q"].clone(), + ) + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + kc, + vc, + out_q=oq, + **kw, + ) + torch.cuda.synchronize() + gkc, gvc, goq = ( + case["key_cache"].clone(), + case["value_cache"].clone(), + case["out_q"].clone(), + ) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + for _ in range( + 2 + ): # warm the JIT module / allocator on the side stream before capture + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + gkc, + gvc, + out_q=goq, + **kw, + ) + gkc.fill_(POISON) + gvc.fill_(POISON) + goq.fill_(POISON) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + gkc, + gvc, + out_q=goq, + **kw, + ) + torch.cuda.synchronize() + graph.replay() + torch.cuda.synchronize() + assert torch.equal(goq.view(torch.int16), oq.view(torch.int16)) + assert torch.equal(gkc.view(torch.int16), kc.view(torch.int16)) + assert torch.equal(gvc.view(torch.int16), vc.view(torch.int16)) + + +def test_rejects_bad_arguments(): + _skip_unless_supported() + case = make_case( + q_lens=[1], + ctx_lens=[5], + num_q_heads=8, + num_kv_heads=1, + page_size=64, + policy=0, + seed=1, + ) + with pytest.raises(ValueError): + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + case["key_cache"], + case["value_cache"], + num_q_heads=4, + num_kv_heads=1, + ) + with pytest.raises(ValueError): + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + case["key_cache"], + case["value_cache"], + num_q_heads=8, + num_kv_heads=1, + qk_norm_policy=2, + ) + with pytest.raises(ValueError): + cake_fused_qk_rmsnorm_rope_append_paged_kv_cache( + case["qkv"], + case["cos_sin"], + case["seq_lens"], + case["q_indptr"], + case["page_indices"], + case["key_cache"], + case["value_cache"], + num_q_heads=8, + num_kv_heads=1, + out_q=torch.empty((1, 8, 64), dtype=torch.bfloat16, device="cuda"), + )