diff --git a/miles_ext_kernels/cake/_impl/gqa/__init__.py b/miles_ext_kernels/cake/_impl/gqa/__init__.py new file mode 100644 index 0000000..20d8694 --- /dev/null +++ b/miles_ext_kernels/cake/_impl/gqa/__init__.py @@ -0,0 +1,18 @@ +""" +Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +Licensed under the Apache License, Version 2.0. +https://www.apache.org/licenses/LICENSE-2.0 + +The implementation of ``miles_ext_kernels.cake.gqa`` on its vendored package ``cake_gqa_train`` (:mod:`.ops`), +imported on first use of an entry point. Importing it imports torch. +""" + +from miles_ext_kernels.cake._runtime import loader + +loader.install_op("gqa") # the runtime loader on the op's package, before it can build + +from miles_ext_kernels.cake.gqa import check_supported, is_supported # noqa: E402 + +from .ops import prebuild, release, varlen_gqa_attention # noqa: E402 + +__all__ = ["check_supported", "is_supported", "prebuild", "release", "varlen_gqa_attention"] diff --git a/miles_ext_kernels/cake/_impl/gqa/ops.py b/miles_ext_kernels/cake/_impl/gqa/ops.py new file mode 100644 index 0000000..e9fd957 --- /dev/null +++ b/miles_ext_kernels/cake/_impl/gqa/ops.py @@ -0,0 +1,198 @@ +""" +Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +Licensed under the Apache License, Version 2.0. +https://www.apache.org/licenses/LICENSE-2.0 + +Packed causal GQA attention, head dimension 256, forward and backward on the vendored ``cake_gqa_train`` package +(FlashInfer's ``flashinfer.gqa_train`` backend), SM100a / SM103a / SM107a. + +Semantics. ``q [T_q, H_q, 256]``, ``k`` / ``v [T_k, H_kv, 256]`` (BF16, ``H_q / H_kv`` even), segments packed by +int32 device ``cu_seqlens_q`` / ``cu_seqlens_k`` (``cu[0] = 0``, non-decreasing, ``cu_q[B] = T_q``, +``cu_k[B] = T_k``). Query row ``i`` of a segment with ``q_len`` rows and ``k_len`` keys sees the keys +``j <= i + k_len - q_len`` (FlashAttention's bottom-right alignment); a row that sees no key gets a zero output and a +zero gradient. + +``max_seqlen_k``. The backward main kernel runs one CTA pair per 128 keys of ``max_seqlen_k``, per key/value head and +segment, and a pair past its segment's last key block retires at once. A value below the longest segment's key count +leaves key blocks out: their dK / dV are never written and their dQ terms are missing, without an error. A larger value +is correct but launches pairs that only retire, which slows the backward down the more segments a call packs (``T_k`` +is always correct, and the slowest). The kernels never read the boundaries on the host, so only the caller keeps this +contract; the option ``cake_gqa_check_seqlens`` (``--cake-gqa-check-seqlens``) checks it at every call, at the price +of one device synchronization. + +Numerics. FP32 accumulation throughout, P and dS rounded to BF16 only as MMA operands. The output, the LSE and the BF16 +dK / dV repeat bit for bit run to run, whatever valid ``max_seqlen_k`` is given; dQ is summed in FP32 with atomic adds, +one partial per 128-key block in a scheduling order, and rounded to BF16 once, so it is not reproducible bit for bit. + +Memory. Each backward takes its scratch (FP32 ``[H_q / 2, rows_pad, 256]`` and two ``[H_q / 2, rows_pad]``, +``rows_pad = round_up(2 T_q + 128 B, 128)``) from the runtime's workspace pool on the current stream, and each forward +zeroes a fresh ``uint32`` tile counter; autograd saves ``(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)`` and scalars +only. Every launch is a plain launch on the current stream, and nothing synchronizes the device. +""" + +from __future__ import annotations + +import math +from pathlib import Path +from typing import Any + +import torch + +from miles_ext_kernels.cake._runtime import config, loader, workspace +from miles_ext_kernels.cake._vendor.cake_gqa_train import cake_backend, cake_jit + +__all__ = [ + "HEAD_DIM", + "PACKAGE", + "STAGES", + "PoolWorkspace", + "check_seqlens", + "libraries", + "prebuild", + "release", + "varlen_gqa_attention", +] + +OP = "gqa" +PACKAGE = "cake_gqa_train" +HEAD_DIM = cake_backend.HEAD_DIM +# The stages one forward and backward of varlen_gqa_attention launch (BF16 dK / dV), in launch order. +STAGES = ("fwd", "bwd_pre", "bwd_bf16", "bwd_post") + + +def libraries(arch: str) -> list[str]: + """The ``build.json`` libraries of :data:`STAGES` on ``arch``.""" + return [cake_backend.library_name(cake_backend.stage_record(stage, arch), arch) for stage in STAGES] + + +def prebuild(*, device: Any = None, max_workers: int | None = None) -> list[Path]: + """Build the libraries of :data:`STAGES` for the architecture of ``device`` (default: the current one). + + No model parameter: every head layout, length and option runs these four libraries. + """ + return loader.prebuild_libraries(PACKAGE, libraries(loader.device_arch(device)), max_workers=max_workers) + + +def release() -> None: + """Nothing to drop here: the backward scratch lives in the runtime's workspace pool, the tile counter per call.""" + + +class PoolWorkspace: + """What the vendored launches ask their ``workspace`` for, holding nothing between calls. + + :meth:`tile_counter` is the forward's zeroed ``uint32`` ticket counter, new at every call (the kernel leaves it + zero). :meth:`scratch` is the backward's ``(dq_accum, lse_log2, dpsum)``: views of one buffer of the workspace + pool, key ``("gqa", "bwd_scratch")``, on the current stream. + """ + + def __init__(self, device: torch.device) -> None: + self.device = device + + def tile_counter(self) -> torch.Tensor: + return torch.zeros(1, dtype=torch.uint32, device=self.device) + + def scratch(self, total_q: int, num_heads_q: int, num_segments: int) -> tuple[torch.Tensor, ...]: + nbytes = cake_jit.scratch_bytes(total_q, num_heads_q, num_segments) + buffer = workspace.take(OP, "bwd_scratch", (nbytes,), torch.uint8, self.device) + return tuple( + buffer[offset : offset + 4 * math.prod(shape)].view(torch.float32).view(shape) + for offset, shape in cake_jit.scratch_layout(total_q, num_heads_q, num_segments).values() + ) + + +def check_seqlens( + cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, total_q: int, total_k: int, max_seqlen_k: int +) -> None: + """``ValueError`` unless the call keeps the contract the kernels cannot check. + + Each boundary vector starts at 0, never decreases and ends at its token count, and no segment has more than + ``max_seqlen_k`` keys. Reads both vectors on the host: one device synchronization. + """ + split = cu_seqlens_q.numel() + bounds = torch.cat((cu_seqlens_q.reshape(-1), cu_seqlens_k.reshape(-1))).tolist() + for name, cu, total in (("cu_seqlens_q", bounds[:split], total_q), ("cu_seqlens_k", bounds[split:], total_k)): + if not cu or cu[0] != 0 or cu[-1] != total or any(b < a for a, b in zip(cu, cu[1:])): + shown = f"{cu[:8]}{' ...' if len(cu) > 8 else ''}" + raise ValueError(f"{name} must start at 0, never decrease and end at {total}; got {shown}") + keys = [b - a for a, b in zip(bounds[split:], bounds[split + 1 :])] + longest = max(keys, default=0) + if longest > max_seqlen_k: + raise ValueError( + f"max_seqlen_k is {max_seqlen_k}, but segment {keys.index(longest)} has {longest} keys: the backward " + "would leave key blocks out and give wrong gradients" + ) + + +class _VarlenGQAAttention(torch.autograd.Function): + @staticmethod + def forward( # type: ignore[override] + ctx: Any, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + softmax_scale: float, + max_seqlen_k: int, + ) -> torch.Tensor: + out, lse = cake_backend.gqa_attention_fwd( + q, k, v, cu_seqlens_q, cu_seqlens_k, softmax_scale=softmax_scale, workspace=PoolWorkspace(q.device) + ) + ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k) + ctx.softmax_scale = softmax_scale + ctx.max_seqlen_k = max_seqlen_k + return out + + @staticmethod + def backward(ctx: Any, dout: torch.Tensor) -> tuple[torch.Tensor | None, ...]: # type: ignore[override] + q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + dq, dk, dv = cake_backend.gqa_attention_bwd( + q, + k, + v, + out, + lse, + dout.contiguous(), + cu_seqlens_q, + cu_seqlens_k, + softmax_scale=ctx.softmax_scale, + max_seqlen_k=ctx.max_seqlen_k, + workspace=PoolWorkspace(q.device), + ) + return dq, dk, dv, None, None, None, None + + +def varlen_gqa_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + *, + max_seqlen_k: int, + softmax_scale: float | None = None, +) -> torch.Tensor: + """Packed causal GQA attention with autograd: BF16 ``out [T_q, H_q, 256]``; ``q``, ``k`` and ``v`` get BF16 + gradients. + + ``cu_seqlens_q`` / ``cu_seqlens_k``: int32 ``[B + 1]`` on the device, ``cu[0] = 0``, non-decreasing, + ``cu_q[B] = T_q``, ``cu_k[B] = T_k``. ``max_seqlen_k`` (required): a host int at least every segment's key + count, as FlashAttention's; a smaller value gives wrong gradients without an error (see the module documentation). + ``softmax_scale``: default ``256 ** -0.5``. + + Raises ``ValueError`` for operands the kernels do not take (dtype, head dim, an odd query group, strides, the + boundary vectors) and in a process configured for deterministic mode (there is no ``det`` variant), + ``TypeError`` when ``max_seqlen_k`` is not an int, ``RuntimeError`` on another GPU. With the option + ``cake_gqa_check_seqlens`` it also checks the boundaries and ``max_seqlen_k`` on the host (:func:`check_seqlens`). + """ + config.variant(OP) # "fast", the only variant; raises in a process configured for deterministic mode + if isinstance(max_seqlen_k, bool) or not isinstance(max_seqlen_k, int): + raise TypeError( + f"max_seqlen_k must be a host int, at least every segment's key count; got {type(max_seqlen_k).__name__}" + ) + if max_seqlen_k < 0: + raise ValueError(f"max_seqlen_k must be >= 0, got {max_seqlen_k}") + if config.option(OP, "cake_gqa_check_seqlens"): + check_seqlens(cu_seqlens_q, cu_seqlens_k, q.shape[0], k.shape[0], max_seqlen_k) + scale = HEAD_DIM**-0.5 if softmax_scale is None else float(softmax_scale) + return _VarlenGQAAttention.apply(q, k, v, cu_seqlens_q, cu_seqlens_k, scale, max_seqlen_k) diff --git a/miles_ext_kernels/cake/_runtime/registry.py b/miles_ext_kernels/cake/_runtime/registry.py index 705bbcb..06f2156 100644 --- a/miles_ext_kernels/cake/_runtime/registry.py +++ b/miles_ext_kernels/cake/_runtime/registry.py @@ -31,6 +31,7 @@ "rmsnorm", "dsa", "mhc", + "gqa", ) # miles_ext_kernels.cake: op modules are its submodules, implementations those of its _impl package diff --git a/miles_ext_kernels/cake/_vendor/VENDORED.json b/miles_ext_kernels/cake/_vendor/VENDORED.json index 3910baf..6e2cf05 100644 --- a/miles_ext_kernels/cake/_vendor/VENDORED.json +++ b/miles_ext_kernels/cake/_vendor/VENDORED.json @@ -967,6 +967,40 @@ } } }, + "gqa": { + "cake_gqa_train": { + "cake_revision": "7ab5de0c9d4e42a367ba3e9e218d98208f9bd8fa", + "files": { + ".gitignore": "e5fd70b13ee2e116ccd3094d1035544ee7f627ed", + "README.md": "d150ca19da18df0c2aeccd77668937c41486348e", + "__init__.py": "a8171e69c4ff0b4b32f5fa80b0a8aca1b0ff0524", + "build.json": "79acfb251bc55693fdad8f8c262d7819069d10e2", + "cake_backend.py": "3586d0322a850e7f75333635cd1ff9583c857871", + "cake_build.py": "00756137eb159e30a77907eda84fd245123d87ea", + "cake_jit.py": "982ebc930de678e9216d7e051ee9796064f89ae8", + "csrc/.clang-format": "47a38a93f2dbdd6d944f9dddc3465e3fb65aec29", + "csrc/cake_ffi_checks.h": "129288f9432e054dc51b77d93194fb0f2ec20d9c", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu": "005e8da7bf8bfc6c20469921b056257da8f1b4e4", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu": "d879f473402e7037a89573c971838e8f6f486252", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu": "263c17eb8fdaf6a5db12f53cff290b929c88b19b", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu": "83f14ab04f97d6a4c8f0f38b548935164e8c0c9c", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu": "3dfd95bee83febd5bb32f1aba95df1c549c876f7", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu": "11dc829a88a1cd6628c31bdc55b2ac4d00dd1e74", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu": "39577c16e0e11e96c9147c2850a0b8717c3180f4", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu": "d6900616e6eaf6118da78a61efd0001d4707ea86", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu": "8cb01566afeb72272e39fcf3e84ff7cc2d8aef90", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu": "00f1909bf8e29c6087769ba20ecf9477492b5fab", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu": "7c8ea54ebbd80e7d1305b279296c7027418ea546", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu": "574e593a9cabf04738c7a68dc30c8bd06cb379f3" + }, + "source": { + "commit": "c5db8d3ccd3fa3842812280f7f208718dc0b8610", + "kind": "flashinfer", + "path": "flashinfer/experimental/cake_gqa_train", + "repo": "https://github.com/flashinfer-ai/flashinfer" + } + } + }, "kda": { "cake_kda_det_train": { "cake_revision": "3462767011d1e0de9c84de41bbb90397fa1ebafb", diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/.gitignore b/miles_ext_kernels/cake/_vendor/cake_gqa_train/.gitignore new file mode 100644 index 0000000..e5fd70b --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/.gitignore @@ -0,0 +1,2 @@ +# Prebuilt libraries and their build_info.json sidecar (cake_build.py); built locally, never tracked. +aot/ diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/README.md b/miles_ext_kernels/cake/_vendor/cake_gqa_train/README.md new file mode 100644 index 0000000..d150ca1 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/README.md @@ -0,0 +1,91 @@ +# Cake packed-varlen GQA training attention backend (head dim 256, SM100 / SM103 / SM107) + +Experimental backend for the forward and first-order backward of causal GQA +attention over packed variable-length segments, head dimension 256: BF16 +`q [T_q, H_q, 256]`, `k`, `v [T_k, H_kv, 256]` (`H_q / H_kv` even), segments +packed by `cu_seqlens_q` / `cu_seqlens_k` (int32 on the device, never read on +the host). Query row `i` of a segment with `q_len` rows and `k_len` keys sees +the keys `j <= i + k_len - q_len` (FlashAttention's bottom-right alignment); +a row that sees no key gets `out = 0`, `lse = -inf` and zero gradients. +`T_q`, `T_k`, the segment count and every length are run-time values: no +shape list, no recompilation, no device-to-host read, no device-wide +synchronization. + +Both the API and the backend are experimental: no compatibility guarantee. +SM100 (B200 / GB200), SM103 (B300 / GB300) and SM107 (R200) are the +supported architectures; another device raises `RuntimeError`. The feature does not +participate in automatic backend selection, autotuning or trace-apply. Its +libraries are loaded prebuilt from `aot/` when a builder compiled them and +JIT-built by FlashInfer from the sources and flags `build.json` records +otherwise. The package reads no environment variable. + +## Opt-in + +```python +from flashinfer.gqa_train import varlen_gqa_attention, gqa_train_workspace + +ws = gqa_train_workspace(q.device) # optional; reused across calls and shapes +out = varlen_gqa_attention( + q, k, v, cu_seqlens_q, cu_seqlens_k, + softmax_scale=0.0625, # default 256 ** -0.5 + max_seqlen_k=max_k, # optional host bound of the longest key segment + workspace=ws, +) +out.backward(dout) # BF16 dq, dk, dv + +# FP32 dK / dV added into caller-owned FP32 destinations [N_dst, H_kv, 256] +# through an int32 row map [T_k] (repeated rows accumulate); the callable runs +# in backward, so nothing is held between the forward and the backward: +out = varlen_gqa_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, + dkv_acc=lambda: (dk_dst, dv_dst, row_map)) +# FP32 window gradients [T_k, H_kv, 256] stored into caller-owned tensors: +out = varlen_gqa_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, + dkv_acc=lambda: (dk_win, dv_win)) +``` + +`varlen_gqa_attention_fwd` / `varlen_gqa_attention_bwd` are the two halves +without autograd (`dkv_dtype=torch.float32` returns FP32 window gradients; +`dk_acc`, `dv_acc`, `row_map` add into destinations). `max_seqlen_k` (a +host int, never read from the device) must be at least every segment's key +count, as FlashAttention's `max_seqlen_k`: a smaller value leaves key blocks +out. It sizes the backward grid (one CTA pair per 128 keys of the longest +segment, per key head and segment); without it the bound is `T_k`, which is +always correct but launches CTAs that retire at once: on a 12-segment layout of +257K keys the backward call took 8 to 15 % longer without it on SM100, SM103 +and SM107 GPUs. The backend package +exposes the same entry points (`varlen_gqa_attention`, `gqa_attention_fwd`, +`gqa_attention_bwd`, `Workspace`). + +## Kernels + +| stage | kernel | +|---|---| +| `fwd` | O and the natural-log LSE `[H_q, T_q]`: a CTA pair per 256 query rows of one (segment, query head), persistent tile tickets from a self-resetting counter | +| `bwd_pre` | `dPsum = sum(O * dO)`, `LSE * log2(e)` and the zeroed FP32 dQ accumulator, two query heads per scratch head | +| `bwd_fp32` / `bwd_rowmap` / `bwd_bf16` | the 5-GEMM backward: a CTA pair per 128 keys of one (segment, key head), looping the key head's query heads two at a time, dK / dV summed in FP32 in tensor memory; FP32 window stores, FP32 row-map adds or BF16 stores | +| `bwd_post` | `dQ = bf16(acc * softmax_scale)` | + +Numerics: FP32 accumulation for the scores, `dP`, dK, dV, the dQ partials and +the row sums; P and dS are rounded to BF16 only as MMA operands. dQ and the +row-map adds use atomic reductions, so they are not bitwise reproducible run +to run. + +## Checkpointing and memory + +Autograd saves `(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)` through +`ctx.save_for_backward` only; `ctx` holds scalars, the `dkv_acc` callable and +a weak reference to the workspace. Under +`torch.utils.checkpoint(..., use_reentrant=False)` a forward retains nothing. +The backward scratch is FP32 `[H_q / 2, rows_pad, 256]` plus two +`[H_q / 2, rows_pad]` with `rows_pad = round_up(2 T_q + 128 B, 128)` +(`gqa_train_workspace_bytes`); a workspace keeps the largest one seen +(`.nbytes`, `.release()`). One workspace serves one stream at a time; +without one, each device stream gets its own default workspace. + +## Limitations + +- Head dimension 256, BF16 inputs, causal only (bottom-right aligned), no + dropout, no sliding window. +- `H_q / H_kv` must be even. +- `dout` of the low-level backward must be contiguous (autograd makes it so). +- `sm_100a`, `sm_103a` and `sm_107a` only. diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/__init__.py b/miles_ext_kernels/cake/_vendor/cake_gqa_train/__init__.py new file mode 100644 index 0000000..a8171e6 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/__init__.py @@ -0,0 +1,51 @@ +""" +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. +""" + +# Experimental Cake backend: packed-varlen causal GQA attention training with +# head dimension 256 on SM100 / SM103 / SM107. Q ``[T_q, H_q, 256]``, K / V +# ``[T_k, H_kv, 256]`` (BF16, ``H_q / H_kv`` even), segments packed by +# ``cu_seqlens_q`` / ``cu_seqlens_k`` (int32 on the device), causal with the +# bottom-right alignment per segment. The entry point with autograd is +# ``varlen_gqa_attention``; ``gqa_attention_fwd`` / ``gqa_attention_bwd`` are +# the two halves. The registry and the host plan live in ``cake_jit``, the +# loader, launches and autograd in ``cake_backend``. Every switch is an +# explicit keyword argument; the package reads no environment variable. + +from .cake_backend import ( + ARCHES, + HEAD_DIM, + DkvAccSource, + VarlenGQAAttentionFunction, + Workspace, + default_workspace, + gqa_attention_bwd, + gqa_attention_fwd, + varlen_gqa_attention, + workspace_bytes, +) + +__all__ = [ + "ARCHES", + "HEAD_DIM", + "DkvAccSource", + "VarlenGQAAttentionFunction", + "Workspace", + "default_workspace", + "gqa_attention_bwd", + "gqa_attention_fwd", + "varlen_gqa_attention", + "workspace_bytes", +] diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/build.json b/miles_ext_kernels/cake/_vendor/cake_gqa_train/build.json new file mode 100644 index 0000000..79acfb2 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/build.json @@ -0,0 +1,432 @@ +{ + "arch_flags": { + "sm_100a": [ + "-gencode=arch=compute_100a,code=sm_100a" + ], + "sm_103a": [ + "-gencode=arch=compute_103a,code=sm_103a" + ], + "sm_107a": [ + "-gencode=arch=compute_107a,code=sm_107a" + ] + }, + "common": { + "include_dirs": [ + "csrc" + ], + "link": { + "cudart": "shared", + "host_linker": "c++", + "tvm_ffi_symbols": "resolved at load; do not pass --no-undefined" + }, + "link_flags": [ + "-shared", + "-lcudart", + "-lcuda" + ], + "nvcc_flags": [ + "-std=c++17", + "-O3", + "-DNDEBUG", + "--compiler-options=-fPIC", + "--expt-relaxed-constexpr", + "-Xfatbin=-compress-all", + "--compress-mode=size" + ], + "nvcc_flags_if": [ + { + "cuda": ">=12.8", + "flags": [ + "-static-global-template-stub=false" + ] + } + ], + "requires": { + "cuda_toolkit": { + "min": "12.8" + }, + "tvm_ffi": { + "include_dirs": [ + "tvm_ffi.libinfo.find_include_path()", + "tvm_ffi.libinfo.find_dlpack_include_path()" + ], + "python": "apache-tvm-ffi>=0.1.11,<0.2" + } + } + }, + "generator": { + "cake_revision": "7ab5de0c9d4e42a367ba3e9e218d98208f9bd8fa", + "cuda": "13.4.92", + "rendered_at": "2026-10-30T00:45:55Z" + }, + "libraries": [ + { + "arch": "sm_100a", + "closure_sha256": "090299fe2ee0bee51aa557b31d9a3dbc29ff428f2e1fe22f0d35960af2ad9fff", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_gqa_train_1afdfdde603ff8806b4d_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_1afdfdde603ff8806b4d_sm_100a.so", + "sha256": { + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu": "faf82b2e096f9f56c5a6c2d0a1b55593858186816b7625ac6ad7833012091a60", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu": "a48e190118d5e764728551e302a98da93e0934968a82bbe7564952b2d02a6fbb" + }, + "sources": [ + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "b4b2f1d8e2efead58ccd95422e5ba035d04e29e7290bab16a061ff89864a6764", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_gqa_train_1afdfdde603ff8806b4d_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_1afdfdde603ff8806b4d_sm_103a.so", + "sha256": { + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu": "faf82b2e096f9f56c5a6c2d0a1b55593858186816b7625ac6ad7833012091a60", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu": "a48e190118d5e764728551e302a98da93e0934968a82bbe7564952b2d02a6fbb" + }, + "sources": [ + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "d65e823c1d5ee8ae415d87d6be23cdf4ea970f2b88d5f3e8476276cc16024aae", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gqa_train_1afdfdde603ff8806b4d_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_1afdfdde603ff8806b4d_sm_107a.so", + "sha256": { + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu": "faf82b2e096f9f56c5a6c2d0a1b55593858186816b7625ac6ad7833012091a60", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu": "a48e190118d5e764728551e302a98da93e0934968a82bbe7564952b2d02a6fbb" + }, + "sources": [ + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu", + "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "6ec9777f60274328bd67011c3decad29c55c91ebb73f35e3cd039d01f673e3e0", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_gqa_train_37f4cd45025705f3cfc8_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_37f4cd45025705f3cfc8_sm_100a.so", + "sha256": { + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu": "938513d81f53ed8f5f374da22de6522300cf86632203c9526428c6c02c699bd5", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu": "9a1e932d14d7c583fd07a30e93af2333dc179ef81a4f962b1ae892f05d0c5708" + }, + "sources": [ + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "22a502a0d0a252bf66a7738f0b009293082dc697680c7905f682c7eac0ffad7a", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_gqa_train_37f4cd45025705f3cfc8_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_37f4cd45025705f3cfc8_sm_103a.so", + "sha256": { + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu": "938513d81f53ed8f5f374da22de6522300cf86632203c9526428c6c02c699bd5", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu": "9a1e932d14d7c583fd07a30e93af2333dc179ef81a4f962b1ae892f05d0c5708" + }, + "sources": [ + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "e6aa357b95d786677bb61c2b6df1b54f9185a463ae6b9471727d6176bdeb82b7", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gqa_train_37f4cd45025705f3cfc8_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_37f4cd45025705f3cfc8_sm_107a.so", + "sha256": { + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu": "938513d81f53ed8f5f374da22de6522300cf86632203c9526428c6c02c699bd5", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu": "9a1e932d14d7c583fd07a30e93af2333dc179ef81a4f962b1ae892f05d0c5708" + }, + "sources": [ + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu", + "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "d3b0836b0abe352bc8e1335d97c45a613b87f46f7670b610d29a39ddbd51d25e", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_gqa_train_493c4a28dafa5fc6b2f1_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_493c4a28dafa5fc6b2f1_sm_100a.so", + "sha256": { + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu": "450aea69af6731ddf7a16b81cc56288254c950263f2549f586364ec7e486db13", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu": "0daf7e3a247817ce4382266bf37d1da4758b903e13189054eae014ddb0e7e052" + }, + "sources": [ + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "4af77e78d4458d48236e6c5fc0b772204547e4a3e18120c093c7afd795a804e0", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_gqa_train_493c4a28dafa5fc6b2f1_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_493c4a28dafa5fc6b2f1_sm_103a.so", + "sha256": { + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu": "450aea69af6731ddf7a16b81cc56288254c950263f2549f586364ec7e486db13", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu": "0daf7e3a247817ce4382266bf37d1da4758b903e13189054eae014ddb0e7e052" + }, + "sources": [ + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "20bb57809d27be25376da4f783bba2540e5f968137218840ceeae20364578eea", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gqa_train_493c4a28dafa5fc6b2f1_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_493c4a28dafa5fc6b2f1_sm_107a.so", + "sha256": { + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu": "450aea69af6731ddf7a16b81cc56288254c950263f2549f586364ec7e486db13", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu": "0daf7e3a247817ce4382266bf37d1da4758b903e13189054eae014ddb0e7e052" + }, + "sources": [ + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu", + "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "773aab9c73e3caf4989dd1194ae56893e7d54e96567fcae675802f980f56eae7", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_gqa_train_6fa13db592acdb66a6c9_sm_100a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_6fa13db592acdb66a6c9_sm_100a.so", + "sha256": { + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu": "5af55cf4d04da534e1ff43c68248399c24ddfb33c1c4a200f42dc60242545d91", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu": "501a7a3ead1fec78f58047e1f64a27c5ecd6849266f640098701e36e0a48e227" + }, + "sources": [ + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "116f17c6c5860f7ea1c3ec417409a17dd8875a328fc3d3fa25052023c0a4542d", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_gqa_train_6fa13db592acdb66a6c9_sm_103a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_6fa13db592acdb66a6c9_sm_103a.so", + "sha256": { + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu": "5af55cf4d04da534e1ff43c68248399c24ddfb33c1c4a200f42dc60242545d91", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu": "501a7a3ead1fec78f58047e1f64a27c5ecd6849266f640098701e36e0a48e227" + }, + "sources": [ + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "24a8cabdd2d154b920a17806e321a30793122e68be856b7072a198057e41003a", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gqa_train_6fa13db592acdb66a6c9_sm_107a", + "nvcc_flags": [ + "--use_fast_math" + ], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_6fa13db592acdb66a6c9_sm_107a.so", + "sha256": { + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu": "5af55cf4d04da534e1ff43c68248399c24ddfb33c1c4a200f42dc60242545d91", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu": "501a7a3ead1fec78f58047e1f64a27c5ecd6849266f640098701e36e0a48e227" + }, + "sources": [ + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu", + "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "a6651874727543a5e6d4e55849dd1da26e3f9d3029b22ffa9c25e756e06c998a", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_gqa_train_76decc9a5d22dc91ddd1_sm_100a", + "nvcc_flags": [], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_76decc9a5d22dc91ddd1_sm_100a.so", + "sha256": { + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu": "b87c3ec594f0165348580d774642269623554e173d135a99d25588214cd7b44b", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu": "7b08add9919582c68bfd1ef6cf2f9ac8f16cdd1b742010543505c456f85b8c68" + }, + "sources": [ + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "3510317840b45d813a2da8e39718f5074f85acb8fd430609511325f2a772268c", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_gqa_train_76decc9a5d22dc91ddd1_sm_103a", + "nvcc_flags": [], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_76decc9a5d22dc91ddd1_sm_103a.so", + "sha256": { + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu": "b87c3ec594f0165348580d774642269623554e173d135a99d25588214cd7b44b", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu": "7b08add9919582c68bfd1ef6cf2f9ac8f16cdd1b742010543505c456f85b8c68" + }, + "sources": [ + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "734eb2baeb406d58e54fc5a2431714f4c5d23dad47d8b3b413100cf68416df05", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gqa_train_76decc9a5d22dc91ddd1_sm_107a", + "nvcc_flags": [], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_76decc9a5d22dc91ddd1_sm_107a.so", + "sha256": { + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu": "b87c3ec594f0165348580d774642269623554e173d135a99d25588214cd7b44b", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu": "7b08add9919582c68bfd1ef6cf2f9ac8f16cdd1b742010543505c456f85b8c68" + }, + "sources": [ + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu", + "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu" + ] + }, + { + "arch": "sm_100a", + "closure_sha256": "6b79da0f28db8858a38ccc72810b2620c1db0a2aaf63b26d78515884a8481f57", + "entry": "run", + "gencode": "-gencode=arch=compute_100a,code=sm_100a", + "min_cuda": "12.8", + "name": "cake_gqa_train_a91752e8a2732bf2e246_sm_100a", + "nvcc_flags": [], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_a91752e8a2732bf2e246_sm_100a.so", + "sha256": { + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu": "99ba60250780c56bfadeb10eef179ea057fce601bfb1bc070b2c397bb18407cb", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu": "608b0b3556801230711a639a02b5bf9ab590e985eb328bf53b5f05ef5871c29d" + }, + "sources": [ + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu" + ] + }, + { + "arch": "sm_103a", + "closure_sha256": "4991ab01491c48fc1a41fef6779e5c0f99dbb1ed90a323fef03a226277ae59ed", + "entry": "run", + "gencode": "-gencode=arch=compute_103a,code=sm_103a", + "min_cuda": "12.9", + "name": "cake_gqa_train_a91752e8a2732bf2e246_sm_103a", + "nvcc_flags": [], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_a91752e8a2732bf2e246_sm_103a.so", + "sha256": { + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu": "99ba60250780c56bfadeb10eef179ea057fce601bfb1bc070b2c397bb18407cb", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu": "608b0b3556801230711a639a02b5bf9ab590e985eb328bf53b5f05ef5871c29d" + }, + "sources": [ + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu" + ] + }, + { + "arch": "sm_107a", + "closure_sha256": "1301cba0673dd92847b36de1b9e03b3e6375b4cf23e12e4d12cf2d2861d22547", + "entry": "run", + "gencode": "-gencode=arch=compute_107a,code=sm_107a", + "min_cuda": "13.4", + "name": "cake_gqa_train_a91752e8a2732bf2e246_sm_107a", + "nvcc_flags": [], + "nvcc_flags_if": [], + "path": "aot/cake_gqa_train_a91752e8a2732bf2e246_sm_107a.so", + "sha256": { + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu": "99ba60250780c56bfadeb10eef179ea057fce601bfb1bc070b2c397bb18407cb", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu": "608b0b3556801230711a639a02b5bf9ab590e985eb328bf53b5f05ef5871c29d" + }, + "sources": [ + "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu", + "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu" + ] + } + ], + "package": "flashinfer.experimental.cake_gqa_train", + "schema": "cake-build/1" +} diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_backend.py b/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_backend.py new file mode 100644 index 0000000..3586d03 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_backend.py @@ -0,0 +1,742 @@ +""" +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. + +Loader, launches and autograd of the Cake packed-varlen GQA training +attention (head dimension 256, SM100 / SM103 / SM107). + +The registry (:mod:`.cake_jit`) is data plus the host plan: ``MODULES`` holds +one record per program (its ``library`` stem; the ``build.json`` entries are +``_``), :func:`.cake_jit.grid` the launch grids and +:func:`.cake_jit.scratch_layout` the backward scratch. :mod:`.cake_build` +serves a library from a prebuilt ``aot/`` file when a builder compiled it and +otherwise through FlashInfer's JIT from the sources and flags of +``build.json``; the prebuilt route imports nothing from ``flashinfer``. + +Every launch is a plain stream-ordered launch on the current torch stream: +nothing cooperative, no PDL, no co-residency requirement. Nothing here reads +an environment variable, reads device memory or synchronizes; grids and +scratch sizes come from tensor shapes, ``len(cu_seqlens) - 1`` and the SM +count. The segment boundaries ``cu_seqlens_q`` / ``cu_seqlens_k`` (int32 on +the device, ``cu[0] == 0``, non-decreasing, ``cu_q[-1] == T_q``, +``cu_k[-1] == T_k``) are read by the kernels only. +""" + +from __future__ import annotations + +import functools +import math +import weakref +from collections.abc import Callable, Mapping +from typing import Any, Optional, Union + +import torch +import tvm_ffi + +from . import cake_build, cake_jit + +__all__ = [ + "ARCHES", + "ARCH_BY_CAPABILITY", + "HEAD_DIM", + "DkvAccSource", + "VarlenGQAAttentionFunction", + "Workspace", + "arch_for_capability", + "default_workspace", + "device_arch", + "gqa_attention_bwd", + "gqa_attention_fwd", + "is_prebuilt", + "library_name", + "load", + "parameters", + "record", + "run_record", + "sm_count", + "specs_for_arch", + "stage_record", + "varlen_gqa_attention", + "workspace_bytes", +] + +# One verified build per architecture: B200 / GB200 (``sm_100a``), B300 / +# GB300 (``sm_103a``) and R200 (``sm_107a``), each compiled with its exact +# ``-gencode`` flag set. +ARCHES: tuple[str, ...] = tuple(cake_jit.ARCHES) +ARCH_BY_CAPABILITY: dict[tuple[int, int], str] = dict(cake_jit.ARCH_BY_CAPABILITY) +HEAD_DIM: int = cake_jit.HEAD_DIM +LOG2_E = math.log2(math.e) + +# The three dK / dV outputs of the backward main kernel, by stage. +_BWD_STAGE = {"bf16": "bwd_bf16", "fp32": "bwd_fp32", "rowmap": "bwd_rowmap"} + + +# ---------------------------------------------------------------------------------------------------------------- +# Devices + + +def arch_for_capability(capability: tuple[int, int]) -> Optional[str]: + """The exported architecture of one device capability, if any.""" + values = tuple(int(value) for value in capability) + if len(values) != 2: + return None + return ARCH_BY_CAPABILITY.get((values[0], values[1])) + + +@functools.cache +def device_arch(device_index: int) -> str: + """Exported architecture of one CUDA device (``RuntimeError`` when the device is not an exported target); cached + per device index.""" + capability = tuple(torch.cuda.get_device_capability(device_index)) + arch = arch_for_capability(capability) + if arch is None: + raise RuntimeError( + f"device {device_index} (capability {capability}) is not an exported Cake GQA training attention " + f"target; the exported architectures are {ARCHES}" + ) + return arch + + +@functools.cache +def sm_count(device_index: int) -> int: + """Streaming multiprocessor count of one CUDA device (immutable per process).""" + return int(torch.cuda.get_device_properties(device_index).multi_processor_count) + + +# ---------------------------------------------------------------------------------------------------------------- +# Records and libraries (``build.json``) + + +def record(name: str) -> dict[str, Any]: + """The registry record ``name`` (a key of ``MODULES``).""" + found = cake_jit.MODULES.get(name) + if found is None: + raise KeyError( + f"{name!r} is not a record of the Cake GQA training attention registry" + ) + return found + + +@functools.cache +def stage_record(stage: str, arch: str) -> str: + """The record serving ``stage`` on ``arch`` (``KeyError`` when none does).""" + return cake_jit.stage_record(stage, arch) + + +def library_name(name: str, arch: str) -> str: + """``build.json`` library of record ``name`` built for ``arch``: ``_``.""" + found = record(name) + if arch not in found["arches"]: + raise RuntimeError( + f"Cake GQA training attention record {name} is not delivered for {arch}; its architectures: " + f"{found['arches']}" + ) + return f"{found['library']}_{arch}" + + +def is_prebuilt(name: str, arch: str) -> bool: + """Whether the library of record ``name`` for ``arch`` is served from ``aot/`` (a builder compiled it).""" + return cake_build.prebuilt_path(library_name(name, arch)).is_file() + + +@functools.cache +def load(name: str, arch: str) -> Any: + """The loaded tvm-ffi module of record ``name`` built for ``arch``: the prebuilt ``aot/`` library when a builder + compiled it, else FlashInfer's JIT build of the same sources and flags (``cake_build.load_library``).""" + return cake_build.load_library(library_name(name, arch)) + + +def specs_for_arch(arch: str) -> list[Any]: + """JIT specifications of every ``arch`` library of ``build.json`` without a prebuilt file (for ahead-of-time + warm-up through FlashInfer's ``build_jit_specs``).""" + return cake_build.jit_specs(arch) + + +@functools.cache +def parameters(name: str) -> tuple[str, ...]: + """The kernel parameter names of record ``name`` in ``run`` order (the argument plan without the grid).""" + names: list[str] = [] + axes: list[str] = [] + for kind, key in record(name)["arg_plan"]: + if kind == "grid": + axes.append(key) + else: + if axes: + raise ValueError( + f"record {name}: kernel parameter {key!r} follows the grid arguments" + ) + names.append(key) + if axes != ["grid_x", "grid_y", "grid_z"]: + raise ValueError( + f"record {name}: the argument plan must end with grid_x, grid_y, grid_z; got {axes}" + ) + return tuple(names) + + +def run_record( + name: str, + values: Mapping[str, Any], + *, + grid: tuple[int, int, int], + device_index: int, +) -> None: + """Launch record ``name`` with its kernel parameters named in ``values`` on ``grid``, on the current torch + stream of device ``device_index`` (the operands' device, made current for the launch).""" + args = [] + for key in parameters(name): + try: + args.append(values[key]) + except KeyError: + raise ValueError(f"record {name} needs argument {key!r}") from None + entry = getattr(load(name, device_arch(device_index)), record(name)["ffi_entry"]) + with torch.cuda.device(device_index), tvm_ffi.use_torch_stream(): + entry(*args, int(grid[0]), int(grid[1]), int(grid[2])) + + +def _launch( + stage: str, values: Mapping[str, Any], dims: Mapping[str, int], device_index: int +) -> None: + grid = cake_jit.grid(stage, sm_count=sm_count(device_index), **dims) + run_record( + stage_record(stage, device_arch(device_index)), + values, + grid=grid, + device_index=device_index, + ) + + +# ---------------------------------------------------------------------------------------------------------------- +# Workspace + + +def workspace_bytes(total_q: int, num_heads_q: int, num_segments: int) -> int: + """Bytes of the backward scratch one call of these sizes needs (FP32 ``dq_accum [H_q / 2, rows_pad, 256]``, + ``lse_log2`` and ``dpsum [H_q / 2, rows_pad]``, ``rows_pad = round_up(2 T_q + 128 B, 128)``).""" + return cake_jit.scratch_bytes(int(total_q), int(num_heads_q), int(num_segments)) + + +class Workspace: + """Device scratch of the GQA training attention, reused across calls and shapes. + + It holds the forward's tile counter (one ``uint32``, zeroed once; every launch leaves it zero) and the backward + scratch, grown to the largest shape seen and never shrunk. One workspace serves one stream at a time: two + launches running concurrently on different streams need two workspaces. Autograd keeps only a weak reference + to it, so a checkpointed forward retains nothing through it. + """ + + def __init__(self, device: Union[torch.device, str, int, None] = None) -> None: + if device is None: + resolved = torch.device("cuda", torch.cuda.current_device()) + elif isinstance(device, int): + resolved = torch.device("cuda", device) + else: + resolved = torch.device(device) + if resolved.type != "cuda": + raise ValueError(f"a workspace lives on a CUDA device, not {resolved}") + if resolved.index is None: + resolved = torch.device("cuda", torch.cuda.current_device()) + self.device = resolved + self._counter: Optional[torch.Tensor] = None + self._scratch: Optional[torch.Tensor] = None + self._views: Optional[tuple[tuple[int, int, int], tuple[torch.Tensor, ...]]] = ( + None + ) + + @property + def nbytes(self) -> int: + """Device bytes the workspace holds now.""" + held = 0 if self._scratch is None else self._scratch.numel() + return held + ( + 0 if self._counter is None else self._counter.untyped_storage().nbytes() + ) + + def tile_counter(self) -> torch.Tensor: + """The forward's ``uint32`` tile counter (zero between launches).""" + if self._counter is None: + self._counter = torch.zeros(1, dtype=torch.uint32, device=self.device) + return self._counter + + def scratch( + self, total_q: int, num_heads_q: int, num_segments: int + ) -> tuple[torch.Tensor, ...]: + """``(dq_accum, lse_log2, dpsum)`` views for one backward call, growing the scratch when it is too small.""" + key = (total_q, num_heads_q, num_segments) + if self._views is not None and self._views[0] == key: + return self._views[1] + need = cake_jit.scratch_bytes(total_q, num_heads_q, num_segments) + if self._scratch is None or self._scratch.numel() < need: + self._views = None + self._scratch = None + self._scratch = torch.empty(need, dtype=torch.uint8, device=self.device) + views = tuple( + self._scratch[offset : offset + 4 * math.prod(shape)] + .view(torch.float32) + .view(shape) + for offset, shape in cake_jit.scratch_layout( + total_q, num_heads_q, num_segments + ).values() + ) + # One entry: the views of the last shape (a forward, its replay and its backward share one shape). + self._views = (key, views) + return views + + def release(self) -> None: + """Drop the backward scratch (the tile counter stays).""" + self._views = None + self._scratch = None + + +_DEFAULT_WORKSPACES: dict[tuple[int, int], Workspace] = {} + + +def default_workspace(device_index: int) -> Workspace: + """The workspace of the device's current stream, made on first use and kept for the process.""" + stream = torch.cuda.current_stream(device_index).stream_id + found = _DEFAULT_WORKSPACES.get((device_index, stream)) + if found is None: + found = _DEFAULT_WORKSPACES[(device_index, stream)] = Workspace(device_index) + return found + + +# ---------------------------------------------------------------------------------------------------------------- +# Checks (shapes, dtypes, strides and devices only: nothing reads tensor contents) + + +def _check_operands( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, +) -> dict[str, int]: + for name, t in (("q", q), ("k", k), ("v", v)): + if t.dtype != torch.bfloat16 or t.dim() != 3 or t.shape[2] != HEAD_DIM: + raise ValueError( + f"{name} must be bfloat16 [T, H, {HEAD_DIM}], got {t.dtype} {tuple(t.shape)}" + ) + if not t.is_cuda or t.device != q.device: + raise ValueError( + f"{name} must be on q's CUDA device {q.device}, got {t.device}" + ) + if t.stride(2) != 1 or t.stride(0) % 8 or t.stride(1) % 8 or t.data_ptr() % 16: + raise ValueError( + f"{name} must have a unit head-dim stride and 16-byte aligned token / head strides and base, got " + f"strides {t.stride()}" + ) + if v.shape != k.shape: + raise ValueError( + f"v must have k's shape {tuple(k.shape)}, got {tuple(v.shape)}" + ) + total_q, heads_q = int(q.shape[0]), int(q.shape[1]) + total_k, heads_kv = int(k.shape[0]), int(k.shape[1]) + if heads_kv == 0 or heads_q % heads_kv or (heads_q // heads_kv) % 2: + raise ValueError( + f"the query heads ({heads_q}) must be an even multiple of the key heads ({heads_kv}): the backward packs " + "two query heads of one key head per tile" + ) + for name, cu in (("cu_seqlens_q", cu_seqlens_q), ("cu_seqlens_k", cu_seqlens_k)): + if ( + cu.dtype != torch.int32 + or cu.dim() != 1 + or cu.device != q.device + or not cu.is_contiguous() + ): + raise ValueError( + f"{name} must be a contiguous int32 vector on {q.device}, got {cu.dtype} {cu.device}" + ) + if cu_seqlens_q.numel() != cu_seqlens_k.numel() or cu_seqlens_q.numel() < 2: + raise ValueError( + "cu_seqlens_q and cu_seqlens_k must both hold num_segments + 1 >= 2 entries, got " + f"{cu_seqlens_q.numel()} and {cu_seqlens_k.numel()}" + ) + return { + "total_q": total_q, + "total_k": total_k, + "num_heads_q": heads_q, + "num_heads_kv": heads_kv, + "num_segments": int(cu_seqlens_q.numel()) - 1, + } + + +def _check_out( + name: str, + t: torch.Tensor, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, +) -> None: + if ( + tuple(t.shape) != shape + or t.dtype != dtype + or t.device != device + or not t.is_contiguous() + ): + raise ValueError( + f"{name} must be a contiguous {dtype} tensor of shape {shape} on {device}, got {t.dtype} " + f"{tuple(t.shape)} on {t.device}" + ) + + +def _scale(softmax_scale: Optional[float]) -> float: + return HEAD_DIM**-0.5 if softmax_scale is None else float(softmax_scale) + + +# ---------------------------------------------------------------------------------------------------------------- +# Forward and backward + + +def gqa_attention_fwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + *, + softmax_scale: Optional[float] = None, + out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + workspace: Optional[Workspace] = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Forward: ``(out [T_q, H_q, 256] bf16, lse [H_q, T_q] fp32)``, causal with bottom-right alignment per segment; + a row that sees no key gets ``out = 0`` and ``lse = -inf``.""" + dims = _check_operands(q, k, v, cu_seqlens_q, cu_seqlens_k) + total_q, heads_q = dims["total_q"], dims["num_heads_q"] + if out is None: + out = torch.empty( + total_q, heads_q, HEAD_DIM, dtype=torch.bfloat16, device=q.device + ) + else: + _check_out("out", out, (total_q, heads_q, HEAD_DIM), torch.bfloat16, q.device) + if lse is None: + lse = torch.empty(heads_q, total_q, dtype=torch.float32, device=q.device) + else: + _check_out("lse", lse, (heads_q, total_q), torch.float32, q.device) + if total_q == 0: + return out, lse + if dims["total_k"] == 0: + out.zero_() + lse.fill_(-math.inf) + return out, lse + device_index = q.device.index + ws = workspace if workspace is not None else default_workspace(device_index) + scale = _scale(softmax_scale) + values = { + "Q": q, + "K": k, + "V": v, + "O": out, + "LSE": lse, + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "tile_counter": ws.tile_counter(), + "num_segments": dims["num_segments"], + "num_heads_q": heads_q, + "group": heads_q // dims["num_heads_kv"], + "total_q": total_q, + "softmax_scale_log2": scale * LOG2_E, + } + _launch("fwd", values, dims, device_index) + return out, lse + + +def gqa_attention_bwd( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + out: torch.Tensor, + lse: torch.Tensor, + dout: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + *, + softmax_scale: Optional[float] = None, + dq: Optional[torch.Tensor] = None, + dk: Optional[torch.Tensor] = None, + dv: Optional[torch.Tensor] = None, + dkv_dtype: torch.dtype = torch.bfloat16, + dk_acc: Optional[torch.Tensor] = None, + dv_acc: Optional[torch.Tensor] = None, + row_map: Optional[torch.Tensor] = None, + max_seqlen_k: Optional[int] = None, + workspace: Optional[Workspace] = None, +) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Backward: ``(dq, dk, dv)``. ``dq`` is BF16 ``[T_q, H_q, 256]``. dK / dV are summed over the query heads of + each key head in FP32 and leave the kernel as: + + - ``dkv_dtype=torch.bfloat16`` (default): BF16 ``[T_k, H_kv, 256]``, rounded once; + - ``dkv_dtype=torch.float32``: FP32 window gradients ``[T_k, H_kv, 256]``, stored (``dk`` / ``dv`` overwritten + when given); + - ``row_map`` given: FP32 adds into the caller's ``dk_acc`` / ``dv_acc`` ``[N_dst, H_kv, 256]``: window row ``j`` + adds into destination row ``row_map[j]`` (int32 ``[T_k]``, rows may repeat; out-of-range rows are not + checked), and the destinations are returned as ``dk`` / ``dv``. + + ``max_seqlen_k`` (host int, optional) must be at least every segment's key count (the caller's contract, as + FlashAttention's ``max_seqlen_k``: a smaller value leaves key blocks out); it sizes the backward grid, one CTA pair + per 128 keys of the longest segment, per key head and segment. ``None`` bounds it by ``T_k``: always correct, + with more CTAs that retire at once. It is never read from the device. + """ + dims = _check_operands(q, k, v, cu_seqlens_q, cu_seqlens_k) + if max_seqlen_k is not None and int(max_seqlen_k) < 0: + raise ValueError(f"max_seqlen_k must be >= 0, got {max_seqlen_k}") + total_q, total_k = dims["total_q"], dims["total_k"] + heads_q, heads_kv = dims["num_heads_q"], dims["num_heads_kv"] + device = q.device + _check_out("out", out, (total_q, heads_q, HEAD_DIM), torch.bfloat16, device) + _check_out("dout", dout, (total_q, heads_q, HEAD_DIM), torch.bfloat16, device) + _check_out("lse", lse, (heads_q, total_q), torch.float32, device) + kv_shape = (total_k, heads_kv, HEAD_DIM) + if row_map is not None: + if dk_acc is None or dv_acc is None: + raise ValueError("row_map needs both dk_acc and dv_acc") + if ( + row_map.dtype != torch.int32 + or tuple(row_map.shape) != (total_k,) + or row_map.device != device + ): + raise ValueError( + f"row_map must be int32 [{total_k}] on {device}, got {row_map.dtype} {tuple(row_map.shape)}" + ) + if not row_map.is_contiguous(): + raise ValueError("row_map must be contiguous") + dst_shape = (int(dk_acc.shape[0]),) + kv_shape[1:] if dk_acc.dim() == 3 else () + _check_out("dk_acc", dk_acc, dst_shape, torch.float32, device) + _check_out("dv_acc", dv_acc, dst_shape, torch.float32, device) + mode, dk_out, dv_out = "rowmap", dk_acc, dv_acc + else: + if dk_acc is not None or dv_acc is not None: + raise ValueError( + "dk_acc / dv_acc need a row_map (use dk / dv with dkv_dtype=torch.float32 to store)" + ) + if dkv_dtype not in (torch.bfloat16, torch.float32): + raise ValueError( + f"dkv_dtype must be torch.bfloat16 or torch.float32, got {dkv_dtype}" + ) + mode = "bf16" if dkv_dtype == torch.bfloat16 else "fp32" + if dk is None: + dk = torch.empty(kv_shape, dtype=dkv_dtype, device=device) + else: + _check_out("dk", dk, kv_shape, dkv_dtype, device) + if dv is None: + dv = torch.empty(kv_shape, dtype=dkv_dtype, device=device) + else: + _check_out("dv", dv, kv_shape, dkv_dtype, device) + dk_out, dv_out = dk, dv + if dq is None: + dq = torch.empty( + total_q, heads_q, HEAD_DIM, dtype=torch.bfloat16, device=device + ) + else: + _check_out("dq", dq, (total_q, heads_q, HEAD_DIM), torch.bfloat16, device) + if total_q == 0 or total_k == 0: + # No (query, key) pair: every gradient is zero; a row-map destination keeps its contents. + dq.zero_() + if mode != "rowmap": + dk_out.zero_() + dv_out.zero_() + return dq, dk_out, dv_out + device_index = device.index + ws = workspace if workspace is not None else default_workspace(device_index) + dq_accum, lse_log2, dpsum = ws.scratch(total_q, heads_q, dims["num_segments"]) + scale = _scale(softmax_scale) + rows = cake_jit.rows_pad(total_q, dims["num_segments"]) + common = { + "cu_seqlens_q": cu_seqlens_q, + "cu_seqlens_k": cu_seqlens_k, + "num_segments": dims["num_segments"], + "rows_pad": rows, + "T_q": total_q, + "H_q": heads_q, + "softmax_scale": scale, + } + _launch( + "bwd_pre", + { + **common, + "O": out, + "dO": dout, + "LSE": lse, + "LSE_log2": lse_log2, + "dPsum": dpsum, + "dQ_accum": dq_accum, + }, + dims, + device_index, + ) + _launch( + _BWD_STAGE[mode], + { + **common, + "Q": q, + "dO": dout, + "K": k, + "V": v, + "lse_log2": lse_log2, + "dpsum": dpsum, + "dq_accum": dq_accum, + "dk_out": dk_out, + "dv_out": dv_out, + "h_kv": heads_kv, + "group": heads_q // heads_kv, + "softmax_scale_log2": scale * LOG2_E, + "row_map": row_map, + }, + { + **dims, + "max_seqlen_k": None if max_seqlen_k is None else int(max_seqlen_k), + }, + device_index, + ) + _launch("bwd_post", {**common, "dQ_accum": dq_accum, "dQ": dq}, dims, device_index) + return dq, dk_out, dv_out + + +# ---------------------------------------------------------------------------------------------------------------- +# Autograd + +# ``dkv_acc``: a zero-argument callable, called in backward, returning ``(dk_acc, dv_acc, row_map)`` (FP32 adds into +# the caller's destinations through the row map) or ``(dk_win, dv_win)`` (FP32 window gradients stored into them). +# Under activation checkpointing a destination passed at forward time would live for every layer at once; one the +# callable produces lives only while its backward runs. +DkvAccSource = Callable[ + [], + Union[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ], +] + + +def _resolve_dkv_acc( + dkv_acc: DkvAccSource, +) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + produced = dkv_acc() + if not isinstance(produced, tuple) or len(produced) not in (2, 3): + raise ValueError( + "a dkv_acc callable returns (dk_acc, dv_acc, row_map) or (dk_win, dv_win)" + ) + if not all(isinstance(t, torch.Tensor) for t in produced): + raise ValueError("a dkv_acc callable returns tensors") + if len(produced) == 2: + return produced[0], produced[1], None + return produced[0], produced[1], produced[2] + + +class VarlenGQAAttentionFunction(torch.autograd.Function): + """``(out, lse) = attention(q, k, v)``; the saved set is ``(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)`` + through ``ctx.save_for_backward`` alone, with scalars, the ``dkv_acc`` callable and a weak reference to the + workspace on ``ctx`` (a checkpointed forward keeps no tensor alive through ``ctx``).""" + + @staticmethod + def forward( # type: ignore[override] + ctx: Any, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + softmax_scale: float, + dkv_acc: Optional[DkvAccSource], + max_seqlen_k: Optional[int], + workspace: Optional[Workspace], + ) -> tuple[torch.Tensor, torch.Tensor]: + out, lse = gqa_attention_fwd( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + softmax_scale=softmax_scale, + workspace=workspace, + ) + ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k) + ctx.softmax_scale = softmax_scale + ctx.max_seqlen_k = max_seqlen_k + ctx.dkv_acc = dkv_acc + ctx.workspace_ref = None if workspace is None else weakref.ref(workspace) + ctx.mark_non_differentiable(lse) + return out, lse + + @staticmethod + def backward( + ctx: Any, dout: Optional[torch.Tensor], _dlse: Any + ) -> tuple[Optional[torch.Tensor], ...]: + q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors + workspace = None if ctx.workspace_ref is None else ctx.workspace_ref() + if dout is None: + dout = torch.zeros_like(out) + elif not dout.is_contiguous(): + dout = dout.contiguous() + kwargs: dict[str, Any] = { + "softmax_scale": ctx.softmax_scale, + "max_seqlen_k": ctx.max_seqlen_k, + "workspace": workspace, + } + if ctx.dkv_acc is None: + dq, dk, dv = gqa_attention_bwd( + q, k, v, out, lse, dout, cu_seqlens_q, cu_seqlens_k, **kwargs + ) + return dq, dk, dv, None, None, None, None, None, None + dk_t, dv_t, row_map = _resolve_dkv_acc(ctx.dkv_acc) + if row_map is None: + kwargs.update(dk=dk_t, dv=dv_t, dkv_dtype=torch.float32) + else: + kwargs.update(dk_acc=dk_t, dv_acc=dv_t, row_map=row_map) + dq = gqa_attention_bwd( + q, k, v, out, lse, dout, cu_seqlens_q, cu_seqlens_k, **kwargs + )[0] + return dq, None, None, None, None, None, None, None, None + + +def varlen_gqa_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + *, + softmax_scale: Optional[float] = None, + causal: bool = True, + dkv_acc: Optional[DkvAccSource] = None, + max_seqlen_k: Optional[int] = None, + workspace: Optional[Workspace] = None, + return_lse: bool = False, +) -> Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + """Packed-varlen causal GQA attention with autograd (see :func:`gqa_attention_fwd` / :func:`gqa_attention_bwd`). + + ``dkv_acc=None``: autograd returns BF16 gradients for ``k`` / ``v``. A ``dkv_acc`` callable is called in + backward; it returns ``(dk_acc, dv_acc, row_map)`` (FP32 adds into the destinations) or ``(dk_win, dv_win)`` (FP32 + window gradients stored into them), and ``k`` / ``v`` then get no gradient. ``max_seqlen_k``: see + :func:`gqa_attention_bwd`. + """ + if not causal: + raise NotImplementedError( + "only causal attention (bottom-right aligned per segment) is implemented" + ) + if dkv_acc is not None and not callable(dkv_acc): + raise TypeError( + "dkv_acc must be None or a zero-argument callable (see DkvAccSource)" + ) + out, lse = VarlenGQAAttentionFunction.apply( + q, + k, + v, + cu_seqlens_q, + cu_seqlens_k, + _scale(softmax_scale), + dkv_acc, + None if max_seqlen_k is None else int(max_seqlen_k), + workspace, + ) + return (out, lse) if return_lse else out diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_build.py b/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_build.py new file mode 100644 index 0000000..0075613 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_build.py @@ -0,0 +1,324 @@ +""" +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. + +Library loader of one Cake-generated FlashInfer package. + +``build.json`` (schema ``cake-build/1``) next to this file is the only record +of what the package compiles: one entry per ``[_]_`` +library with its package-relative CUDA sources, nvcc flags, toolkit-conditional +flags and digests. A library is served from ``aot/.so`` when a builder +(a vendoring consumer's build rules, or FlashInfer's +``scripts/check_cake_build_manifests.py``) has compiled it from the manifest, +and otherwise through FlashInfer's JIT from the same sources and flags. + +The prebuilt path never imports ``flashinfer``, so the package stays usable in +a process without a FlashInfer installation. A builder that fills ``aot/`` +also writes the sidecar ``aot/build_info.json`` (schema ``cake-build-info/1``: +``{"schema", "cuda": "13.0.88", "cudart_version": 13000, "arches": [...]}``, +``arches`` naming every architecture the builder compiled). Every prebuilt +library exports ``build_cuda_version()``, the ``CUDART_VERSION`` it was +compiled against, and :func:`load_library` refuses a library whose value +disagrees with the sidecar. + +Public surface (``__all__``): :func:`manifest`, :func:`libraries_by_name`, +:func:`prebuilt_path`, :func:`library_path`, :func:`load_library`, +:func:`jit_spec`, :func:`jit_specs` and :func:`record_cuda_version`. + +The Cake exporter writes this file byte-identically into every package from +its loader template; do not edit it in place. +""" + +from __future__ import annotations + +import functools +import hashlib +import json +import re +from pathlib import Path +from typing import Any + +__all__ = [ + "jit_spec", + "jit_specs", + "libraries_by_name", + "library_path", + "load_library", + "manifest", + "prebuilt_path", + "record_cuda_version", +] + +_PACKAGE = Path(__file__).parent +_MANIFEST_SCHEMA = "cake-build/1" +_BUILD_INFO_SCHEMA = "cake-build-info/1" +_MANIFEST_PATH = _PACKAGE / "build.json" +_BUILD_INFO_PATH = _PACKAGE / "aot" / "build_info.json" +_CUDA_CONSTRAINT = re.compile(r"(==|>=)([0-9]+)\.([0-9]+)") + + +@functools.cache +def manifest() -> dict[str, Any]: + """The parsed ``build.json`` of this package (read on first use, not at import). + + One cached document shared by every caller: read it, do not modify it. + """ + try: + text = _MANIFEST_PATH.read_text(encoding="utf-8") + except FileNotFoundError as exc: + raise RuntimeError( + f"{_MANIFEST_PATH} is missing: the package was not exported with its {_MANIFEST_SCHEMA} manifest", + ) from exc + document = json.loads(text) + schema = document.get("schema") + if schema != _MANIFEST_SCHEMA: + raise RuntimeError( + f"{_MANIFEST_PATH} has schema {schema!r}; this loader reads {_MANIFEST_SCHEMA!r}", + ) + return document + + +@functools.cache +def libraries_by_name() -> dict[str, dict[str, Any]]: + """``name -> entry`` over the ``libraries`` of ``build.json``.""" + index: dict[str, dict[str, Any]] = {} + for entry in manifest()["libraries"]: + if entry["name"] in index: + raise RuntimeError(f"{_MANIFEST_PATH} lists {entry['name']!r} twice") + index[entry["name"]] = entry + return index + + +def _entry(name: str) -> dict[str, Any]: + try: + return libraries_by_name()[name] + except KeyError: + raise KeyError(f"{name!r} is not a library of {_MANIFEST_PATH}") from None + + +def prebuilt_path(name: str) -> Path: + """Where a builder places the compiled library of ``name`` (``aot/.so``).""" + return _PACKAGE / _entry(name)["path"] + + +def _cuda_constraint_matches(constraint: str, cuda: tuple[int, int]) -> bool: + """Does toolkit ``cuda`` (major, minor) satisfy ``"==X.Y"`` or ``">=X.Y"``?""" + match = _CUDA_CONSTRAINT.fullmatch(constraint) + if match is None: + raise ValueError( + f"unsupported CUDA constraint {constraint!r}; expected '==X.Y' or '>=X.Y'", + ) + operator, major, minor = match.groups() + version = (int(major), int(minor)) + if operator == "==": + return tuple(cuda) == version + return tuple(cuda) >= version + + +def _conditional_flags(items: list[Any], cuda: tuple[int, int]) -> list[str]: + """The flags of every ``nvcc_flags_if`` item whose ``cuda`` constraint matches.""" + flags: list[str] = [] + for item in items: + if _cuda_constraint_matches(item["cuda"], cuda): + flags.extend(item["flags"]) + return flags + + +def _jit_toolchain(name: str) -> tuple[Any, Any]: + """``(gen_jit_spec, get_cuda_version)`` of the local FlashInfer, or the error.""" + try: + from flashinfer.jit.core import gen_jit_spec + from flashinfer.jit.cpp_ext import get_cuda_version + except ImportError as exc: + raise RuntimeError( + f"{name} is not prebuilt ({prebuilt_path(name)}) and flashinfer's JIT is not available", + ) from exc + return gen_jit_spec, get_cuda_version + + +def _local_cuda_version(get_cuda_version: Any) -> tuple[int, int]: + version = get_cuda_version() + return (int(version.major), int(version.minor)) + + +def _version(text: str) -> tuple[int, int]: + """``(major, minor)`` of a dotted version such as ``"13.0"`` or ``"13.0.88"``.""" + parts = str(text).split(".") + if len(parts) < 2 or not all(part.isdigit() for part in parts[:2]): + raise RuntimeError(f"malformed CUDA version {text!r}") + return (int(parts[0]), int(parts[1])) + + +def _toolkit_builds(entry: dict[str, Any], cuda: tuple[int, int]) -> bool: + """Does toolkit ``cuda`` reach the entry's ``min_cuda`` floor?""" + return cuda >= _version(entry["min_cuda"]) + + +def _verified_sources(entry: dict[str, Any]) -> list[Path]: + """The entry's sources as package paths, each checked against its digest.""" + sources: list[Path] = [] + for relative in entry["sources"]: + path = _PACKAGE / relative + digest = hashlib.sha256(path.read_bytes()).hexdigest() + expected = entry["sha256"][relative] + if digest != expected: + raise RuntimeError( + f"{path} does not match its build.json digest for {entry['name']}: {digest} != {expected}", + ) + sources.append(path) + return sources + + +def jit_spec(name: str) -> Any: + """FlashInfer's ``JitSpecNvcc`` of ``name``, built from the manifest entry. + + Sources are digest-verified here and only here, so a modified or partially + delivered source tree fails closed before anything compiles. The spec name + carries the closure digest: a changed closure never reuses a stale JIT cache + entry. Flags are ``arch_flags[arch]`` + the entry's ``nvcc_flags`` + the + ``nvcc_flags_if`` items matching the local toolkit; FlashInfer's JIT adds + its own defaults (``-O3``, ``-fPIC``, ...), so ``common.nvcc_flags`` are not + repeated here. A toolkit below the entry's ``min_cuda`` raises. + """ + entry = _entry(name) + gen_jit_spec, get_cuda_version = _jit_toolchain(name) + document = manifest() + cuda = _local_cuda_version(get_cuda_version) + if not _toolkit_builds(entry, cuda): + raise RuntimeError( + f"{name} needs CUDA {entry['min_cuda']} or newer; the local toolkit is {cuda[0]}.{cuda[1]}", + ) + flags = [ + *document["arch_flags"][entry["arch"]], + *entry["nvcc_flags"], + *_conditional_flags(entry["nvcc_flags_if"], cuda), + ] + include_dirs = [_PACKAGE / d for d in document["common"]["include_dirs"]] + return gen_jit_spec( + f"{name}_{entry['closure_sha256'][:20]}", + _verified_sources(entry), + extra_cuda_cflags=flags, + extra_ldflags=list(document["common"]["link_flags"]), + extra_include_paths=include_dirs, + use_fast_math=False, + ) + + +def _names_of(arch: str) -> list[str]: + index = libraries_by_name() + return [name for name in index if index[name]["arch"] == arch] + + +def jit_specs(arch: str) -> list[Any]: + """JIT specifications of the ``arch`` libraries still to build locally. + + The entries without a prebuilt file that the local toolkit can build: the + list feeds ``flashinfer.jit.build_jit_specs`` warm-ups, so an entry whose + ``min_cuda`` exceeds the local toolkit is left out (the manifest check + script's "skip min_cuda > toolkit" rule), while :func:`jit_spec` on that + entry by name still raises. When every ``arch`` library is prebuilt the + list is empty and ``flashinfer`` is not imported. + """ + pending = [name for name in _names_of(arch) if not prebuilt_path(name).is_file()] + if not pending: + return [] + _, get_cuda_version = _jit_toolchain(pending[0]) + cuda = _local_cuda_version(get_cuda_version) + return [jit_spec(name) for name in pending if _toolkit_builds(_entry(name), cuda)] + + +def library_path(name: str) -> Path: + """The prebuilt library of ``name`` when present, else its JIT cache path.""" + path = prebuilt_path(name) + if path.is_file(): + return path + return jit_spec(name).get_library_path() + + +@functools.cache +def _build_info() -> dict[str, Any]: + """The builder-written ``aot/build_info.json`` sidecar (prebuilt path only).""" + try: + text = _BUILD_INFO_PATH.read_text(encoding="utf-8") + except FileNotFoundError as exc: + raise RuntimeError( + f"{_BUILD_INFO_PATH} is missing although prebuilt libraries exist under {_PACKAGE / 'aot'}: the builder must write the {_BUILD_INFO_SCHEMA} sidecar", + ) from exc + info = json.loads(text) + schema = info.get("schema") + if schema != _BUILD_INFO_SCHEMA: + raise RuntimeError( + f"{_BUILD_INFO_PATH} has schema {schema!r}; this loader reads {_BUILD_INFO_SCHEMA!r}", + ) + return info + + +def _prebuilt_info(path: Path, arch: str) -> dict[str, Any]: + """The sidecar, which must record a build of ``arch`` for the prebuilt ``path``.""" + info = _build_info() + if arch not in info["arches"]: + raise RuntimeError( + f"{path} exists but {_BUILD_INFO_PATH} records no {arch} build (arches {info['arches']})", + ) + return info + + +def load_library(name: str) -> Any: + """The loaded tvm-ffi module of ``name``: the prebuilt library, else a JIT build.""" + entry = _entry(name) + path = _PACKAGE / entry["path"] + if not path.is_file(): + return jit_spec(name).build_and_load() + info = _prebuilt_info(path, entry["arch"]) + import tvm_ffi + + module = tvm_ffi.load_module(str(path)) + built = int(module.build_cuda_version()) + if built != int(info["cudart_version"]): + raise RuntimeError( + f"{path} was compiled against CUDART {built} but {_BUILD_INFO_PATH} records {info['cudart_version']}", + ) + return module + + +def record_cuda_version(arch: str) -> tuple[int, int]: + """``(major, minor)`` of the toolkit that builds the ``arch`` libraries. + + Prebuilt path, taken as soon as any library of the manifest exists under + ``aot/``: the sidecar ``aot/build_info.json`` is mandatory (missing -> + ``RuntimeError``; no local nvcc is consulted) and its ``cuda`` is the + answer whenever its ``arches`` list ``arch`` -- the builder compiled that + architecture, whether or not every one of its libraries is present. An + ``arch`` library under ``aot/`` that the sidecar does not list raises, as + :func:`load_library` does. + + JIT path, when nothing is prebuilt or the builder did not compile + ``arch``: the local toolkit of FlashInfer's JIT (``RuntimeError`` in a + process without a FlashInfer installation). Call it for the architecture + you are about to load. + """ + names = _names_of(arch) + if not names: + raise KeyError(f"{arch!r} has no libraries in {_MANIFEST_PATH}") + index = libraries_by_name() + prebuilt = [name for name in index if (_PACKAGE / index[name]["path"]).is_file()] + if prebuilt: + of_arch = [name for name in prebuilt if index[name]["arch"] == arch] + if of_arch: + return _version(_prebuilt_info(prebuilt_path(of_arch[0]), arch)["cuda"]) + info = _build_info() + if arch in info["arches"]: + return _version(info["cuda"]) + _, get_cuda_version = _jit_toolchain(names[0]) + return _local_cuda_version(get_cuda_version) diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_jit.py b/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_jit.py new file mode 100644 index 0000000..982ebc9 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/cake_jit.py @@ -0,0 +1,384 @@ +""" +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. + +Registry of the Cake packed-varlen GQA training attention kernels, head +dimension 256 (SM100 / SM103 / SM107). The loader and the launch plan are +:mod:`.cake_backend`. + +Every record of ``MODULES`` is one mechanically exported program: a CUDA +device translation unit plus a tvm-ffi binding, compiled once per +architecture. A record is keyed by its stage name (or ``_`` +for a stage that ships a different program per architecture) and carries the +host plan: the ``stage`` it serves, the architectures it is delivered for and +the ``run`` entry. The derived leaves (``library``, ``kernel_symbol``, +``launch``, ``arg_plan``) are written by the exporter from the rendered +binding; ``library`` is the program identity (``cake_gqa_train_<20 hex>``), +so a kernel swap renames it. ``MODULES`` is a literal table (the exporter +and the manifest checker read it with ``ast.literal_eval``). + +The stages, in launch order: ``fwd`` (O and the natural-log LSE), +``bwd_pre`` (``dPsum``, ``LSE * log2(e)`` and +the zeroed FP32 dQ accumulator in the pack-2 scratch layout), the backward +main kernel in its three dK / dV outputs ``bwd_fp32`` (FP32 window +gradients), ``bwd_rowmap`` (FP32 adds into caller-owned destinations through +a row map) and ``bwd_bf16``, and ``bwd_post`` (dQ from the accumulator). +Every launch is a plain stream-ordered launch: nothing cooperative, no PDL. +""" + +from __future__ import annotations + +from typing import Any, Optional + +ARCHES = ("sm_100a", "sm_103a", "sm_107a") +ARCH_BY_CAPABILITY = {(10, 0): "sm_100a", (10, 3): "sm_103a", (10, 7): "sm_107a"} +HEAD_DIM = 256 +PACKAGE = "flashinfer.experimental.cake_gqa_train" +LIBRARY_PREFIX = "cake_gqa_train" +MANIFEST = "build.json" + +# The stages of one training step, in launch order; every stage is served on +# every architecture of ``ARCHES`` by exactly one record. +STAGES = ("fwd", "bwd_pre", "bwd_fp32", "bwd_rowmap", "bwd_bf16", "bwd_post") + +MODULES: dict[str, dict[str, Any]] = { + "fwd": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "Q"], + ["buffer", "K"], + ["buffer", "V"], + ["buffer", "O"], + ["buffer", "LSE"], + ["buffer", "cu_seqlens_q"], + ["buffer", "cu_seqlens_k"], + ["buffer", "tile_counter"], + ["parameter", "num_segments"], + ["parameter", "num_heads_q"], + ["parameter", "group"], + ["parameter", "total_q"], + ["parameter", "softmax_scale_log2"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_gqa_train_493c4a28dafa5fc6b2f1", + "launch": { + "block": (384, 1, 1), + "cluster": (2, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 232064, + "use_pdl": False, + }, + "library": "cake_gqa_train_493c4a28dafa5fc6b2f1", + "stage": "fwd", + }, + "bwd_pre": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "O"], + ["buffer", "dO"], + ["buffer", "LSE"], + ["buffer", "cu_seqlens_q"], + ["buffer", "LSE_log2"], + ["buffer", "dPsum"], + ["buffer", "dQ_accum"], + ["parameter", "T_q"], + ["parameter", "H_q"], + ["parameter", "num_segments"], + ["parameter", "rows_pad"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_gqa_train_a91752e8a2732bf2e246", + "launch": { + "block": (256, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 0, + "use_pdl": False, + }, + "library": "cake_gqa_train_a91752e8a2732bf2e246", + "stage": "bwd_pre", + }, + "bwd_fp32": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "Q"], + ["buffer", "dO"], + ["buffer", "K"], + ["buffer", "V"], + ["buffer", "lse_log2"], + ["buffer", "dpsum"], + ["buffer", "dq_accum"], + ["buffer", "cu_seqlens_q"], + ["buffer", "cu_seqlens_k"], + ["buffer", "dk_out"], + ["buffer", "dv_out"], + ["parameter", "num_segments"], + ["parameter", "h_kv"], + ["parameter", "group"], + ["parameter", "rows_pad"], + ["parameter", "softmax_scale"], + ["parameter", "softmax_scale_log2"], + ["buffer", "row_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_gqa_train_6fa13db592acdb66a6c9", + "launch": { + "block": (512, 1, 1), + "cluster": (2, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 231424, + "use_pdl": False, + }, + "library": "cake_gqa_train_6fa13db592acdb66a6c9", + "stage": "bwd_fp32", + }, + "bwd_rowmap": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "Q"], + ["buffer", "dO"], + ["buffer", "K"], + ["buffer", "V"], + ["buffer", "lse_log2"], + ["buffer", "dpsum"], + ["buffer", "dq_accum"], + ["buffer", "cu_seqlens_q"], + ["buffer", "cu_seqlens_k"], + ["buffer", "dk_out"], + ["buffer", "dv_out"], + ["parameter", "num_segments"], + ["parameter", "h_kv"], + ["parameter", "group"], + ["parameter", "rows_pad"], + ["parameter", "softmax_scale"], + ["parameter", "softmax_scale_log2"], + ["buffer", "row_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_gqa_train_37f4cd45025705f3cfc8", + "launch": { + "block": (512, 1, 1), + "cluster": (2, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 231424, + "use_pdl": False, + }, + "library": "cake_gqa_train_37f4cd45025705f3cfc8", + "stage": "bwd_rowmap", + }, + "bwd_bf16": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "Q"], + ["buffer", "dO"], + ["buffer", "K"], + ["buffer", "V"], + ["buffer", "lse_log2"], + ["buffer", "dpsum"], + ["buffer", "dq_accum"], + ["buffer", "cu_seqlens_q"], + ["buffer", "cu_seqlens_k"], + ["buffer", "dk_out"], + ["buffer", "dv_out"], + ["parameter", "num_segments"], + ["parameter", "h_kv"], + ["parameter", "group"], + ["parameter", "rows_pad"], + ["parameter", "softmax_scale"], + ["parameter", "softmax_scale_log2"], + ["buffer", "row_map"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_gqa_train_1afdfdde603ff8806b4d", + "launch": { + "block": (512, 1, 1), + "cluster": (2, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 231424, + "use_pdl": False, + }, + "library": "cake_gqa_train_1afdfdde603ff8806b4d", + "stage": "bwd_bf16", + }, + "bwd_post": { + "arches": ["sm_100a", "sm_103a", "sm_107a"], + "arg_plan": [ + ["buffer", "dQ_accum"], + ["buffer", "cu_seqlens_q"], + ["buffer", "dQ"], + ["parameter", "T_q"], + ["parameter", "H_q"], + ["parameter", "num_segments"], + ["parameter", "rows_pad"], + ["parameter", "softmax_scale"], + ["grid", "grid_x"], + ["grid", "grid_y"], + ["grid", "grid_z"], + ], + "ffi_entry": "run", + "kernel_symbol": "kernel_cake_gqa_train_76decc9a5d22dc91ddd1", + "launch": { + "block": (256, 1, 1), + "cluster": (1, 1, 1), + "cooperative": False, + "dynamic_smem_bytes": 0, + "use_pdl": False, + }, + "library": "cake_gqa_train_76decc9a5d22dc91ddd1", + "stage": "bwd_post", + }, +} + +# --------------------------------------------------------------------------- +# The host plan: grids and scratch sizes from host-known shapes only (the +# tensors' shapes, the segment count ``len(cu_seqlens) - 1`` and the SM +# count); nothing here reads a tensor. + +# Forward: a CTA pair covers 256 query rows of one (segment, query head); the +# grid is persistent (tickets from a self-resetting counter), one pair per two +# SMs at most, any count >= 1 (no co-residency is assumed). +FWD_TILE_ROWS = 256 +FWD_CLUSTER = 2 +# Backward main kernel: a CTA pair covers 128 keys of one (segment, key head) +# and loops the key head's query heads two at a time (``group`` even). The +# grid is ``(2 N, H_kv, B)``: the pair at ``(2 n, h, b)`` owns key block ``n`` +# of segment ``b``; ``N`` is at least the longest segment's key-block count +# (from ``max_seqlen_k`` when the caller gives it, else from ``T_k``, which +# bounds every segment), and a pair past its segment's last block does nothing. +BWD_TILE_KEYS = 128 +BWD_CLUSTER = 2 +# Pack-2 scratch: scratch head ``s`` holds query heads ``2 s`` and ``2 s + 1``; +# every segment starts on a 128-row tile. The preprocess runs one CTA per 32 +# packed rows of one scratch head, the postprocess one per 64 (each program's +# ``rows_per_cta``; a grid past ``rows_pad`` would read past the scratch). +SCRATCH_TILE_ROWS = 128 +PRE_ROWS_PER_CTA = 32 +POST_ROWS_PER_CTA = 64 +# Byte alignment of every scratch buffer (TMA / bulk copies need 16 B, the +# 256-bit loads and stores of the pre/postprocess 32 B). +SCRATCH_ALIGN = 256 + + +def rows_pad(total_q: int, num_segments: int) -> int: + """Packed rows per scratch head: ``round_up(2 T_q + 128 B, 128)``.""" + rows = 2 * total_q + SCRATCH_TILE_ROWS * num_segments + return -(-rows // SCRATCH_TILE_ROWS) * SCRATCH_TILE_ROWS + + +def scratch_layout( + total_q: int, num_heads_q: int, num_segments: int +) -> dict[str, tuple[int, tuple[int, ...]]]: + """The backward scratch as ``name -> (byte offset, FP32 shape)``: ``lse_log2`` and + ``dpsum`` ``[H_q / 2, rows_pad]``, ``dq_accum`` ``[H_q / 2, rows_pad, 256]``.""" + heads = num_heads_q // 2 + rows = rows_pad(total_q, num_segments) + layout: dict[str, tuple[int, tuple[int, ...]]] = {} + offset = 0 + for name, shape in ( + ("dq_accum", (heads, rows, HEAD_DIM)), + ("lse_log2", (heads, rows)), + ("dpsum", (heads, rows)), + ): + layout[name] = (offset, shape) + count = 1 + for extent in shape: + count *= extent + offset += -(-4 * count // SCRATCH_ALIGN) * SCRATCH_ALIGN + return layout + + +def scratch_bytes(total_q: int, num_heads_q: int, num_segments: int) -> int: + """Bytes of the backward scratch of one call (:func:`scratch_layout`).""" + if total_q == 0: + return 0 + offset, shape = scratch_layout(total_q, num_heads_q, num_segments)["dpsum"] + return offset + -(-4 * shape[0] * shape[1] // SCRATCH_ALIGN) * SCRATCH_ALIGN + + +def grid( + stage: str, + *, + total_q: int, + total_k: int, + num_heads_q: int, + num_heads_kv: int, + num_segments: int, + sm_count: int, + max_seqlen_k: Optional[int] = None, +) -> tuple[int, int, int]: + """The launch grid of ``stage``. Tile counts are upper bounds from host + values (``sum_b ceil(len_b / n) <= T // n + B``; the longest segment has + at most ``max_seqlen_k`` or ``T_k`` keys); the kernels decode the segments + on the device and retire the CTAs past the last tile.""" + if stage == "fwd": + tiles = num_heads_q * (total_q // FWD_TILE_ROWS + num_segments) + pairs = max(1, min(sm_count // FWD_CLUSTER, tiles)) + return (FWD_CLUSTER * pairs, 1, 1) + if stage in ("bwd_pre", "bwd_post"): + rows = rows_pad(total_q, num_segments) + per_cta = PRE_ROWS_PER_CTA if stage == "bwd_pre" else POST_ROWS_PER_CTA + return (rows // per_cta, num_heads_q // 2, 1) + if stage in ("bwd_fp32", "bwd_rowmap", "bwd_bf16"): + longest = total_k if max_seqlen_k is None else min(max_seqlen_k, total_k) + blocks = max(1, -(-longest // BWD_TILE_KEYS)) + return (BWD_CLUSTER * blocks, num_heads_kv, num_segments) + raise KeyError(f"no stage {stage!r}; the stages are {STAGES}") + + +def stage_record(stage: str, arch: str) -> str: + """The record of ``MODULES`` serving ``stage`` on ``arch``.""" + for name, found in MODULES.items(): + if found["stage"] == stage and arch in found["arches"]: + return name + raise KeyError(f"no record serves stage {stage!r} on {arch}") + + +__all__ = [ + "ARCHES", + "ARCH_BY_CAPABILITY", + "BWD_CLUSTER", + "BWD_TILE_KEYS", + "FWD_CLUSTER", + "FWD_TILE_ROWS", + "HEAD_DIM", + "LIBRARY_PREFIX", + "MANIFEST", + "MODULES", + "PACKAGE", + "SCRATCH_ALIGN", + "POST_ROWS_PER_CTA", + "PRE_ROWS_PER_CTA", + "SCRATCH_TILE_ROWS", + "STAGES", + "grid", + "rows_pad", + "scratch_bytes", + "scratch_layout", + "stage_record", +] diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/.clang-format b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/.clang-format new file mode 100644 index 0000000..47a38a9 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/.clang-format @@ -0,0 +1,2 @@ +DisableFormat: true +SortIncludes: Never diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_ffi_checks.h b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_ffi_checks.h new file mode 100644 index 0000000..129288f --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_ffi_checks.h @@ -0,0 +1,89 @@ +/* + * 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. + */ + +// Host surface of the generated Cake launchers (``*_binding.cu``) for a package built without +// FlashInfer's csrc: the TVM-FFI / DLPack headers the launchers use, the named tensor-argument +// checks and the CUDA runtime version the library was compiled against. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +inline void check_cuda_tensor(const tvm::ffi::TensorView& t, const char* name) { + TVM_FFI_CHECK(t.device().device_type == kDLCUDA, ValueError) + << name << " must be a CUDA tensor, got device_type=" << static_cast(t.device().device_type); +} + +inline void check_dtype(const tvm::ffi::TensorView& t, DLDataType expected, const char* name) { + DLDataType d = t.dtype(); + TVM_FFI_CHECK(d.code == expected.code && d.bits == expected.bits && d.lanes == expected.lanes, TypeError) + << name << " dtype mismatch: expected DLDataType(code=" << static_cast(expected.code) + << ", bits=" << static_cast(expected.bits) << ", lanes=" << static_cast(expected.lanes) + << "), got (code=" << static_cast(d.code) << ", bits=" << static_cast(d.bits) + << ", lanes=" << static_cast(d.lanes) << ")"; +} + +inline void check_contiguous(const tvm::ffi::TensorView& t, const char* name) { + TVM_FFI_CHECK(t.IsContiguous(), ValueError) << name << " must be contiguous"; +} + +inline void check_same_device(const tvm::ffi::TensorView& t, const tvm::ffi::TensorView& reference, + const char* name, const char* reference_name) { + TVM_FFI_CHECK(t.device().device_type == reference.device().device_type && + t.device().device_id == reference.device().device_id, + ValueError) + << name << " must be on the same device as " << reference_name << ": got device_type=" + << static_cast(t.device().device_type) << " id=" << t.device().device_id + << " versus device_type=" << static_cast(reference.device().device_type) + << " id=" << reference.device().device_id; +} + +// Require the dimensions above the last `trailing` dimensions to form one dense +// row-major chain, so they can be folded into a single logical dimension whose +// step is stride(-(trailing + 1)). Shape products are stride-independent, so +// this verifies the physical layout instead of inventing a folded stride. +inline void check_dense_leading_dims(const tvm::ffi::TensorView& t, int trailing, const char* name) { + TVM_FFI_CHECK(trailing > 0 && t.ndim() >= trailing, ValueError) + << name << " cannot fold leading dimensions above " << trailing << " trailing dims from ndim=" << t.ndim(); + int outer_last = t.ndim() - trailing - 1; + if (outer_last <= 0) { + return; + } + int64_t step = t.stride(outer_last); + TVM_FFI_CHECK(step > 0, ValueError) << name << " physical strides must be positive"; + int64_t expected = step; + for (int axis = outer_last - 1; axis >= 0; --axis) { + expected *= t.size(axis + 1); + if (t.size(axis) > 1) { + TVM_FFI_CHECK(t.stride(axis) == expected, ValueError) + << name << " leading dims are not physically foldable above " << trailing + << " trailing dims: stride(" << axis << ")=" << t.stride(axis) << ", expected " << expected; + } + } +} + +// The CUDA runtime this library was compiled against (``CUDART_VERSION``, e.g. 13000 for 13.0); every +// binding exports it as ``build_cuda_version`` so the package loader can compare it with the prebuilt +// sidecar before the first launch. +inline int64_t cake_ffi_build_cuda_version() { return CUDART_VERSION; } diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu new file mode 100644 index 0000000..005e8da --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu @@ -0,0 +1,350 @@ +/* + * 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 Cake: the tvm-ffi host binding of kernel_cake_gqa_train_1afdfdde603ff8806b4d, for the device source that defines it. Do not edit. +#include +#include +#include + +#include "cake_ffi_checks.h" +#include + +#include + +extern "C" __global__ void kernel_cake_gqa_train_1afdfdde603ff8806b4d(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap dO, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, float* __restrict__ lse_log2, float* __restrict__ dpsum, float* __restrict__ dq_accum, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, __nv_bfloat16* __restrict__ dk_out, __nv_bfloat16* __restrict__ dv_out, int num_segments, int h_kv, int group, int rows_pad, float softmax_scale, float softmax_scale_log2, int* __restrict__ row_map); + +namespace cake_binding_kernel_cake_gqa_train_1afdfdde603ff8806b4d { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gqa_train_1afdfdde603ff8806b4d"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CMod(int64_t a, int64_t b) { + if (b == 0) NoInt64(a, "%", b); + return b == -1 ? 0 : a % b; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_Q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'Q' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "Q"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'Q' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "Q"), Extent(d2, "Q"), Extent(d3, "Q")}; + cuuint64_t global_strides[2] = {Stride(s2, "Q"), Stride(s3, "Q")}; + cuuint32_t box_dim[3] = {64u, 2u, 32u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'Q' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'Q' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'Q' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of Q"); + return map; +} + +CUtensorMap EncodeTma_dO(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'dO' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "dO"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'dO' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "dO"), Extent(d2, "dO"), Extent(d3, "dO")}; + cuuint64_t global_strides[2] = {Stride(s2, "dO"), Stride(s3, "dO")}; + cuuint32_t box_dim[3] = {64u, 2u, 32u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'dO' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'dO' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'dO' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of dO"); + return map; +} + +CUtensorMap EncodeTma_K(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'K' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "K"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'K' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "K"), Extent(d2, "K"), Extent(d3, "K")}; + cuuint64_t global_strides[2] = {Stride(s2, "K"), Stride(s3, "K")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'K' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'K' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'K' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of K"); + return map; +} + +CUtensorMap EncodeTma_V(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'V' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "V"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'V' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "V"), Extent(d2, "V"), Extent(d3, "V")}; + cuuint64_t global_strides[2] = {Stride(s2, "V"), Stride(s3, "V")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'V' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'V' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'V' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of V"); + return map; +} + +void Run(TensorView a_Q, TensorView a_dO, TensorView a_K, TensorView a_V, TensorView a_lse_log2, TensorView a_dpsum, TensorView a_dq_accum, TensorView a_cu_seqlens_q, TensorView a_cu_seqlens_k, TensorView a_dk_out, TensorView a_dv_out, int64_t a_num_segments, int64_t a_h_kv, int64_t a_group, int64_t a_rows_pad, double a_softmax_scale, double a_softmax_scale_log2, Optional a_row_map, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_Q.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_Q, "Q"); + check_dtype(a_Q, DLDataType{kDLBfloat, 16, 1}, "Q"); + check_cuda_tensor(a_dO, "dO"); + check_dtype(a_dO, DLDataType{kDLBfloat, 16, 1}, "dO"); + check_cuda_tensor(a_K, "K"); + check_dtype(a_K, DLDataType{kDLBfloat, 16, 1}, "K"); + check_cuda_tensor(a_V, "V"); + check_dtype(a_V, DLDataType{kDLBfloat, 16, 1}, "V"); + check_cuda_tensor(a_lse_log2, "lse_log2"); + check_dtype(a_lse_log2, DLDataType{kDLFloat, 32, 1}, "lse_log2"); + check_cuda_tensor(a_dpsum, "dpsum"); + check_dtype(a_dpsum, DLDataType{kDLFloat, 32, 1}, "dpsum"); + check_cuda_tensor(a_dq_accum, "dq_accum"); + check_dtype(a_dq_accum, DLDataType{kDLFloat, 32, 1}, "dq_accum"); + check_cuda_tensor(a_cu_seqlens_q, "cu_seqlens_q"); + check_dtype(a_cu_seqlens_q, DLDataType{kDLInt, 32, 1}, "cu_seqlens_q"); + check_cuda_tensor(a_cu_seqlens_k, "cu_seqlens_k"); + check_dtype(a_cu_seqlens_k, DLDataType{kDLInt, 32, 1}, "cu_seqlens_k"); + check_cuda_tensor(a_dk_out, "dk_out"); + check_dtype(a_dk_out, DLDataType{kDLBfloat, 16, 1}, "dk_out"); + check_cuda_tensor(a_dv_out, "dv_out"); + check_dtype(a_dv_out, DLDataType{kDLBfloat, 16, 1}, "dv_out"); + TVM_FFI_CHECK(a_num_segments >= -2147483648LL && a_num_segments <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_segments' = " << a_num_segments + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_h_kv >= -2147483648LL && a_h_kv <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'h_kv' = " << a_h_kv + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_group >= -2147483648LL && a_group <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'group' = " << a_group + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_pad >= -2147483648LL && a_rows_pad <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_pad' = " << a_rows_pad + << " is outside the i32 range [-2147483648, 2147483647]"; + if (a_row_map.has_value()) check_cuda_tensor(a_row_map.value(), "row_map"); + if (a_row_map.has_value()) check_dtype(a_row_map.value(), DLDataType{kDLInt, 32, 1}, "row_map"); + check_same_device(a_dO, a_Q, "dO", "Q"); + check_same_device(a_K, a_Q, "K", "Q"); + check_same_device(a_V, a_Q, "V", "Q"); + check_contiguous(a_lse_log2, "lse_log2"); + check_same_device(a_lse_log2, a_Q, "lse_log2", "Q"); + check_contiguous(a_dpsum, "dpsum"); + check_same_device(a_dpsum, a_Q, "dpsum", "Q"); + check_contiguous(a_dq_accum, "dq_accum"); + check_same_device(a_dq_accum, a_Q, "dq_accum", "Q"); + check_contiguous(a_cu_seqlens_q, "cu_seqlens_q"); + check_same_device(a_cu_seqlens_q, a_Q, "cu_seqlens_q", "Q"); + check_contiguous(a_cu_seqlens_k, "cu_seqlens_k"); + check_same_device(a_cu_seqlens_k, a_Q, "cu_seqlens_k", "Q"); + check_contiguous(a_dk_out, "dk_out"); + check_same_device(a_dk_out, a_Q, "dk_out", "Q"); + check_contiguous(a_dv_out, "dv_out"); + check_same_device(a_dv_out, a_Q, "dv_out", "Q"); + if (a_row_map.has_value()) check_contiguous(a_row_map.value(), "row_map"); + if (a_row_map.has_value()) check_same_device(a_row_map.value(), a_Q, "row_map", "Q"); + TVM_FFI_CHECK((a_group >= 2LL), ValueError) + << KERNEL_NAME << ": requires (group >= 2)"; + TVM_FFI_CHECK((a_group <= 64LL), ValueError) + << KERNEL_NAME << ": requires (group <= 64)"; + TVM_FFI_CHECK((CMod(a_group, 2LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(group, 2) == 0)"; + TVM_FFI_CHECK((CMod(a_rows_pad, 128LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(rows_pad, 128) == 0)"; + TVM_FFI_CHECK((a_num_segments >= 1LL), ValueError) + << KERNEL_NAME << ": requires (num_segments >= 1)"; + TVM_FFI_CHECK((a_h_kv >= 1LL), ValueError) + << KERNEL_NAME << ": requires (h_kv >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") is not a multiple of the cluster shape (2, 1, 1)"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_gqa_train_1afdfdde603ff8806b4d), cudaFuncAttributeMaxDynamicSharedMemorySize, 231424), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_Q = EncodeTma_Q(a_Q); + CUtensorMap v_dO = EncodeTma_dO(a_dO); + CUtensorMap v_K = EncodeTma_K(a_K); + CUtensorMap v_V = EncodeTma_V(a_V); + void* v_lse_log2 = DataPtr(a_lse_log2); + void* v_dpsum = DataPtr(a_dpsum); + void* v_dq_accum = DataPtr(a_dq_accum); + void* v_cu_seqlens_q = DataPtr(a_cu_seqlens_q); + void* v_cu_seqlens_k = DataPtr(a_cu_seqlens_k); + void* v_dk_out = DataPtr(a_dk_out); + void* v_dv_out = DataPtr(a_dv_out); + int32_t v_num_segments = static_cast(a_num_segments); + int32_t v_h_kv = static_cast(a_h_kv); + int32_t v_group = static_cast(a_group); + int32_t v_rows_pad = static_cast(a_rows_pad); + float v_softmax_scale = static_cast(a_softmax_scale); + float v_softmax_scale_log2 = static_cast(a_softmax_scale_log2); + void* v_row_map = a_row_map.has_value() ? DataPtr(a_row_map.value()) : nullptr; + void* kernel_args[] = {&v_Q, &v_dO, &v_K, &v_V, &v_lse_log2, &v_dpsum, &v_dq_accum, &v_cu_seqlens_q, &v_cu_seqlens_k, &v_dk_out, &v_dv_out, &v_num_segments, &v_h_kv, &v_group, &v_rows_pad, &v_softmax_scale, &v_softmax_scale_log2, &v_row_map}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(512, 1, 1); + config.dynamicSmemBytes = 231424; + config.stream = stream; + cudaLaunchAttribute attrs[2] = {}; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = 2; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + config.attrs = attrs; + config.numAttrs = 2; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gqa_train_1afdfdde603ff8806b4d), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gqa_train_1afdfdde603ff8806b4d + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gqa_train_1afdfdde603ff8806b4d::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu new file mode 100644 index 0000000..d879f47 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu @@ -0,0 +1,1609 @@ +/* + * 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 TMEM_NCOLS 512 +#define TMEM_DV_OFFSET 0 +#define TMEM_DK_OFFSET 128 +#define TMEM_S_OFFSET 256 +#define TMEM_DQ_OFFSET 384 +#define TMEM_DP_OFFSET 320 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_DO_PIPE_STAGES 1 +#define SMEM_SQ_OFF 1024 +#define SMEM_SQ_STAGE_BYTES 32768 +#define SMEM_SQ_STRIDE 32768 +#define SMEM_SQT_OFF 1024 +#define SMEM_SQT_STAGE_BYTES 32768 +#define SMEM_SQT_STRIDE 32768 +#define SMEM_SDO_OFF 33792 +#define SMEM_SDO_STAGE_BYTES 32768 +#define SMEM_SDO_STRIDE 32768 +#define SMEM_SDOT_OFF 33792 +#define SMEM_SDOT_STAGE_BYTES 32768 +#define SMEM_SDOT_STRIDE 32768 +#define SMEM_SK_OFF 66560 +#define SMEM_SK_STAGE_BYTES 32768 +#define SMEM_SK_STRIDE 32768 +#define SMEM_SV_OFF 99328 +#define SMEM_SV_STAGE_BYTES 32768 +#define SMEM_SV_STRIDE 32768 +#define SMEM_SKT_OFF 132096 +#define SMEM_SKT_STAGE_BYTES 32768 +#define SMEM_SKT_STRIDE 32768 +#define SMEM_SP_OFF 164864 +#define SMEM_SP_STAGE_BYTES 16384 +#define SMEM_SP_STRIDE 16384 +#define SMEM_SDST_OFF 181248 +#define SMEM_SDST_STAGE_BYTES 16384 +#define SMEM_SDST_STRIDE 16384 +#define SMEM_SDS_OFF 197632 +#define SMEM_SDS_STAGE_BYTES 16384 +#define SMEM_SDS_STRIDE 16384 +#define SMEM_SDQ_OFF 214016 +#define SMEM_SDQ_STAGE_BYTES 8192 +#define SMEM_SDQ_STRIDE 8192 +#define SMEM_SLSE_OFF 230400 +#define SMEM_SLSE_STAGE_BYTES 512 +#define SMEM_SLSE_STRIDE 512 +#define SMEM_SDPS_OFF 230912 +#define SMEM_SDPS_STAGE_BYTES 512 +#define SMEM_SDPS_STRIDE 512 +#define SMEM_TOTAL 231424 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, LAUNCH_MIN_BLOCKS) __cluster_dims__(2,1,1) void +kernel_cake_gqa_train_1afdfdde603ff8806b4d(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap dO, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, float* __restrict__ lse_log2, float* __restrict__ dpsum, float* __restrict__ dq_accum, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, __nv_bfloat16* __restrict__ dk_out, __nv_bfloat16* __restrict__ dv_out, int num_segments, int h_kv, int group, int rows_pad, float softmax_scale, float softmax_scale_log2, int* __restrict__ row_map) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define kv_full_addr (mbar_base + 0) + #define q_full_addr (mbar_base + 8) + #define q_empty_addr (mbar_base + 16) + #define do_full_addr (mbar_base + 24) + #define do_empty_addr (mbar_base + 32) + #define lse_full_addr (mbar_base + 40) + #define lse_empty_addr (mbar_base + 48) + #define dpsum_full_addr (mbar_base + 56) + #define dpsum_empty_addr (mbar_base + 64) + #define s_full_addr (mbar_base + 72) + #define dp_full_addr (mbar_base + 80) + #define ds_full_addr (mbar_base + 88) + #define ds_empty_addr (mbar_base + 96) + #define xchg_full_addr (mbar_base + 104) + #define xchg_done_addr (mbar_base + 112) + #define dq_full_addr (mbar_base + 120) + #define dq_empty_addr (mbar_base + 128) + #define dqlo_empty_addr (mbar_base + 136) + #define dkv_full_addr (mbar_base + 144) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* sq = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sq_addr = smem + 1024; + __nv_bfloat16* sqt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sqt_addr = smem + 1024; + __nv_bfloat16* sdo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int sdo_addr = smem + 33792; + __nv_bfloat16* sdot = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int sdot_addr = smem + 33792; + __nv_bfloat16* sk = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int sk_addr = smem + 66560; + __nv_bfloat16* sv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int sv_addr = smem + 99328; + __nv_bfloat16* skt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); + const int skt_addr = smem + 132096; + __nv_bfloat16* sp = reinterpret_cast<__nv_bfloat16*>(smem_raw + 164864); + const int sp_addr = smem + 164864; + __nv_bfloat16* sdst = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); + const int sdst_addr = smem + 181248; + __nv_bfloat16* sds = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int sds_addr = smem + 197632; + float* sdq = reinterpret_cast(smem_raw + 214016); + const int sdq_addr = smem + 214016; + float* slse = reinterpret_cast(smem_raw + 230400); + const int slse_addr = smem + 230400; + float* sdps = reinterpret_cast(smem_raw + 230912); + const int sdps_addr = smem + 230912; + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&Q))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&dO))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&K))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&V))) : "memory"); } + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 19 barriers) + // Mbarriers at smem_raw[0..152) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // kv_full: 1 barriers, init_count=2 + mbarrier_init(smem + 0, 2); + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=2 + mbarrier_init(smem + 8, 2); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // --- pipeline 'do_pipe' --- + // do_full: 1 barriers, init_count=2 + mbarrier_init(smem + 24, 2); + // do_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + // lse_full: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // lse_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 48, 256); + // dpsum_full: 1 barriers, init_count=1 + mbarrier_init(smem + 56, 1); + // dpsum_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 64, 256); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // dp_full: 1 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + // ds_full: 1 barriers, init_count=512 + mbarrier_init(smem + 88, 512); + // ds_empty: 1 barriers, init_count=2 + mbarrier_init(smem + 96, 2); + // xchg_full: 1 barriers, init_count=1 + mbarrier_init(smem + 104, 1); + // xchg_done: 1 barriers, init_count=2 + mbarrier_init(smem + 112, 2); + // dq_full: 1 barriers, init_count=1 + mbarrier_init(smem + 120, 1); + // dq_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 128, 256); + // dqlo_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 136, 256); + // dkv_full: 1 barriers, init_count=1 + mbarrier_init(smem + 144, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 152); + if (warp == 12) { + int _tmem_hold = smem + 152; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_dv = taddr; + const int tmem_dk = taddr + 128; + const int tmem_s = taddr + 256; + const int tmem_dq = taddr + 384; + const int tmem_dp = taddr + 320; + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); + } + + // ---- Role: reduce ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // reduce_main + int seg = blockIdx.z; + int n_blk = blockIdx.x >> 1; + int q0 = cu_seqlens_q[seg]; + int q_len = cu_seqlens_q[seg + 1] - q0; + int k0 = cu_seqlens_k[seg]; + int k_len = cu_seqlens_k[seg + 1] - k0; + int found = ((n_blk * 128 < k_len) ? 1 : 0); + int off = k_len - q_len; + int first_key = n_blk * 128 - off; + int m_min = ((first_key > 0) ? first_key : 0) >> 6; + int m_max = q_len + 63 >> 6; + int num_m = ((found != 0) ? m_max - m_min : 0); + int row_off = q0 * 2 + seg * 128 >> 7 << 7; + int rank = cta_rank & 1; + int h = blockIdx.y; + int n_groups = group >> 1; + const int quarter = warp % 4; + const int tl = quarter * 32 + lane; + const unsigned int lane_base = (unsigned int)(quarter * 32 << 16); + int t = 0; + #pragma unroll 1 + for (int g = 0; g < n_groups; g++) { + long long head_rows = (long long)((h * n_groups + g) * rows_pad + row_off); + #pragma unroll 1 + for (int mi = 0; mi < num_m; mi++) { + long long tile = (head_rows + (long long)((m_min + mi) * 128)) * 256 + (long long)(rank * 8 * 2048); + int _parity_dq_full = (int)((unsigned int)t % 2); + mbarrier_wait(dq_full_addr, _parity_dq_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + float dq_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_0[0]), "=f"(dq_0[1]), "=f"(dq_0[2]), "=f"(dq_0[3]), "=f"(dq_0[4]), "=f"(dq_0[5]), "=f"(dq_0[6]), "=f"(dq_0[7]), "=f"(dq_0[8]), "=f"(dq_0[9]), "=f"(dq_0[10]), "=f"(dq_0[11]), "=f"(dq_0[12]), "=f"(dq_0[13]), "=f"(dq_0[14]), "=f"(dq_0[15]), "=f"(dq_0[16]), "=f"(dq_0[17]), "=f"(dq_0[18]), "=f"(dq_0[19]), "=f"(dq_0[20]), "=f"(dq_0[21]), "=f"(dq_0[22]), "=f"(dq_0[23]), "=f"(dq_0[24]), "=f"(dq_0[25]), "=f"(dq_0[26]), "=f"(dq_0[27]), "=f"(dq_0[28]), "=f"(dq_0[29]), "=f"(dq_0[30]), "=f"(dq_0[31]) + : "r"(taddr + 448 + lane_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_0[32]), "=f"(dq_0[33]), "=f"(dq_0[34]), "=f"(dq_0[35]), "=f"(dq_0[36]), "=f"(dq_0[37]), "=f"(dq_0[38]), "=f"(dq_0[39]), "=f"(dq_0[40]), "=f"(dq_0[41]), "=f"(dq_0[42]), "=f"(dq_0[43]), "=f"(dq_0[44]), "=f"(dq_0[45]), "=f"(dq_0[46]), "=f"(dq_0[47]), "=f"(dq_0[48]), "=f"(dq_0[49]), "=f"(dq_0[50]), "=f"(dq_0[51]), "=f"(dq_0[52]), "=f"(dq_0[53]), "=f"(dq_0[54]), "=f"(dq_0[55]), "=f"(dq_0[56]), "=f"(dq_0[57]), "=f"(dq_0[58]), "=f"(dq_0[59]), "=f"(dq_0[60]), "=f"(dq_0[61]), "=f"(dq_0[62]), "=f"(dq_0[63]) + : "r"(taddr + 480 + lane_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((dq_empty_addr) & 0xFEFFFFFF) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_0[0]), "f"(dq_0[1]), "f"(dq_0[2]), "f"(dq_0[3]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_0[4]), "f"(dq_0[5]), "f"(dq_0[6]), "f"(dq_0[7]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_0[8]), "f"(dq_0[9]), "f"(dq_0[10]), "f"(dq_0[11]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_0[12]), "f"(dq_0[13]), "f"(dq_0[14]), "f"(dq_0[15]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_0 = reinterpret_cast(dq_accum + (tile + 8192)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_0), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_0[16]), "f"(dq_0[17]), "f"(dq_0[18]), "f"(dq_0[19]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_0[20]), "f"(dq_0[21]), "f"(dq_0[22]), "f"(dq_0[23]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_0[24]), "f"(dq_0[25]), "f"(dq_0[26]), "f"(dq_0[27]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_0[28]), "f"(dq_0[29]), "f"(dq_0[30]), "f"(dq_0[31]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_1 = reinterpret_cast(dq_accum + (tile + 10240)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_1), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_0[32]), "f"(dq_0[33]), "f"(dq_0[34]), "f"(dq_0[35]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_0[36]), "f"(dq_0[37]), "f"(dq_0[38]), "f"(dq_0[39]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_0[40]), "f"(dq_0[41]), "f"(dq_0[42]), "f"(dq_0[43]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_0[44]), "f"(dq_0[45]), "f"(dq_0[46]), "f"(dq_0[47]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_2 = reinterpret_cast(dq_accum + (tile + 12288)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_2), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_0[48]), "f"(dq_0[49]), "f"(dq_0[50]), "f"(dq_0[51]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_0[52]), "f"(dq_0[53]), "f"(dq_0[54]), "f"(dq_0[55]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_0[56]), "f"(dq_0[57]), "f"(dq_0[58]), "f"(dq_0[59]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_0[60]), "f"(dq_0[61]), "f"(dq_0[62]), "f"(dq_0[63]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_3 = reinterpret_cast(dq_accum + (tile + 14336)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_3), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float dq_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_1[0]), "=f"(dq_1[1]), "=f"(dq_1[2]), "=f"(dq_1[3]), "=f"(dq_1[4]), "=f"(dq_1[5]), "=f"(dq_1[6]), "=f"(dq_1[7]), "=f"(dq_1[8]), "=f"(dq_1[9]), "=f"(dq_1[10]), "=f"(dq_1[11]), "=f"(dq_1[12]), "=f"(dq_1[13]), "=f"(dq_1[14]), "=f"(dq_1[15]), "=f"(dq_1[16]), "=f"(dq_1[17]), "=f"(dq_1[18]), "=f"(dq_1[19]), "=f"(dq_1[20]), "=f"(dq_1[21]), "=f"(dq_1[22]), "=f"(dq_1[23]), "=f"(dq_1[24]), "=f"(dq_1[25]), "=f"(dq_1[26]), "=f"(dq_1[27]), "=f"(dq_1[28]), "=f"(dq_1[29]), "=f"(dq_1[30]), "=f"(dq_1[31]) + : "r"(taddr + 384 + lane_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_1[32]), "=f"(dq_1[33]), "=f"(dq_1[34]), "=f"(dq_1[35]), "=f"(dq_1[36]), "=f"(dq_1[37]), "=f"(dq_1[38]), "=f"(dq_1[39]), "=f"(dq_1[40]), "=f"(dq_1[41]), "=f"(dq_1[42]), "=f"(dq_1[43]), "=f"(dq_1[44]), "=f"(dq_1[45]), "=f"(dq_1[46]), "=f"(dq_1[47]), "=f"(dq_1[48]), "=f"(dq_1[49]), "=f"(dq_1[50]), "=f"(dq_1[51]), "=f"(dq_1[52]), "=f"(dq_1[53]), "=f"(dq_1[54]), "=f"(dq_1[55]), "=f"(dq_1[56]), "=f"(dq_1[57]), "=f"(dq_1[58]), "=f"(dq_1[59]), "=f"(dq_1[60]), "=f"(dq_1[61]), "=f"(dq_1[62]), "=f"(dq_1[63]) + : "r"(taddr + 416 + lane_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((dqlo_empty_addr) & 0xFEFFFFFF) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_1[0]), "f"(dq_1[1]), "f"(dq_1[2]), "f"(dq_1[3]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_1[4]), "f"(dq_1[5]), "f"(dq_1[6]), "f"(dq_1[7]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_1[8]), "f"(dq_1[9]), "f"(dq_1[10]), "f"(dq_1[11]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_1[12]), "f"(dq_1[13]), "f"(dq_1[14]), "f"(dq_1[15]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_4 = reinterpret_cast(dq_accum + (tile + 0)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_4), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_1[16]), "f"(dq_1[17]), "f"(dq_1[18]), "f"(dq_1[19]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_1[20]), "f"(dq_1[21]), "f"(dq_1[22]), "f"(dq_1[23]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_1[24]), "f"(dq_1[25]), "f"(dq_1[26]), "f"(dq_1[27]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_1[28]), "f"(dq_1[29]), "f"(dq_1[30]), "f"(dq_1[31]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_5 = reinterpret_cast(dq_accum + (tile + 2048)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_5), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_1[32]), "f"(dq_1[33]), "f"(dq_1[34]), "f"(dq_1[35]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_1[36]), "f"(dq_1[37]), "f"(dq_1[38]), "f"(dq_1[39]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_1[40]), "f"(dq_1[41]), "f"(dq_1[42]), "f"(dq_1[43]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_1[44]), "f"(dq_1[45]), "f"(dq_1[46]), "f"(dq_1[47]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_6 = reinterpret_cast(dq_accum + (tile + 4096)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_6), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_1[48]), "f"(dq_1[49]), "f"(dq_1[50]), "f"(dq_1[51]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_1[52]), "f"(dq_1[53]), "f"(dq_1[54]), "f"(dq_1[55]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_1[56]), "f"(dq_1[57]), "f"(dq_1[58]), "f"(dq_1[59]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_1[60]), "f"(dq_1[61]), "f"(dq_1[62]), "f"(dq_1[63]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_7 = reinterpret_cast(dq_accum + (tile + 6144)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_7), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + t += 1; + } + } + if (warp == 0) { + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + } + } + } + } + // ---- Role: compute ---- + if (warp >= 4 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + { // compute_main + int seg_1 = blockIdx.z; + int n_blk_1 = blockIdx.x >> 1; + int q0_1 = cu_seqlens_q[seg_1]; + int q_len_1 = cu_seqlens_q[seg_1 + 1] - q0_1; + int k0_1 = cu_seqlens_k[seg_1]; + int k_len_1 = cu_seqlens_k[seg_1 + 1] - k0_1; + int found_1 = ((n_blk_1 * 128 < k_len_1) ? 1 : 0); + int off_1 = k_len_1 - q_len_1; + int first_key_1 = n_blk_1 * 128 - off_1; + int m_min_1 = ((first_key_1 > 0) ? first_key_1 : 0) >> 6; + int m_max_1 = q_len_1 + 63 >> 6; + int num_m_1 = ((found_1 != 0) ? m_max_1 - m_min_1 : 0); + int row_off_1 = q0_1 * 2 + seg_1 * 128 >> 7 << 7; + int rank_1 = cta_rank & 1; + int h_1 = blockIdx.y; + int n_groups_1 = group >> 1; + const int quarter_1 = warp % 4; + const int wg = warp - 4 >> 2; + const int row = (quarter_1 & 1) * 32 + lane; + const int half = quarter_1 >> 1; + const int col0 = half * 64 + wg * 32; + const unsigned int lane_base_1 = (unsigned int)(quarter_1 * 32 << 16); + int key = n_blk_1 * 128 + rank_1 * 64 + row; + int prow = half * 64 + row; + int t_1 = 0; + #pragma unroll 1 + for (int g_1 = 0; g_1 < n_groups_1; g_1++) { + #pragma unroll 1 + for (int mi_1 = 0; mi_1 < num_m_1; mi_1++) { + int m = m_min_1 + mi_1; + int _parity_lse_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(lse_full_addr, _parity_lse_full); + int _parity_s_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(s_full_addr, _parity_s_full); + int _parity_ds_empty = (int)((unsigned int)t_1 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); + float s[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(s[0]), "=f"(s[1]), "=f"(s[2]), "=f"(s[3]), "=f"(s[4]), "=f"(s[5]), "=f"(s[6]), "=f"(s[7]), "=f"(s[8]), "=f"(s[9]), "=f"(s[10]), "=f"(s[11]), "=f"(s[12]), "=f"(s[13]), "=f"(s[14]), "=f"(s[15]), "=f"(s[16]), "=f"(s[17]), "=f"(s[18]), "=f"(s[19]), "=f"(s[20]), "=f"(s[21]), "=f"(s[22]), "=f"(s[23]), "=f"(s[24]), "=f"(s[25]), "=f"(s[26]), "=f"(s[27]), "=f"(s[28]), "=f"(s[29]), "=f"(s[30]), "=f"(s[31]) + : "r"(taddr + (unsigned int)(256 + wg * 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int lo = (key - off_1 - m * 64) * 2 - col0; + int hi = (q_len_1 - m * 64) * 2 - col0; + int hi_k = ((key < k_len_1) ? hi : 0); + uint32_t _thresh_mask_0; + { + int _lim_1 = hi_k; + if (_lim_1 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_1)); + } + } + uint32_t _thresh_mask_2; + { + int _lim_3 = lo; + if (_lim_3 <= 0) { _thresh_mask_2 = 0u; } + else if (_lim_3 >= 32) { _thresh_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_2) : "r"(_lim_3)); + } + } + unsigned int keep = _thresh_mask_0 & ~_thresh_mask_2; + if (!(keep & (1u << 0))) s[0] = -CUDART_INF_F; + if (!(keep & (1u << 1))) s[1] = -CUDART_INF_F; + if (!(keep & (1u << 2))) s[2] = -CUDART_INF_F; + if (!(keep & (1u << 3))) s[3] = -CUDART_INF_F; + if (!(keep & (1u << 4))) s[4] = -CUDART_INF_F; + if (!(keep & (1u << 5))) s[5] = -CUDART_INF_F; + if (!(keep & (1u << 6))) s[6] = -CUDART_INF_F; + if (!(keep & (1u << 7))) s[7] = -CUDART_INF_F; + if (!(keep & (1u << 8))) s[8] = -CUDART_INF_F; + if (!(keep & (1u << 9))) s[9] = -CUDART_INF_F; + if (!(keep & (1u << 10))) s[10] = -CUDART_INF_F; + if (!(keep & (1u << 11))) s[11] = -CUDART_INF_F; + if (!(keep & (1u << 12))) s[12] = -CUDART_INF_F; + if (!(keep & (1u << 13))) s[13] = -CUDART_INF_F; + if (!(keep & (1u << 14))) s[14] = -CUDART_INF_F; + if (!(keep & (1u << 15))) s[15] = -CUDART_INF_F; + if (!(keep & (1u << 16))) s[16] = -CUDART_INF_F; + if (!(keep & (1u << 17))) s[17] = -CUDART_INF_F; + if (!(keep & (1u << 18))) s[18] = -CUDART_INF_F; + if (!(keep & (1u << 19))) s[19] = -CUDART_INF_F; + if (!(keep & (1u << 20))) s[20] = -CUDART_INF_F; + if (!(keep & (1u << 21))) s[21] = -CUDART_INF_F; + if (!(keep & (1u << 22))) s[22] = -CUDART_INF_F; + if (!(keep & (1u << 23))) s[23] = -CUDART_INF_F; + if (!(keep & (1u << 24))) s[24] = -CUDART_INF_F; + if (!(keep & (1u << 25))) s[25] = -CUDART_INF_F; + if (!(keep & (1u << 26))) s[26] = -CUDART_INF_F; + if (!(keep & (1u << 27))) s[27] = -CUDART_INF_F; + if (!(keep & (1u << 28))) s[28] = -CUDART_INF_F; + if (!(keep & (1u << 29))) s[29] = -CUDART_INF_F; + if (!(keep & (1u << 30))) s[30] = -CUDART_INF_F; + if (!(keep & (1u << 31))) s[31] = -CUDART_INF_F; + float lse[32]; + #pragma unroll + for (int j = 0; j < 32; j++) { + lse[j] = slse[col0 + j]; + } + float2 scale2 = make_float2(softmax_scale_log2, softmax_scale_log2); + #pragma unroll + for (int pair = 0; pair < 16; pair++) { + float2 x = fma_sub_f32x2_rn_ftz(make_float2(s[pair * 2], s[pair * 2 + 1]), scale2, make_float2(lse[pair * 2], lse[pair * 2 + 1])); + s[pair * 2] = x.x; + s[pair * 2 + 1] = x.y; + } + unsigned int p16[16]; + softmax_frag_exp2_cast(&s[0], p16, 0); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 0 ^ prow & 7) * 16))), "r"(p16[0]), "r"(p16[1]), "r"(p16[2]), "r"(p16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 1 ^ prow & 7) * 16))), "r"(p16[4]), "r"(p16[5]), "r"(p16[6]), "r"(p16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 2 ^ prow & 7) * 16))), "r"(p16[8]), "r"(p16[9]), "r"(p16[10]), "r"(p16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 3 ^ prow & 7) * 16))), "r"(p16[12]), "r"(p16[13]), "r"(p16[14]), "r"(p16[15]) : "memory"); + mbarrier_arrive(lse_empty_addr); + int _parity_dpsum_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(dpsum_full_addr, _parity_dpsum_full); + int _parity_dp_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(dp_full_addr, _parity_dp_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + float dp[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dp[0]), "=f"(dp[1]), "=f"(dp[2]), "=f"(dp[3]), "=f"(dp[4]), "=f"(dp[5]), "=f"(dp[6]), "=f"(dp[7]), "=f"(dp[8]), "=f"(dp[9]), "=f"(dp[10]), "=f"(dp[11]), "=f"(dp[12]), "=f"(dp[13]), "=f"(dp[14]), "=f"(dp[15]), "=f"(dp[16]), "=f"(dp[17]), "=f"(dp[18]), "=f"(dp[19]), "=f"(dp[20]), "=f"(dp[21]), "=f"(dp[22]), "=f"(dp[23]), "=f"(dp[24]), "=f"(dp[25]), "=f"(dp[26]), "=f"(dp[27]), "=f"(dp[28]), "=f"(dp[29]), "=f"(dp[30]), "=f"(dp[31]) + : "r"(taddr + (unsigned int)(320 + wg * 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float dps[32]; + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1++) { + dps[j_1] = sdps[col0 + j_1]; + } + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + sub_f32x2_inplace(&reinterpret_cast(dp)[_ls], reinterpret_cast(dps)[_ls]); + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(dp)[_ls], reinterpret_cast(s)[_ls]); + unsigned int ds16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dp[_lp*2 + 0], dp[_lp*2+1 + 0])); + ds16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 0 ^ prow & 7) * 16))), "r"(ds16[0]), "r"(ds16[1]), "r"(ds16[2]), "r"(ds16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 1 ^ prow & 7) * 16))), "r"(ds16[4]), "r"(ds16[5]), "r"(ds16[6]), "r"(ds16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 2 ^ prow & 7) * 16))), "r"(ds16[8]), "r"(ds16[9]), "r"(ds16[10]), "r"(ds16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 3 ^ prow & 7) * 16))), "r"(ds16[12]), "r"(ds16[13]), "r"(ds16[14]), "r"(ds16[15]) : "memory"); + if (half == rank_1) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 0 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[0]), "r"(ds16[1]), "r"(ds16[2]), "r"(ds16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 1 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[4]), "r"(ds16[5]), "r"(ds16[6]), "r"(ds16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 2 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[8]), "r"(ds16[9]), "r"(ds16[10]), "r"(ds16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 3 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[12]), "r"(ds16[13]), "r"(ds16[14]), "r"(ds16[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((ds_full_addr) & 0xFEFFFFFF) : "memory"); + mbarrier_arrive(dpsum_empty_addr); + asm volatile("barrier.sync 1, 256;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(xchg_full_addr), "r"(1 - rank_1), "r"((uint32_t)(8192)) : "memory"); + uint32_t _mbar_peer_4; + asm( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mbar_peer_4) : "r"(xchg_full_addr), "r"(1 - rank_1)); + uint32_t _cluster_mapa_5; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_cluster_mapa_5) : "r"(sds_addr + (unsigned int)(rank_1 * 8192)), "r"(1 - rank_1)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_cluster_mapa_5), "r"(sdst_addr + (unsigned int)((1 - rank_1) * 8192)), "r"((uint32_t)(8192)), "r"(_mbar_peer_4) + : "memory"); + } + } + t_1 += 1; + } + } + int _parity_ds_empty_1 = (int)((unsigned int)t_1 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty_1, 10000000); + mbarrier_wait(dkv_full_addr, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int hd0 = half * 128 + wg * 64; + int store_row = ((key < k_len_1) ? 1 : 0); + if (found_1 == 0) { + store_row = 0; + } + long long out_row = (long long)(k0_1 + key); + float acc_0[64]; + if (num_m_1 > 0) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_0[0]), "=f"(acc_0[1]), "=f"(acc_0[2]), "=f"(acc_0[3]), "=f"(acc_0[4]), "=f"(acc_0[5]), "=f"(acc_0[6]), "=f"(acc_0[7]), "=f"(acc_0[8]), "=f"(acc_0[9]), "=f"(acc_0[10]), "=f"(acc_0[11]), "=f"(acc_0[12]), "=f"(acc_0[13]), "=f"(acc_0[14]), "=f"(acc_0[15]), "=f"(acc_0[16]), "=f"(acc_0[17]), "=f"(acc_0[18]), "=f"(acc_0[19]), "=f"(acc_0[20]), "=f"(acc_0[21]), "=f"(acc_0[22]), "=f"(acc_0[23]), "=f"(acc_0[24]), "=f"(acc_0[25]), "=f"(acc_0[26]), "=f"(acc_0[27]), "=f"(acc_0[28]), "=f"(acc_0[29]), "=f"(acc_0[30]), "=f"(acc_0[31]) + : "r"(taddr + (unsigned int)(0 + wg * 64) + lane_base_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_0[32]), "=f"(acc_0[33]), "=f"(acc_0[34]), "=f"(acc_0[35]), "=f"(acc_0[36]), "=f"(acc_0[37]), "=f"(acc_0[38]), "=f"(acc_0[39]), "=f"(acc_0[40]), "=f"(acc_0[41]), "=f"(acc_0[42]), "=f"(acc_0[43]), "=f"(acc_0[44]), "=f"(acc_0[45]), "=f"(acc_0[46]), "=f"(acc_0[47]), "=f"(acc_0[48]), "=f"(acc_0[49]), "=f"(acc_0[50]), "=f"(acc_0[51]), "=f"(acc_0[52]), "=f"(acc_0[53]), "=f"(acc_0[54]), "=f"(acc_0[55]), "=f"(acc_0[56]), "=f"(acc_0[57]), "=f"(acc_0[58]), "=f"(acc_0[59]), "=f"(acc_0[60]), "=f"(acc_0[61]), "=f"(acc_0[62]), "=f"(acc_0[63]) + : "r"(taddr + (unsigned int)(0 + wg * 64 + 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } else { + #pragma unroll + for (int i_0 = 0; i_0 < 64; i_0++) { + acc_0[i_0] = 0.0f; + } + } + long long base_0 = (out_row * (long long)h_kv + (long long)h_1) * 256 + (long long)hd0; + if (store_row != 0) { + unsigned int packed_0[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(acc_0[_lp*2 + 0], acc_0[_lp*2+1 + 0])); + packed_0[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int i_1 = 0; i_1 < 8; i_1++) { + reinterpret_cast(dv_out + (base_0 + (long long)(i_1 * 8)))[0] = reinterpret_cast(packed_0 + (i_1 * 4))[0]; + } + } + float acc_1[64]; + if (num_m_1 > 0) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_1[0]), "=f"(acc_1[1]), "=f"(acc_1[2]), "=f"(acc_1[3]), "=f"(acc_1[4]), "=f"(acc_1[5]), "=f"(acc_1[6]), "=f"(acc_1[7]), "=f"(acc_1[8]), "=f"(acc_1[9]), "=f"(acc_1[10]), "=f"(acc_1[11]), "=f"(acc_1[12]), "=f"(acc_1[13]), "=f"(acc_1[14]), "=f"(acc_1[15]), "=f"(acc_1[16]), "=f"(acc_1[17]), "=f"(acc_1[18]), "=f"(acc_1[19]), "=f"(acc_1[20]), "=f"(acc_1[21]), "=f"(acc_1[22]), "=f"(acc_1[23]), "=f"(acc_1[24]), "=f"(acc_1[25]), "=f"(acc_1[26]), "=f"(acc_1[27]), "=f"(acc_1[28]), "=f"(acc_1[29]), "=f"(acc_1[30]), "=f"(acc_1[31]) + : "r"(taddr + (unsigned int)(128 + wg * 64) + lane_base_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_1[32]), "=f"(acc_1[33]), "=f"(acc_1[34]), "=f"(acc_1[35]), "=f"(acc_1[36]), "=f"(acc_1[37]), "=f"(acc_1[38]), "=f"(acc_1[39]), "=f"(acc_1[40]), "=f"(acc_1[41]), "=f"(acc_1[42]), "=f"(acc_1[43]), "=f"(acc_1[44]), "=f"(acc_1[45]), "=f"(acc_1[46]), "=f"(acc_1[47]), "=f"(acc_1[48]), "=f"(acc_1[49]), "=f"(acc_1[50]), "=f"(acc_1[51]), "=f"(acc_1[52]), "=f"(acc_1[53]), "=f"(acc_1[54]), "=f"(acc_1[55]), "=f"(acc_1[56]), "=f"(acc_1[57]), "=f"(acc_1[58]), "=f"(acc_1[59]), "=f"(acc_1[60]), "=f"(acc_1[61]), "=f"(acc_1[62]), "=f"(acc_1[63]) + : "r"(taddr + (unsigned int)(128 + wg * 64 + 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } else { + #pragma unroll + for (int i_2 = 0; i_2 < 64; i_2++) { + acc_1[i_2] = 0.0f; + } + } + const float2 _scale2_6 = {softmax_scale, softmax_scale}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(acc_1)[_ls], _scale2_6); + long long base_1 = (out_row * (long long)h_kv + (long long)h_1) * 256 + (long long)hd0; + if (store_row != 0) { + unsigned int packed_1[32]; + #pragma unroll + for (int _lp = 0; _lp < 32; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(acc_1[_lp*2 + 0], acc_1[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int i_3 = 0; i_3 < 8; i_3++) { + reinterpret_cast(dk_out + (base_1 + (long long)(i_3 * 8)))[0] = reinterpret_cast(packed_1 + (i_3 * 4))[0]; + } + } + } + } + // ---- Role: mma ---- + if (warp == 12) { + { // mma_main + int seg_2 = blockIdx.z; + int n_blk_2 = blockIdx.x >> 1; + int q0_2 = cu_seqlens_q[seg_2]; + int q_len_2 = cu_seqlens_q[seg_2 + 1] - q0_2; + int k0_2 = cu_seqlens_k[seg_2]; + int k_len_2 = cu_seqlens_k[seg_2 + 1] - k0_2; + int found_2 = ((n_blk_2 * 128 < k_len_2) ? 1 : 0); + int off_2 = k_len_2 - q_len_2; + int first_key_2 = n_blk_2 * 128 - off_2; + int m_min_2 = ((first_key_2 > 0) ? first_key_2 : 0) >> 6; + int m_max_2 = q_len_2 + 63 >> 6; + int num_m_2 = ((found_2 != 0) ? m_max_2 - m_min_2 : 0); + int row_off_2 = q0_2 * 2 + seg_2 * 128 >> 7 << 7; + int n_groups_2 = group >> 1; + int t_2 = 0; + unsigned int _phase_q_full = 0; + unsigned int _phase_do_full = 0; + if ((cta_rank & 1) == 0) { + if (elect_sync()) { + mbarrier_wait_cluster_hint(kv_full_addr, 0, 10000000); + #pragma unroll 1 + for (int g_2 = 0; g_2 < n_groups_2; g_2++) { + #pragma unroll 1 + for (int mi_2 = 0; mi_2 < num_m_2; mi_2++) { + int first = ((t_2 == 0) ? 1 : 0); + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = ((sk_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_1 = ((sq_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_0 = ((sk_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_1 = ((sq_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_s), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(s_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + mbarrier_wait_cluster_hint(do_full_addr, _phase_do_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = ((sv_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_3 = ((sdo_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_2 = ((sv_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_3 = ((sdo_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_dp), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(dp_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(do_empty_addr, (uint16_t)(3)); + _phase_do_full ^= 1; + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + int _parity_ds_full = (int)((unsigned int)t_2 % 2); + mbarrier_wait_cluster_hint(ds_full_addr, _parity_ds_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_4 = ((sdst_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_5 = (((sqt_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_4 = ((sdst_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_5 = (((sqt_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_4), "r"(_smem_desc_lo_5), "r"(tmem_dk), "r"(((first) ? 0 : 1)) : "memory"); + tcgen05_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + mbarrier_wait_cluster_hint(do_full_addr, _phase_do_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_6 = ((sp_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_7 = (((sdot_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_6 = ((sp_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_7 = (((sdot_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_6), "r"(_smem_desc_lo_7), "r"(tmem_dv), "r"(((first) ? 0 : 1)) : "memory"); + tcgen05_commit_cg2_multicast(do_empty_addr, (uint16_t)(3)); + _phase_do_full ^= 1; + int _parity_xchg_done = (int)((unsigned int)t_2 % 2); + mbarrier_wait_cluster_hint(xchg_done_addr, _parity_xchg_done, 10000000); + int _parity_ds_empty_2 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty_2, 10000000); + int _parity_dq_empty = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dq_empty_addr, _parity_dq_empty, 10000000); + int _parity_dqlo_empty = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dqlo_empty_addr, _parity_dqlo_empty, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_8 = (((sds_addr) >> 4) & 0x7FFF) | 0x4000000; + int _smem_desc_lo_9 = (((skt_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_8 = (((sds_addr) >> 4) & 0x3FFF) | 0x4000000; + int _smem_desc_lo_9 = (((skt_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138511504;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_8), "r"(_smem_desc_lo_9), "r"(tmem_dq), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(dq_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(ds_empty_addr, (uint16_t)(3)); + t_2 += 1; + } + } + tcgen05_commit_cg2_multicast(dkv_full_addr, (uint16_t)(3)); + int _parity_dq_empty_1 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dq_empty_addr, _parity_dq_empty_1, 10000000); + int _parity_dqlo_empty_1 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dqlo_empty_addr, _parity_dqlo_empty_1, 10000000); + } + } + } + } + // ---- Role: load ---- + if (warp == 13) { + { // load_main + int seg_3 = blockIdx.z; + int n_blk_3 = blockIdx.x >> 1; + int q0_3 = cu_seqlens_q[seg_3]; + int q_len_3 = cu_seqlens_q[seg_3 + 1] - q0_3; + int k0_3 = cu_seqlens_k[seg_3]; + int k_len_3 = cu_seqlens_k[seg_3 + 1] - k0_3; + int found_3 = ((n_blk_3 * 128 < k_len_3) ? 1 : 0); + int off_3 = k_len_3 - q_len_3; + int first_key_3 = n_blk_3 * 128 - off_3; + int m_min_3 = ((first_key_3 > 0) ? first_key_3 : 0) >> 6; + int m_max_3 = q_len_3 + 63 >> 6; + int num_m_3 = ((found_3 != 0) ? m_max_3 - m_min_3 : 0); + int row_off_3 = q0_3 * 2 + seg_3 * 128 >> 7 << 7; + int rank_2 = cta_rank & 1; + int h_2 = blockIdx.y; + int kv_row0 = k0_3 + n_blk_3 * 128; + int n_groups_3 = group >> 1; + int t_3 = 0; + unsigned int _phase_q_empty = 1; + unsigned int _phase_do_empty = 1; + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(98304)) : "memory"); + tma_3d_gmem2smem_cta2(sk_addr + 0, (&K), 0, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 0, (&V), 0, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 8192, (&K), 64, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 8192, (&V), 64, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 16384, (&K), 128, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 16384, (&V), 128, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 24576, (&K), 192, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 24576, (&V), 192, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 0, (&K), rank_2 * 128 + 0, h_2, kv_row0 + 0, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 8192, (&K), rank_2 * 128 + 0, h_2, kv_row0 + 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 16384, (&K), rank_2 * 128 + 64, h_2, kv_row0 + 0, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 24576, (&K), rank_2 * 128 + 64, h_2, kv_row0 + 64, ((kv_full_addr) & 0xFEFFFFFF)); + #pragma unroll 1 + for (int g_3 = 0; g_3 < n_groups_3; g_3++) { + int head0 = h_2 * group + g_3 * 2; + int srow = (h_2 * n_groups_3 + g_3) * rows_pad + row_off_3; + #pragma unroll 1 + for (int mi_3 = 0; mi_3 < num_m_3; mi_3++) { + int tok = q0_3 + (m_min_3 + mi_3) * 64; + int stat_row = srow + (m_min_3 + mi_3) * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sq_addr + 0, (&Q), 0, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 8192, (&Q), 64, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 16384, (&Q), 128, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 24576, (&Q), 192, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + _phase_q_empty ^= 1; + int _parity_lse_empty = (int)((unsigned int)t_3 % 2) ^ 1; + mbarrier_wait(lse_empty_addr, _parity_lse_empty); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive_expect_tx(lse_full_addr, 512); + cp_async_bulk_gmem2smem(slse_addr, reinterpret_cast(reinterpret_cast(lse_log2) + ((unsigned long long)stat_row * 4)), 512, lse_full_addr); + mbarrier_wait(do_empty_addr, _phase_do_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((do_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sdo_addr + 0, (&dO), 0, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 8192, (&dO), 64, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 16384, (&dO), 128, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 24576, (&dO), 192, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + _phase_do_empty ^= 1; + int _parity_dpsum_empty = (int)((unsigned int)t_3 % 2) ^ 1; + mbarrier_wait(dpsum_empty_addr, _parity_dpsum_empty); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive_expect_tx(dpsum_full_addr, 512); + cp_async_bulk_gmem2smem(sdps_addr, reinterpret_cast(reinterpret_cast(dpsum) + ((unsigned long long)stat_row * 4)), 512, dpsum_full_addr); + mbarrier_wait(q_empty_addr, _phase_q_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sqt_addr + 0, (&Q), rank_2 * 128 + 0, head0, tok + 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 8192, (&Q), rank_2 * 128 + 0, head0, tok + 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 16384, (&Q), rank_2 * 128 + 64, head0, tok + 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 24576, (&Q), rank_2 * 128 + 64, head0, tok + 32, ((q_full_addr) & 0xFEFFFFFF)); + _phase_q_empty ^= 1; + mbarrier_wait(do_empty_addr, _phase_do_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((do_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sdot_addr + 0, (&dO), rank_2 * 128 + 0, head0, tok + 0, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 8192, (&dO), rank_2 * 128 + 0, head0, tok + 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 16384, (&dO), rank_2 * 128 + 64, head0, tok + 0, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 24576, (&dO), rank_2 * 128 + 64, head0, tok + 32, ((do_full_addr) & 0xFEFFFFFF)); + _phase_do_empty ^= 1; + t_3 += 1; + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty); + mbarrier_wait(do_empty_addr, _phase_do_empty); + } + } + } + // ---- Role: relay ---- + if (warp == 14) { + { // relay_main + int seg_4 = blockIdx.z; + int n_blk_4 = blockIdx.x >> 1; + int q0_4 = cu_seqlens_q[seg_4]; + int q_len_4 = cu_seqlens_q[seg_4 + 1] - q0_4; + int k0_4 = cu_seqlens_k[seg_4]; + int k_len_4 = cu_seqlens_k[seg_4 + 1] - k0_4; + int found_4 = ((n_blk_4 * 128 < k_len_4) ? 1 : 0); + int off_4 = k_len_4 - q_len_4; + int first_key_4 = n_blk_4 * 128 - off_4; + int m_min_4 = ((first_key_4 > 0) ? first_key_4 : 0) >> 6; + int m_max_4 = q_len_4 + 63 >> 6; + int num_m_4 = ((found_4 != 0) ? m_max_4 - m_min_4 : 0); + int row_off_4 = q0_4 * 2 + seg_4 * 128 >> 7 << 7; + int rank_3 = cta_rank & 1; + int n_groups_4 = group >> 1; + int t_4 = 0; + #pragma unroll 1 + for (int g_4 = 0; g_4 < n_groups_4; g_4++) { + #pragma unroll 1 + for (int mi_4 = 0; mi_4 < num_m_4; mi_4++) { + int _parity_xchg_full = (int)((unsigned int)t_4 % 2); + mbarrier_wait_cluster_hint(xchg_full_addr, _parity_xchg_full, 10000000); + if (elect_sync()) { + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((xchg_done_addr) & 0xFEFFFFFF) : "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(ds_empty_addr), "r"(1 - rank_3) : "memory"); + } + t_4 += 1; + } + } + } + } + + // Cleanup + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + if (warp == 12) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + } +} + +} // extern "C" diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu new file mode 100644 index 0000000..263c17e --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu @@ -0,0 +1,349 @@ +/* + * 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 Cake: the tvm-ffi host binding of kernel_cake_gqa_train_37f4cd45025705f3cfc8, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" +#include + +#include + +extern "C" __global__ void kernel_cake_gqa_train_37f4cd45025705f3cfc8(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap dO, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, float* __restrict__ lse_log2, float* __restrict__ dpsum, float* __restrict__ dq_accum, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, float* __restrict__ dk_out, float* __restrict__ dv_out, int num_segments, int h_kv, int group, int rows_pad, float softmax_scale, float softmax_scale_log2, int* __restrict__ row_map); + +namespace cake_binding_kernel_cake_gqa_train_37f4cd45025705f3cfc8 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gqa_train_37f4cd45025705f3cfc8"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CMod(int64_t a, int64_t b) { + if (b == 0) NoInt64(a, "%", b); + return b == -1 ? 0 : a % b; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_Q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'Q' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "Q"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'Q' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "Q"), Extent(d2, "Q"), Extent(d3, "Q")}; + cuuint64_t global_strides[2] = {Stride(s2, "Q"), Stride(s3, "Q")}; + cuuint32_t box_dim[3] = {64u, 2u, 32u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'Q' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'Q' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'Q' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of Q"); + return map; +} + +CUtensorMap EncodeTma_dO(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'dO' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "dO"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'dO' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "dO"), Extent(d2, "dO"), Extent(d3, "dO")}; + cuuint64_t global_strides[2] = {Stride(s2, "dO"), Stride(s3, "dO")}; + cuuint32_t box_dim[3] = {64u, 2u, 32u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'dO' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'dO' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'dO' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of dO"); + return map; +} + +CUtensorMap EncodeTma_K(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'K' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "K"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'K' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "K"), Extent(d2, "K"), Extent(d3, "K")}; + cuuint64_t global_strides[2] = {Stride(s2, "K"), Stride(s3, "K")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'K' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'K' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'K' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of K"); + return map; +} + +CUtensorMap EncodeTma_V(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'V' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "V"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'V' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "V"), Extent(d2, "V"), Extent(d3, "V")}; + cuuint64_t global_strides[2] = {Stride(s2, "V"), Stride(s3, "V")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'V' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'V' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'V' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of V"); + return map; +} + +void Run(TensorView a_Q, TensorView a_dO, TensorView a_K, TensorView a_V, TensorView a_lse_log2, TensorView a_dpsum, TensorView a_dq_accum, TensorView a_cu_seqlens_q, TensorView a_cu_seqlens_k, TensorView a_dk_out, TensorView a_dv_out, int64_t a_num_segments, int64_t a_h_kv, int64_t a_group, int64_t a_rows_pad, double a_softmax_scale, double a_softmax_scale_log2, Optional a_row_map, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_Q.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_Q, "Q"); + check_dtype(a_Q, DLDataType{kDLBfloat, 16, 1}, "Q"); + check_cuda_tensor(a_dO, "dO"); + check_dtype(a_dO, DLDataType{kDLBfloat, 16, 1}, "dO"); + check_cuda_tensor(a_K, "K"); + check_dtype(a_K, DLDataType{kDLBfloat, 16, 1}, "K"); + check_cuda_tensor(a_V, "V"); + check_dtype(a_V, DLDataType{kDLBfloat, 16, 1}, "V"); + check_cuda_tensor(a_lse_log2, "lse_log2"); + check_dtype(a_lse_log2, DLDataType{kDLFloat, 32, 1}, "lse_log2"); + check_cuda_tensor(a_dpsum, "dpsum"); + check_dtype(a_dpsum, DLDataType{kDLFloat, 32, 1}, "dpsum"); + check_cuda_tensor(a_dq_accum, "dq_accum"); + check_dtype(a_dq_accum, DLDataType{kDLFloat, 32, 1}, "dq_accum"); + check_cuda_tensor(a_cu_seqlens_q, "cu_seqlens_q"); + check_dtype(a_cu_seqlens_q, DLDataType{kDLInt, 32, 1}, "cu_seqlens_q"); + check_cuda_tensor(a_cu_seqlens_k, "cu_seqlens_k"); + check_dtype(a_cu_seqlens_k, DLDataType{kDLInt, 32, 1}, "cu_seqlens_k"); + check_cuda_tensor(a_dk_out, "dk_out"); + check_dtype(a_dk_out, DLDataType{kDLFloat, 32, 1}, "dk_out"); + check_cuda_tensor(a_dv_out, "dv_out"); + check_dtype(a_dv_out, DLDataType{kDLFloat, 32, 1}, "dv_out"); + TVM_FFI_CHECK(a_num_segments >= -2147483648LL && a_num_segments <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_segments' = " << a_num_segments + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_h_kv >= -2147483648LL && a_h_kv <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'h_kv' = " << a_h_kv + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_group >= -2147483648LL && a_group <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'group' = " << a_group + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_pad >= -2147483648LL && a_rows_pad <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_pad' = " << a_rows_pad + << " is outside the i32 range [-2147483648, 2147483647]"; + if (a_row_map.has_value()) check_cuda_tensor(a_row_map.value(), "row_map"); + if (a_row_map.has_value()) check_dtype(a_row_map.value(), DLDataType{kDLInt, 32, 1}, "row_map"); + check_same_device(a_dO, a_Q, "dO", "Q"); + check_same_device(a_K, a_Q, "K", "Q"); + check_same_device(a_V, a_Q, "V", "Q"); + check_contiguous(a_lse_log2, "lse_log2"); + check_same_device(a_lse_log2, a_Q, "lse_log2", "Q"); + check_contiguous(a_dpsum, "dpsum"); + check_same_device(a_dpsum, a_Q, "dpsum", "Q"); + check_contiguous(a_dq_accum, "dq_accum"); + check_same_device(a_dq_accum, a_Q, "dq_accum", "Q"); + check_contiguous(a_cu_seqlens_q, "cu_seqlens_q"); + check_same_device(a_cu_seqlens_q, a_Q, "cu_seqlens_q", "Q"); + check_contiguous(a_cu_seqlens_k, "cu_seqlens_k"); + check_same_device(a_cu_seqlens_k, a_Q, "cu_seqlens_k", "Q"); + check_contiguous(a_dk_out, "dk_out"); + check_same_device(a_dk_out, a_Q, "dk_out", "Q"); + check_contiguous(a_dv_out, "dv_out"); + check_same_device(a_dv_out, a_Q, "dv_out", "Q"); + if (a_row_map.has_value()) check_contiguous(a_row_map.value(), "row_map"); + if (a_row_map.has_value()) check_same_device(a_row_map.value(), a_Q, "row_map", "Q"); + TVM_FFI_CHECK((a_group >= 2LL), ValueError) + << KERNEL_NAME << ": requires (group >= 2)"; + TVM_FFI_CHECK((a_group <= 64LL), ValueError) + << KERNEL_NAME << ": requires (group <= 64)"; + TVM_FFI_CHECK((CMod(a_group, 2LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(group, 2) == 0)"; + TVM_FFI_CHECK((CMod(a_rows_pad, 128LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(rows_pad, 128) == 0)"; + TVM_FFI_CHECK((a_num_segments >= 1LL), ValueError) + << KERNEL_NAME << ": requires (num_segments >= 1)"; + TVM_FFI_CHECK((a_h_kv >= 1LL), ValueError) + << KERNEL_NAME << ": requires (h_kv >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") is not a multiple of the cluster shape (2, 1, 1)"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_gqa_train_37f4cd45025705f3cfc8), cudaFuncAttributeMaxDynamicSharedMemorySize, 231424), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_Q = EncodeTma_Q(a_Q); + CUtensorMap v_dO = EncodeTma_dO(a_dO); + CUtensorMap v_K = EncodeTma_K(a_K); + CUtensorMap v_V = EncodeTma_V(a_V); + void* v_lse_log2 = DataPtr(a_lse_log2); + void* v_dpsum = DataPtr(a_dpsum); + void* v_dq_accum = DataPtr(a_dq_accum); + void* v_cu_seqlens_q = DataPtr(a_cu_seqlens_q); + void* v_cu_seqlens_k = DataPtr(a_cu_seqlens_k); + void* v_dk_out = DataPtr(a_dk_out); + void* v_dv_out = DataPtr(a_dv_out); + int32_t v_num_segments = static_cast(a_num_segments); + int32_t v_h_kv = static_cast(a_h_kv); + int32_t v_group = static_cast(a_group); + int32_t v_rows_pad = static_cast(a_rows_pad); + float v_softmax_scale = static_cast(a_softmax_scale); + float v_softmax_scale_log2 = static_cast(a_softmax_scale_log2); + void* v_row_map = a_row_map.has_value() ? DataPtr(a_row_map.value()) : nullptr; + void* kernel_args[] = {&v_Q, &v_dO, &v_K, &v_V, &v_lse_log2, &v_dpsum, &v_dq_accum, &v_cu_seqlens_q, &v_cu_seqlens_k, &v_dk_out, &v_dv_out, &v_num_segments, &v_h_kv, &v_group, &v_rows_pad, &v_softmax_scale, &v_softmax_scale_log2, &v_row_map}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(512, 1, 1); + config.dynamicSmemBytes = 231424; + config.stream = stream; + cudaLaunchAttribute attrs[2] = {}; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = 2; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + config.attrs = attrs; + config.numAttrs = 2; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gqa_train_37f4cd45025705f3cfc8), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gqa_train_37f4cd45025705f3cfc8 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gqa_train_37f4cd45025705f3cfc8::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu new file mode 100644 index 0000000..83f14ab --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu @@ -0,0 +1,1604 @@ +/* + * 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 TMEM_NCOLS 512 +#define TMEM_DV_OFFSET 0 +#define TMEM_DK_OFFSET 128 +#define TMEM_S_OFFSET 256 +#define TMEM_DQ_OFFSET 384 +#define TMEM_DP_OFFSET 320 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_DO_PIPE_STAGES 1 +#define SMEM_SQ_OFF 1024 +#define SMEM_SQ_STAGE_BYTES 32768 +#define SMEM_SQ_STRIDE 32768 +#define SMEM_SQT_OFF 1024 +#define SMEM_SQT_STAGE_BYTES 32768 +#define SMEM_SQT_STRIDE 32768 +#define SMEM_SDO_OFF 33792 +#define SMEM_SDO_STAGE_BYTES 32768 +#define SMEM_SDO_STRIDE 32768 +#define SMEM_SDOT_OFF 33792 +#define SMEM_SDOT_STAGE_BYTES 32768 +#define SMEM_SDOT_STRIDE 32768 +#define SMEM_SK_OFF 66560 +#define SMEM_SK_STAGE_BYTES 32768 +#define SMEM_SK_STRIDE 32768 +#define SMEM_SV_OFF 99328 +#define SMEM_SV_STAGE_BYTES 32768 +#define SMEM_SV_STRIDE 32768 +#define SMEM_SKT_OFF 132096 +#define SMEM_SKT_STAGE_BYTES 32768 +#define SMEM_SKT_STRIDE 32768 +#define SMEM_SP_OFF 164864 +#define SMEM_SP_STAGE_BYTES 16384 +#define SMEM_SP_STRIDE 16384 +#define SMEM_SDST_OFF 181248 +#define SMEM_SDST_STAGE_BYTES 16384 +#define SMEM_SDST_STRIDE 16384 +#define SMEM_SDS_OFF 197632 +#define SMEM_SDS_STAGE_BYTES 16384 +#define SMEM_SDS_STRIDE 16384 +#define SMEM_SDQ_OFF 214016 +#define SMEM_SDQ_STAGE_BYTES 8192 +#define SMEM_SDQ_STRIDE 8192 +#define SMEM_SLSE_OFF 230400 +#define SMEM_SLSE_STAGE_BYTES 512 +#define SMEM_SLSE_STRIDE 512 +#define SMEM_SDPS_OFF 230912 +#define SMEM_SDPS_STAGE_BYTES 512 +#define SMEM_SDPS_STRIDE 512 +#define SMEM_TOTAL 231424 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, LAUNCH_MIN_BLOCKS) __cluster_dims__(2,1,1) void +kernel_cake_gqa_train_37f4cd45025705f3cfc8(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap dO, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, float* __restrict__ lse_log2, float* __restrict__ dpsum, float* __restrict__ dq_accum, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, float* __restrict__ dk_out, float* __restrict__ dv_out, int num_segments, int h_kv, int group, int rows_pad, float softmax_scale, float softmax_scale_log2, int* __restrict__ row_map) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define kv_full_addr (mbar_base + 0) + #define q_full_addr (mbar_base + 8) + #define q_empty_addr (mbar_base + 16) + #define do_full_addr (mbar_base + 24) + #define do_empty_addr (mbar_base + 32) + #define lse_full_addr (mbar_base + 40) + #define lse_empty_addr (mbar_base + 48) + #define dpsum_full_addr (mbar_base + 56) + #define dpsum_empty_addr (mbar_base + 64) + #define s_full_addr (mbar_base + 72) + #define dp_full_addr (mbar_base + 80) + #define ds_full_addr (mbar_base + 88) + #define ds_empty_addr (mbar_base + 96) + #define xchg_full_addr (mbar_base + 104) + #define xchg_done_addr (mbar_base + 112) + #define dq_full_addr (mbar_base + 120) + #define dq_empty_addr (mbar_base + 128) + #define dqlo_empty_addr (mbar_base + 136) + #define dkv_full_addr (mbar_base + 144) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* sq = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sq_addr = smem + 1024; + __nv_bfloat16* sqt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sqt_addr = smem + 1024; + __nv_bfloat16* sdo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int sdo_addr = smem + 33792; + __nv_bfloat16* sdot = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int sdot_addr = smem + 33792; + __nv_bfloat16* sk = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int sk_addr = smem + 66560; + __nv_bfloat16* sv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int sv_addr = smem + 99328; + __nv_bfloat16* skt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); + const int skt_addr = smem + 132096; + __nv_bfloat16* sp = reinterpret_cast<__nv_bfloat16*>(smem_raw + 164864); + const int sp_addr = smem + 164864; + __nv_bfloat16* sdst = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); + const int sdst_addr = smem + 181248; + __nv_bfloat16* sds = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int sds_addr = smem + 197632; + float* sdq = reinterpret_cast(smem_raw + 214016); + const int sdq_addr = smem + 214016; + float* slse = reinterpret_cast(smem_raw + 230400); + const int slse_addr = smem + 230400; + float* sdps = reinterpret_cast(smem_raw + 230912); + const int sdps_addr = smem + 230912; + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&Q))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&dO))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&K))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&V))) : "memory"); } + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 19 barriers) + // Mbarriers at smem_raw[0..152) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // kv_full: 1 barriers, init_count=2 + mbarrier_init(smem + 0, 2); + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=2 + mbarrier_init(smem + 8, 2); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // --- pipeline 'do_pipe' --- + // do_full: 1 barriers, init_count=2 + mbarrier_init(smem + 24, 2); + // do_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + // lse_full: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // lse_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 48, 256); + // dpsum_full: 1 barriers, init_count=1 + mbarrier_init(smem + 56, 1); + // dpsum_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 64, 256); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // dp_full: 1 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + // ds_full: 1 barriers, init_count=512 + mbarrier_init(smem + 88, 512); + // ds_empty: 1 barriers, init_count=2 + mbarrier_init(smem + 96, 2); + // xchg_full: 1 barriers, init_count=1 + mbarrier_init(smem + 104, 1); + // xchg_done: 1 barriers, init_count=2 + mbarrier_init(smem + 112, 2); + // dq_full: 1 barriers, init_count=1 + mbarrier_init(smem + 120, 1); + // dq_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 128, 256); + // dqlo_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 136, 256); + // dkv_full: 1 barriers, init_count=1 + mbarrier_init(smem + 144, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 152); + if (warp == 12) { + int _tmem_hold = smem + 152; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_dv = taddr; + const int tmem_dk = taddr + 128; + const int tmem_s = taddr + 256; + const int tmem_dq = taddr + 384; + const int tmem_dp = taddr + 320; + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); + } + + // ---- Role: reduce ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // reduce_main + int seg = blockIdx.z; + int n_blk = blockIdx.x >> 1; + int q0 = cu_seqlens_q[seg]; + int q_len = cu_seqlens_q[seg + 1] - q0; + int k0 = cu_seqlens_k[seg]; + int k_len = cu_seqlens_k[seg + 1] - k0; + int found = ((n_blk * 128 < k_len) ? 1 : 0); + int off = k_len - q_len; + int first_key = n_blk * 128 - off; + int m_min = ((first_key > 0) ? first_key : 0) >> 6; + int m_max = q_len + 63 >> 6; + int num_m = ((found != 0) ? m_max - m_min : 0); + int row_off = q0 * 2 + seg * 128 >> 7 << 7; + int rank = cta_rank & 1; + int h = blockIdx.y; + int n_groups = group >> 1; + const int quarter = warp % 4; + const int tl = quarter * 32 + lane; + const unsigned int lane_base = (unsigned int)(quarter * 32 << 16); + int t = 0; + #pragma unroll 1 + for (int g = 0; g < n_groups; g++) { + long long head_rows = (long long)((h * n_groups + g) * rows_pad + row_off); + #pragma unroll 1 + for (int mi = 0; mi < num_m; mi++) { + long long tile = (head_rows + (long long)((m_min + mi) * 128)) * 256 + (long long)(rank * 8 * 2048); + int _parity_dq_full = (int)((unsigned int)t % 2); + mbarrier_wait(dq_full_addr, _parity_dq_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + float dq_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_0[0]), "=f"(dq_0[1]), "=f"(dq_0[2]), "=f"(dq_0[3]), "=f"(dq_0[4]), "=f"(dq_0[5]), "=f"(dq_0[6]), "=f"(dq_0[7]), "=f"(dq_0[8]), "=f"(dq_0[9]), "=f"(dq_0[10]), "=f"(dq_0[11]), "=f"(dq_0[12]), "=f"(dq_0[13]), "=f"(dq_0[14]), "=f"(dq_0[15]), "=f"(dq_0[16]), "=f"(dq_0[17]), "=f"(dq_0[18]), "=f"(dq_0[19]), "=f"(dq_0[20]), "=f"(dq_0[21]), "=f"(dq_0[22]), "=f"(dq_0[23]), "=f"(dq_0[24]), "=f"(dq_0[25]), "=f"(dq_0[26]), "=f"(dq_0[27]), "=f"(dq_0[28]), "=f"(dq_0[29]), "=f"(dq_0[30]), "=f"(dq_0[31]) + : "r"(taddr + 448 + lane_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_0[32]), "=f"(dq_0[33]), "=f"(dq_0[34]), "=f"(dq_0[35]), "=f"(dq_0[36]), "=f"(dq_0[37]), "=f"(dq_0[38]), "=f"(dq_0[39]), "=f"(dq_0[40]), "=f"(dq_0[41]), "=f"(dq_0[42]), "=f"(dq_0[43]), "=f"(dq_0[44]), "=f"(dq_0[45]), "=f"(dq_0[46]), "=f"(dq_0[47]), "=f"(dq_0[48]), "=f"(dq_0[49]), "=f"(dq_0[50]), "=f"(dq_0[51]), "=f"(dq_0[52]), "=f"(dq_0[53]), "=f"(dq_0[54]), "=f"(dq_0[55]), "=f"(dq_0[56]), "=f"(dq_0[57]), "=f"(dq_0[58]), "=f"(dq_0[59]), "=f"(dq_0[60]), "=f"(dq_0[61]), "=f"(dq_0[62]), "=f"(dq_0[63]) + : "r"(taddr + 480 + lane_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((dq_empty_addr) & 0xFEFFFFFF) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_0[0]), "f"(dq_0[1]), "f"(dq_0[2]), "f"(dq_0[3]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_0[4]), "f"(dq_0[5]), "f"(dq_0[6]), "f"(dq_0[7]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_0[8]), "f"(dq_0[9]), "f"(dq_0[10]), "f"(dq_0[11]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_0[12]), "f"(dq_0[13]), "f"(dq_0[14]), "f"(dq_0[15]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_0 = reinterpret_cast(dq_accum + (tile + 8192)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_0), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_0[16]), "f"(dq_0[17]), "f"(dq_0[18]), "f"(dq_0[19]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_0[20]), "f"(dq_0[21]), "f"(dq_0[22]), "f"(dq_0[23]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_0[24]), "f"(dq_0[25]), "f"(dq_0[26]), "f"(dq_0[27]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_0[28]), "f"(dq_0[29]), "f"(dq_0[30]), "f"(dq_0[31]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_1 = reinterpret_cast(dq_accum + (tile + 10240)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_1), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_0[32]), "f"(dq_0[33]), "f"(dq_0[34]), "f"(dq_0[35]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_0[36]), "f"(dq_0[37]), "f"(dq_0[38]), "f"(dq_0[39]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_0[40]), "f"(dq_0[41]), "f"(dq_0[42]), "f"(dq_0[43]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_0[44]), "f"(dq_0[45]), "f"(dq_0[46]), "f"(dq_0[47]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_2 = reinterpret_cast(dq_accum + (tile + 12288)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_2), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_0[48]), "f"(dq_0[49]), "f"(dq_0[50]), "f"(dq_0[51]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_0[52]), "f"(dq_0[53]), "f"(dq_0[54]), "f"(dq_0[55]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_0[56]), "f"(dq_0[57]), "f"(dq_0[58]), "f"(dq_0[59]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_0[60]), "f"(dq_0[61]), "f"(dq_0[62]), "f"(dq_0[63]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_3 = reinterpret_cast(dq_accum + (tile + 14336)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_3), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float dq_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_1[0]), "=f"(dq_1[1]), "=f"(dq_1[2]), "=f"(dq_1[3]), "=f"(dq_1[4]), "=f"(dq_1[5]), "=f"(dq_1[6]), "=f"(dq_1[7]), "=f"(dq_1[8]), "=f"(dq_1[9]), "=f"(dq_1[10]), "=f"(dq_1[11]), "=f"(dq_1[12]), "=f"(dq_1[13]), "=f"(dq_1[14]), "=f"(dq_1[15]), "=f"(dq_1[16]), "=f"(dq_1[17]), "=f"(dq_1[18]), "=f"(dq_1[19]), "=f"(dq_1[20]), "=f"(dq_1[21]), "=f"(dq_1[22]), "=f"(dq_1[23]), "=f"(dq_1[24]), "=f"(dq_1[25]), "=f"(dq_1[26]), "=f"(dq_1[27]), "=f"(dq_1[28]), "=f"(dq_1[29]), "=f"(dq_1[30]), "=f"(dq_1[31]) + : "r"(taddr + 384 + lane_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_1[32]), "=f"(dq_1[33]), "=f"(dq_1[34]), "=f"(dq_1[35]), "=f"(dq_1[36]), "=f"(dq_1[37]), "=f"(dq_1[38]), "=f"(dq_1[39]), "=f"(dq_1[40]), "=f"(dq_1[41]), "=f"(dq_1[42]), "=f"(dq_1[43]), "=f"(dq_1[44]), "=f"(dq_1[45]), "=f"(dq_1[46]), "=f"(dq_1[47]), "=f"(dq_1[48]), "=f"(dq_1[49]), "=f"(dq_1[50]), "=f"(dq_1[51]), "=f"(dq_1[52]), "=f"(dq_1[53]), "=f"(dq_1[54]), "=f"(dq_1[55]), "=f"(dq_1[56]), "=f"(dq_1[57]), "=f"(dq_1[58]), "=f"(dq_1[59]), "=f"(dq_1[60]), "=f"(dq_1[61]), "=f"(dq_1[62]), "=f"(dq_1[63]) + : "r"(taddr + 416 + lane_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((dqlo_empty_addr) & 0xFEFFFFFF) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_1[0]), "f"(dq_1[1]), "f"(dq_1[2]), "f"(dq_1[3]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_1[4]), "f"(dq_1[5]), "f"(dq_1[6]), "f"(dq_1[7]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_1[8]), "f"(dq_1[9]), "f"(dq_1[10]), "f"(dq_1[11]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_1[12]), "f"(dq_1[13]), "f"(dq_1[14]), "f"(dq_1[15]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_4 = reinterpret_cast(dq_accum + (tile + 0)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_4), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_1[16]), "f"(dq_1[17]), "f"(dq_1[18]), "f"(dq_1[19]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_1[20]), "f"(dq_1[21]), "f"(dq_1[22]), "f"(dq_1[23]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_1[24]), "f"(dq_1[25]), "f"(dq_1[26]), "f"(dq_1[27]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_1[28]), "f"(dq_1[29]), "f"(dq_1[30]), "f"(dq_1[31]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_5 = reinterpret_cast(dq_accum + (tile + 2048)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_5), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_1[32]), "f"(dq_1[33]), "f"(dq_1[34]), "f"(dq_1[35]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_1[36]), "f"(dq_1[37]), "f"(dq_1[38]), "f"(dq_1[39]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_1[40]), "f"(dq_1[41]), "f"(dq_1[42]), "f"(dq_1[43]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_1[44]), "f"(dq_1[45]), "f"(dq_1[46]), "f"(dq_1[47]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_6 = reinterpret_cast(dq_accum + (tile + 4096)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_6), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_1[48]), "f"(dq_1[49]), "f"(dq_1[50]), "f"(dq_1[51]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_1[52]), "f"(dq_1[53]), "f"(dq_1[54]), "f"(dq_1[55]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_1[56]), "f"(dq_1[57]), "f"(dq_1[58]), "f"(dq_1[59]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_1[60]), "f"(dq_1[61]), "f"(dq_1[62]), "f"(dq_1[63]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_7 = reinterpret_cast(dq_accum + (tile + 6144)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_7), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + t += 1; + } + } + if (warp == 0) { + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + } + } + } + } + // ---- Role: compute ---- + if (warp >= 4 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + { // compute_main + int seg_1 = blockIdx.z; + int n_blk_1 = blockIdx.x >> 1; + int q0_1 = cu_seqlens_q[seg_1]; + int q_len_1 = cu_seqlens_q[seg_1 + 1] - q0_1; + int k0_1 = cu_seqlens_k[seg_1]; + int k_len_1 = cu_seqlens_k[seg_1 + 1] - k0_1; + int found_1 = ((n_blk_1 * 128 < k_len_1) ? 1 : 0); + int off_1 = k_len_1 - q_len_1; + int first_key_1 = n_blk_1 * 128 - off_1; + int m_min_1 = ((first_key_1 > 0) ? first_key_1 : 0) >> 6; + int m_max_1 = q_len_1 + 63 >> 6; + int num_m_1 = ((found_1 != 0) ? m_max_1 - m_min_1 : 0); + int row_off_1 = q0_1 * 2 + seg_1 * 128 >> 7 << 7; + int rank_1 = cta_rank & 1; + int h_1 = blockIdx.y; + int n_groups_1 = group >> 1; + const int quarter_1 = warp % 4; + const int wg = warp - 4 >> 2; + const int row = (quarter_1 & 1) * 32 + lane; + const int half = quarter_1 >> 1; + const int col0 = half * 64 + wg * 32; + const unsigned int lane_base_1 = (unsigned int)(quarter_1 * 32 << 16); + int key = n_blk_1 * 128 + rank_1 * 64 + row; + int prow = half * 64 + row; + int t_1 = 0; + #pragma unroll 1 + for (int g_1 = 0; g_1 < n_groups_1; g_1++) { + #pragma unroll 1 + for (int mi_1 = 0; mi_1 < num_m_1; mi_1++) { + int m = m_min_1 + mi_1; + int _parity_lse_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(lse_full_addr, _parity_lse_full); + int _parity_s_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(s_full_addr, _parity_s_full); + int _parity_ds_empty = (int)((unsigned int)t_1 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); + float s[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(s[0]), "=f"(s[1]), "=f"(s[2]), "=f"(s[3]), "=f"(s[4]), "=f"(s[5]), "=f"(s[6]), "=f"(s[7]), "=f"(s[8]), "=f"(s[9]), "=f"(s[10]), "=f"(s[11]), "=f"(s[12]), "=f"(s[13]), "=f"(s[14]), "=f"(s[15]), "=f"(s[16]), "=f"(s[17]), "=f"(s[18]), "=f"(s[19]), "=f"(s[20]), "=f"(s[21]), "=f"(s[22]), "=f"(s[23]), "=f"(s[24]), "=f"(s[25]), "=f"(s[26]), "=f"(s[27]), "=f"(s[28]), "=f"(s[29]), "=f"(s[30]), "=f"(s[31]) + : "r"(taddr + (unsigned int)(256 + wg * 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int lo = (key - off_1 - m * 64) * 2 - col0; + int hi = (q_len_1 - m * 64) * 2 - col0; + int hi_k = ((key < k_len_1) ? hi : 0); + uint32_t _thresh_mask_0; + { + int _lim_1 = hi_k; + if (_lim_1 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_1)); + } + } + uint32_t _thresh_mask_2; + { + int _lim_3 = lo; + if (_lim_3 <= 0) { _thresh_mask_2 = 0u; } + else if (_lim_3 >= 32) { _thresh_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_2) : "r"(_lim_3)); + } + } + unsigned int keep = _thresh_mask_0 & ~_thresh_mask_2; + if (!(keep & (1u << 0))) s[0] = -CUDART_INF_F; + if (!(keep & (1u << 1))) s[1] = -CUDART_INF_F; + if (!(keep & (1u << 2))) s[2] = -CUDART_INF_F; + if (!(keep & (1u << 3))) s[3] = -CUDART_INF_F; + if (!(keep & (1u << 4))) s[4] = -CUDART_INF_F; + if (!(keep & (1u << 5))) s[5] = -CUDART_INF_F; + if (!(keep & (1u << 6))) s[6] = -CUDART_INF_F; + if (!(keep & (1u << 7))) s[7] = -CUDART_INF_F; + if (!(keep & (1u << 8))) s[8] = -CUDART_INF_F; + if (!(keep & (1u << 9))) s[9] = -CUDART_INF_F; + if (!(keep & (1u << 10))) s[10] = -CUDART_INF_F; + if (!(keep & (1u << 11))) s[11] = -CUDART_INF_F; + if (!(keep & (1u << 12))) s[12] = -CUDART_INF_F; + if (!(keep & (1u << 13))) s[13] = -CUDART_INF_F; + if (!(keep & (1u << 14))) s[14] = -CUDART_INF_F; + if (!(keep & (1u << 15))) s[15] = -CUDART_INF_F; + if (!(keep & (1u << 16))) s[16] = -CUDART_INF_F; + if (!(keep & (1u << 17))) s[17] = -CUDART_INF_F; + if (!(keep & (1u << 18))) s[18] = -CUDART_INF_F; + if (!(keep & (1u << 19))) s[19] = -CUDART_INF_F; + if (!(keep & (1u << 20))) s[20] = -CUDART_INF_F; + if (!(keep & (1u << 21))) s[21] = -CUDART_INF_F; + if (!(keep & (1u << 22))) s[22] = -CUDART_INF_F; + if (!(keep & (1u << 23))) s[23] = -CUDART_INF_F; + if (!(keep & (1u << 24))) s[24] = -CUDART_INF_F; + if (!(keep & (1u << 25))) s[25] = -CUDART_INF_F; + if (!(keep & (1u << 26))) s[26] = -CUDART_INF_F; + if (!(keep & (1u << 27))) s[27] = -CUDART_INF_F; + if (!(keep & (1u << 28))) s[28] = -CUDART_INF_F; + if (!(keep & (1u << 29))) s[29] = -CUDART_INF_F; + if (!(keep & (1u << 30))) s[30] = -CUDART_INF_F; + if (!(keep & (1u << 31))) s[31] = -CUDART_INF_F; + float lse[32]; + #pragma unroll + for (int j = 0; j < 32; j++) { + lse[j] = slse[col0 + j]; + } + float2 scale2 = make_float2(softmax_scale_log2, softmax_scale_log2); + #pragma unroll + for (int pair = 0; pair < 16; pair++) { + float2 x = fma_sub_f32x2_rn_ftz(make_float2(s[pair * 2], s[pair * 2 + 1]), scale2, make_float2(lse[pair * 2], lse[pair * 2 + 1])); + s[pair * 2] = x.x; + s[pair * 2 + 1] = x.y; + } + unsigned int p16[16]; + softmax_frag_exp2_cast(&s[0], p16, 0); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 0 ^ prow & 7) * 16))), "r"(p16[0]), "r"(p16[1]), "r"(p16[2]), "r"(p16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 1 ^ prow & 7) * 16))), "r"(p16[4]), "r"(p16[5]), "r"(p16[6]), "r"(p16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 2 ^ prow & 7) * 16))), "r"(p16[8]), "r"(p16[9]), "r"(p16[10]), "r"(p16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 3 ^ prow & 7) * 16))), "r"(p16[12]), "r"(p16[13]), "r"(p16[14]), "r"(p16[15]) : "memory"); + mbarrier_arrive(lse_empty_addr); + int _parity_dpsum_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(dpsum_full_addr, _parity_dpsum_full); + int _parity_dp_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(dp_full_addr, _parity_dp_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + float dp[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dp[0]), "=f"(dp[1]), "=f"(dp[2]), "=f"(dp[3]), "=f"(dp[4]), "=f"(dp[5]), "=f"(dp[6]), "=f"(dp[7]), "=f"(dp[8]), "=f"(dp[9]), "=f"(dp[10]), "=f"(dp[11]), "=f"(dp[12]), "=f"(dp[13]), "=f"(dp[14]), "=f"(dp[15]), "=f"(dp[16]), "=f"(dp[17]), "=f"(dp[18]), "=f"(dp[19]), "=f"(dp[20]), "=f"(dp[21]), "=f"(dp[22]), "=f"(dp[23]), "=f"(dp[24]), "=f"(dp[25]), "=f"(dp[26]), "=f"(dp[27]), "=f"(dp[28]), "=f"(dp[29]), "=f"(dp[30]), "=f"(dp[31]) + : "r"(taddr + (unsigned int)(320 + wg * 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float dps[32]; + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1++) { + dps[j_1] = sdps[col0 + j_1]; + } + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + sub_f32x2_inplace(&reinterpret_cast(dp)[_ls], reinterpret_cast(dps)[_ls]); + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(dp)[_ls], reinterpret_cast(s)[_ls]); + unsigned int ds16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dp[_lp*2 + 0], dp[_lp*2+1 + 0])); + ds16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 0 ^ prow & 7) * 16))), "r"(ds16[0]), "r"(ds16[1]), "r"(ds16[2]), "r"(ds16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 1 ^ prow & 7) * 16))), "r"(ds16[4]), "r"(ds16[5]), "r"(ds16[6]), "r"(ds16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 2 ^ prow & 7) * 16))), "r"(ds16[8]), "r"(ds16[9]), "r"(ds16[10]), "r"(ds16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 3 ^ prow & 7) * 16))), "r"(ds16[12]), "r"(ds16[13]), "r"(ds16[14]), "r"(ds16[15]) : "memory"); + if (half == rank_1) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 0 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[0]), "r"(ds16[1]), "r"(ds16[2]), "r"(ds16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 1 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[4]), "r"(ds16[5]), "r"(ds16[6]), "r"(ds16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 2 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[8]), "r"(ds16[9]), "r"(ds16[10]), "r"(ds16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 3 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[12]), "r"(ds16[13]), "r"(ds16[14]), "r"(ds16[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((ds_full_addr) & 0xFEFFFFFF) : "memory"); + mbarrier_arrive(dpsum_empty_addr); + asm volatile("barrier.sync 1, 256;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(xchg_full_addr), "r"(1 - rank_1), "r"((uint32_t)(8192)) : "memory"); + uint32_t _mbar_peer_4; + asm( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mbar_peer_4) : "r"(xchg_full_addr), "r"(1 - rank_1)); + uint32_t _cluster_mapa_5; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_cluster_mapa_5) : "r"(sds_addr + (unsigned int)(rank_1 * 8192)), "r"(1 - rank_1)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_cluster_mapa_5), "r"(sdst_addr + (unsigned int)((1 - rank_1) * 8192)), "r"((uint32_t)(8192)), "r"(_mbar_peer_4) + : "memory"); + } + } + t_1 += 1; + } + } + int _parity_ds_empty_1 = (int)((unsigned int)t_1 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty_1, 10000000); + mbarrier_wait(dkv_full_addr, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int hd0 = half * 128 + wg * 64; + int store_row = ((key < k_len_1) ? 1 : 0); + if (found_1 == 0) { + store_row = 0; + } + long long out_row = (long long)(k0_1 + key); + if (store_row != 0) { + out_row = (long long)row_map[k0_1 + key]; + } + float acc_0[64]; + if (num_m_1 > 0) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_0[0]), "=f"(acc_0[1]), "=f"(acc_0[2]), "=f"(acc_0[3]), "=f"(acc_0[4]), "=f"(acc_0[5]), "=f"(acc_0[6]), "=f"(acc_0[7]), "=f"(acc_0[8]), "=f"(acc_0[9]), "=f"(acc_0[10]), "=f"(acc_0[11]), "=f"(acc_0[12]), "=f"(acc_0[13]), "=f"(acc_0[14]), "=f"(acc_0[15]), "=f"(acc_0[16]), "=f"(acc_0[17]), "=f"(acc_0[18]), "=f"(acc_0[19]), "=f"(acc_0[20]), "=f"(acc_0[21]), "=f"(acc_0[22]), "=f"(acc_0[23]), "=f"(acc_0[24]), "=f"(acc_0[25]), "=f"(acc_0[26]), "=f"(acc_0[27]), "=f"(acc_0[28]), "=f"(acc_0[29]), "=f"(acc_0[30]), "=f"(acc_0[31]) + : "r"(taddr + (unsigned int)(0 + wg * 64) + lane_base_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_0[32]), "=f"(acc_0[33]), "=f"(acc_0[34]), "=f"(acc_0[35]), "=f"(acc_0[36]), "=f"(acc_0[37]), "=f"(acc_0[38]), "=f"(acc_0[39]), "=f"(acc_0[40]), "=f"(acc_0[41]), "=f"(acc_0[42]), "=f"(acc_0[43]), "=f"(acc_0[44]), "=f"(acc_0[45]), "=f"(acc_0[46]), "=f"(acc_0[47]), "=f"(acc_0[48]), "=f"(acc_0[49]), "=f"(acc_0[50]), "=f"(acc_0[51]), "=f"(acc_0[52]), "=f"(acc_0[53]), "=f"(acc_0[54]), "=f"(acc_0[55]), "=f"(acc_0[56]), "=f"(acc_0[57]), "=f"(acc_0[58]), "=f"(acc_0[59]), "=f"(acc_0[60]), "=f"(acc_0[61]), "=f"(acc_0[62]), "=f"(acc_0[63]) + : "r"(taddr + (unsigned int)(0 + wg * 64 + 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } else { + #pragma unroll + for (int i_0 = 0; i_0 < 64; i_0++) { + acc_0[i_0] = 0.0f; + } + } + long long base_0 = (out_row * (long long)h_kv + (long long)h_1) * 256 + (long long)hd0; + if (store_row != 0) { + if (num_m_1 > 0) { + #pragma unroll + for (int i_1 = 0; i_1 < 16; i_1++) { + asm volatile("red.global.add.v4.f32 [%0], {%1, %2, %3, %4};" :: "l"(reinterpret_cast(&(dv_out + (base_0 + (long long)(i_1 * 4)))[0])), "f"(acc_0[i_1 * 4]), "f"(acc_0[i_1 * 4 + 1]), "f"(acc_0[i_1 * 4 + 2]), "f"(acc_0[i_1 * 4 + 3]) : "memory"); + } + } + } + float acc_1[64]; + if (num_m_1 > 0) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_1[0]), "=f"(acc_1[1]), "=f"(acc_1[2]), "=f"(acc_1[3]), "=f"(acc_1[4]), "=f"(acc_1[5]), "=f"(acc_1[6]), "=f"(acc_1[7]), "=f"(acc_1[8]), "=f"(acc_1[9]), "=f"(acc_1[10]), "=f"(acc_1[11]), "=f"(acc_1[12]), "=f"(acc_1[13]), "=f"(acc_1[14]), "=f"(acc_1[15]), "=f"(acc_1[16]), "=f"(acc_1[17]), "=f"(acc_1[18]), "=f"(acc_1[19]), "=f"(acc_1[20]), "=f"(acc_1[21]), "=f"(acc_1[22]), "=f"(acc_1[23]), "=f"(acc_1[24]), "=f"(acc_1[25]), "=f"(acc_1[26]), "=f"(acc_1[27]), "=f"(acc_1[28]), "=f"(acc_1[29]), "=f"(acc_1[30]), "=f"(acc_1[31]) + : "r"(taddr + (unsigned int)(128 + wg * 64) + lane_base_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_1[32]), "=f"(acc_1[33]), "=f"(acc_1[34]), "=f"(acc_1[35]), "=f"(acc_1[36]), "=f"(acc_1[37]), "=f"(acc_1[38]), "=f"(acc_1[39]), "=f"(acc_1[40]), "=f"(acc_1[41]), "=f"(acc_1[42]), "=f"(acc_1[43]), "=f"(acc_1[44]), "=f"(acc_1[45]), "=f"(acc_1[46]), "=f"(acc_1[47]), "=f"(acc_1[48]), "=f"(acc_1[49]), "=f"(acc_1[50]), "=f"(acc_1[51]), "=f"(acc_1[52]), "=f"(acc_1[53]), "=f"(acc_1[54]), "=f"(acc_1[55]), "=f"(acc_1[56]), "=f"(acc_1[57]), "=f"(acc_1[58]), "=f"(acc_1[59]), "=f"(acc_1[60]), "=f"(acc_1[61]), "=f"(acc_1[62]), "=f"(acc_1[63]) + : "r"(taddr + (unsigned int)(128 + wg * 64 + 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } else { + #pragma unroll + for (int i_2 = 0; i_2 < 64; i_2++) { + acc_1[i_2] = 0.0f; + } + } + const float2 _scale2_6 = {softmax_scale, softmax_scale}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(acc_1)[_ls], _scale2_6); + long long base_1 = (out_row * (long long)h_kv + (long long)h_1) * 256 + (long long)hd0; + if (store_row != 0) { + if (num_m_1 > 0) { + #pragma unroll + for (int i_3 = 0; i_3 < 16; i_3++) { + asm volatile("red.global.add.v4.f32 [%0], {%1, %2, %3, %4};" :: "l"(reinterpret_cast(&(dk_out + (base_1 + (long long)(i_3 * 4)))[0])), "f"(acc_1[i_3 * 4]), "f"(acc_1[i_3 * 4 + 1]), "f"(acc_1[i_3 * 4 + 2]), "f"(acc_1[i_3 * 4 + 3]) : "memory"); + } + } + } + } + } + // ---- Role: mma ---- + if (warp == 12) { + { // mma_main + int seg_2 = blockIdx.z; + int n_blk_2 = blockIdx.x >> 1; + int q0_2 = cu_seqlens_q[seg_2]; + int q_len_2 = cu_seqlens_q[seg_2 + 1] - q0_2; + int k0_2 = cu_seqlens_k[seg_2]; + int k_len_2 = cu_seqlens_k[seg_2 + 1] - k0_2; + int found_2 = ((n_blk_2 * 128 < k_len_2) ? 1 : 0); + int off_2 = k_len_2 - q_len_2; + int first_key_2 = n_blk_2 * 128 - off_2; + int m_min_2 = ((first_key_2 > 0) ? first_key_2 : 0) >> 6; + int m_max_2 = q_len_2 + 63 >> 6; + int num_m_2 = ((found_2 != 0) ? m_max_2 - m_min_2 : 0); + int row_off_2 = q0_2 * 2 + seg_2 * 128 >> 7 << 7; + int n_groups_2 = group >> 1; + int t_2 = 0; + unsigned int _phase_q_full = 0; + unsigned int _phase_do_full = 0; + if ((cta_rank & 1) == 0) { + if (elect_sync()) { + mbarrier_wait_cluster_hint(kv_full_addr, 0, 10000000); + #pragma unroll 1 + for (int g_2 = 0; g_2 < n_groups_2; g_2++) { + #pragma unroll 1 + for (int mi_2 = 0; mi_2 < num_m_2; mi_2++) { + int first = ((t_2 == 0) ? 1 : 0); + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = ((sk_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_1 = ((sq_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_0 = ((sk_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_1 = ((sq_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_s), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(s_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + mbarrier_wait_cluster_hint(do_full_addr, _phase_do_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = ((sv_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_3 = ((sdo_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_2 = ((sv_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_3 = ((sdo_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_dp), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(dp_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(do_empty_addr, (uint16_t)(3)); + _phase_do_full ^= 1; + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + int _parity_ds_full = (int)((unsigned int)t_2 % 2); + mbarrier_wait_cluster_hint(ds_full_addr, _parity_ds_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_4 = ((sdst_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_5 = (((sqt_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_4 = ((sdst_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_5 = (((sqt_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_4), "r"(_smem_desc_lo_5), "r"(tmem_dk), "r"(((first) ? 0 : 1)) : "memory"); + tcgen05_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + mbarrier_wait_cluster_hint(do_full_addr, _phase_do_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_6 = ((sp_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_7 = (((sdot_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_6 = ((sp_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_7 = (((sdot_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_6), "r"(_smem_desc_lo_7), "r"(tmem_dv), "r"(((first) ? 0 : 1)) : "memory"); + tcgen05_commit_cg2_multicast(do_empty_addr, (uint16_t)(3)); + _phase_do_full ^= 1; + int _parity_xchg_done = (int)((unsigned int)t_2 % 2); + mbarrier_wait_cluster_hint(xchg_done_addr, _parity_xchg_done, 10000000); + int _parity_ds_empty_2 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty_2, 10000000); + int _parity_dq_empty = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dq_empty_addr, _parity_dq_empty, 10000000); + int _parity_dqlo_empty = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dqlo_empty_addr, _parity_dqlo_empty, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_8 = (((sds_addr) >> 4) & 0x7FFF) | 0x4000000; + int _smem_desc_lo_9 = (((skt_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_8 = (((sds_addr) >> 4) & 0x3FFF) | 0x4000000; + int _smem_desc_lo_9 = (((skt_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138511504;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_8), "r"(_smem_desc_lo_9), "r"(tmem_dq), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(dq_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(ds_empty_addr, (uint16_t)(3)); + t_2 += 1; + } + } + tcgen05_commit_cg2_multicast(dkv_full_addr, (uint16_t)(3)); + int _parity_dq_empty_1 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dq_empty_addr, _parity_dq_empty_1, 10000000); + int _parity_dqlo_empty_1 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dqlo_empty_addr, _parity_dqlo_empty_1, 10000000); + } + } + } + } + // ---- Role: load ---- + if (warp == 13) { + { // load_main + int seg_3 = blockIdx.z; + int n_blk_3 = blockIdx.x >> 1; + int q0_3 = cu_seqlens_q[seg_3]; + int q_len_3 = cu_seqlens_q[seg_3 + 1] - q0_3; + int k0_3 = cu_seqlens_k[seg_3]; + int k_len_3 = cu_seqlens_k[seg_3 + 1] - k0_3; + int found_3 = ((n_blk_3 * 128 < k_len_3) ? 1 : 0); + int off_3 = k_len_3 - q_len_3; + int first_key_3 = n_blk_3 * 128 - off_3; + int m_min_3 = ((first_key_3 > 0) ? first_key_3 : 0) >> 6; + int m_max_3 = q_len_3 + 63 >> 6; + int num_m_3 = ((found_3 != 0) ? m_max_3 - m_min_3 : 0); + int row_off_3 = q0_3 * 2 + seg_3 * 128 >> 7 << 7; + int rank_2 = cta_rank & 1; + int h_2 = blockIdx.y; + int kv_row0 = k0_3 + n_blk_3 * 128; + int n_groups_3 = group >> 1; + int t_3 = 0; + unsigned int _phase_q_empty = 1; + unsigned int _phase_do_empty = 1; + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(98304)) : "memory"); + tma_3d_gmem2smem_cta2(sk_addr + 0, (&K), 0, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 0, (&V), 0, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 8192, (&K), 64, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 8192, (&V), 64, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 16384, (&K), 128, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 16384, (&V), 128, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 24576, (&K), 192, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 24576, (&V), 192, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 0, (&K), rank_2 * 128 + 0, h_2, kv_row0 + 0, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 8192, (&K), rank_2 * 128 + 0, h_2, kv_row0 + 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 16384, (&K), rank_2 * 128 + 64, h_2, kv_row0 + 0, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 24576, (&K), rank_2 * 128 + 64, h_2, kv_row0 + 64, ((kv_full_addr) & 0xFEFFFFFF)); + #pragma unroll 1 + for (int g_3 = 0; g_3 < n_groups_3; g_3++) { + int head0 = h_2 * group + g_3 * 2; + int srow = (h_2 * n_groups_3 + g_3) * rows_pad + row_off_3; + #pragma unroll 1 + for (int mi_3 = 0; mi_3 < num_m_3; mi_3++) { + int tok = q0_3 + (m_min_3 + mi_3) * 64; + int stat_row = srow + (m_min_3 + mi_3) * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sq_addr + 0, (&Q), 0, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 8192, (&Q), 64, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 16384, (&Q), 128, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 24576, (&Q), 192, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + _phase_q_empty ^= 1; + int _parity_lse_empty = (int)((unsigned int)t_3 % 2) ^ 1; + mbarrier_wait(lse_empty_addr, _parity_lse_empty); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive_expect_tx(lse_full_addr, 512); + cp_async_bulk_gmem2smem(slse_addr, reinterpret_cast(reinterpret_cast(lse_log2) + ((unsigned long long)stat_row * 4)), 512, lse_full_addr); + mbarrier_wait(do_empty_addr, _phase_do_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((do_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sdo_addr + 0, (&dO), 0, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 8192, (&dO), 64, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 16384, (&dO), 128, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 24576, (&dO), 192, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + _phase_do_empty ^= 1; + int _parity_dpsum_empty = (int)((unsigned int)t_3 % 2) ^ 1; + mbarrier_wait(dpsum_empty_addr, _parity_dpsum_empty); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive_expect_tx(dpsum_full_addr, 512); + cp_async_bulk_gmem2smem(sdps_addr, reinterpret_cast(reinterpret_cast(dpsum) + ((unsigned long long)stat_row * 4)), 512, dpsum_full_addr); + mbarrier_wait(q_empty_addr, _phase_q_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sqt_addr + 0, (&Q), rank_2 * 128 + 0, head0, tok + 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 8192, (&Q), rank_2 * 128 + 0, head0, tok + 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 16384, (&Q), rank_2 * 128 + 64, head0, tok + 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 24576, (&Q), rank_2 * 128 + 64, head0, tok + 32, ((q_full_addr) & 0xFEFFFFFF)); + _phase_q_empty ^= 1; + mbarrier_wait(do_empty_addr, _phase_do_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((do_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sdot_addr + 0, (&dO), rank_2 * 128 + 0, head0, tok + 0, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 8192, (&dO), rank_2 * 128 + 0, head0, tok + 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 16384, (&dO), rank_2 * 128 + 64, head0, tok + 0, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 24576, (&dO), rank_2 * 128 + 64, head0, tok + 32, ((do_full_addr) & 0xFEFFFFFF)); + _phase_do_empty ^= 1; + t_3 += 1; + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty); + mbarrier_wait(do_empty_addr, _phase_do_empty); + } + } + } + // ---- Role: relay ---- + if (warp == 14) { + { // relay_main + int seg_4 = blockIdx.z; + int n_blk_4 = blockIdx.x >> 1; + int q0_4 = cu_seqlens_q[seg_4]; + int q_len_4 = cu_seqlens_q[seg_4 + 1] - q0_4; + int k0_4 = cu_seqlens_k[seg_4]; + int k_len_4 = cu_seqlens_k[seg_4 + 1] - k0_4; + int found_4 = ((n_blk_4 * 128 < k_len_4) ? 1 : 0); + int off_4 = k_len_4 - q_len_4; + int first_key_4 = n_blk_4 * 128 - off_4; + int m_min_4 = ((first_key_4 > 0) ? first_key_4 : 0) >> 6; + int m_max_4 = q_len_4 + 63 >> 6; + int num_m_4 = ((found_4 != 0) ? m_max_4 - m_min_4 : 0); + int row_off_4 = q0_4 * 2 + seg_4 * 128 >> 7 << 7; + int rank_3 = cta_rank & 1; + int n_groups_4 = group >> 1; + int t_4 = 0; + #pragma unroll 1 + for (int g_4 = 0; g_4 < n_groups_4; g_4++) { + #pragma unroll 1 + for (int mi_4 = 0; mi_4 < num_m_4; mi_4++) { + int _parity_xchg_full = (int)((unsigned int)t_4 % 2); + mbarrier_wait_cluster_hint(xchg_full_addr, _parity_xchg_full, 10000000); + if (elect_sync()) { + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((xchg_done_addr) & 0xFEFFFFFF) : "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(ds_empty_addr), "r"(1 - rank_3) : "memory"); + } + t_4 += 1; + } + } + } + } + + // Cleanup + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + if (warp == 12) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + } +} + +} // extern "C" diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu new file mode 100644 index 0000000..3dfd95b --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu @@ -0,0 +1,281 @@ +/* + * 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 Cake: the tvm-ffi host binding of kernel_cake_gqa_train_493c4a28dafa5fc6b2f1, for the device source that defines it. Do not edit. +#include +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_gqa_train_493c4a28dafa5fc6b2f1(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, unsigned int* __restrict__ tile_counter, int num_segments, int num_heads_q, int group, int total_q, float softmax_scale_log2); + +namespace cake_binding_kernel_cake_gqa_train_493c4a28dafa5fc6b2f1 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gqa_train_493c4a28dafa5fc6b2f1"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CDiv(int64_t a, int64_t b) { + if (b == 0 || (a == INT64_MIN && b == -1)) NoInt64(a, "/", b); + return a / b; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_Q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'Q' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "Q"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'Q' requires (d1 == 256)"; + cuuint64_t global_dim[4] = {Extent(64LL, "Q"), Extent(d3, "Q"), Extent(d2, "Q"), Extent(CDiv(d1, 64LL), "Q")}; + cuuint64_t global_strides[3] = {Stride(s3, "Q"), Stride(s2, "Q"), Stride(128LL, "Q")}; + cuuint32_t box_dim[4] = {64u, 128u, 1u, 2u}; + cuuint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + const bool allow_oob_box[4] = {false, true, false, false}; + for (int axis = 0; axis < 4; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'Q' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0 && global_strides[2] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'Q' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'Q' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of Q"); + return map; +} + +CUtensorMap EncodeTma_K(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'K' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "K"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'K' requires (d1 == 256)"; + cuuint64_t global_dim[4] = {Extent(64LL, "K"), Extent(d3, "K"), Extent(d2, "K"), Extent(CDiv(d1, 64LL), "K")}; + cuuint64_t global_strides[3] = {Stride(s3, "K"), Stride(s2, "K"), Stride(128LL, "K")}; + cuuint32_t box_dim[4] = {64u, 64u, 1u, 2u}; + cuuint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + const bool allow_oob_box[4] = {false, true, false, false}; + for (int axis = 0; axis < 4; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'K' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0 && global_strides[2] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'K' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'K' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of K"); + return map; +} + +CUtensorMap EncodeTma_V(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'V' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "V"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'V' requires (d1 == 256)"; + cuuint64_t global_dim[4] = {Extent(64LL, "V"), Extent(d3, "V"), Extent(CDiv(d1, 64LL), "V"), Extent(d2, "V")}; + cuuint64_t global_strides[3] = {Stride(s3, "V"), Stride(128LL, "V"), Stride(s2, "V")}; + cuuint32_t box_dim[4] = {64u, 128u, 1u, 1u}; + cuuint32_t elem_strides[4] = {1u, 1u, 1u, 1u}; + const bool allow_oob_box[4] = {false, true, false, false}; + for (int axis = 0; axis < 4; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'V' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0 && global_strides[2] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'V' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'V' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of V"); + return map; +} + +void Run(TensorView a_Q, TensorView a_K, TensorView a_V, TensorView a_O, TensorView a_LSE, TensorView a_cu_seqlens_q, TensorView a_cu_seqlens_k, TensorView a_tile_counter, int64_t a_num_segments, int64_t a_num_heads_q, int64_t a_group, int64_t a_total_q, double a_softmax_scale_log2, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_Q.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_Q, "Q"); + check_dtype(a_Q, DLDataType{kDLBfloat, 16, 1}, "Q"); + check_cuda_tensor(a_K, "K"); + check_dtype(a_K, DLDataType{kDLBfloat, 16, 1}, "K"); + check_cuda_tensor(a_V, "V"); + check_dtype(a_V, DLDataType{kDLBfloat, 16, 1}, "V"); + check_cuda_tensor(a_O, "O"); + check_dtype(a_O, DLDataType{kDLBfloat, 16, 1}, "O"); + check_cuda_tensor(a_LSE, "LSE"); + check_dtype(a_LSE, DLDataType{kDLFloat, 32, 1}, "LSE"); + check_cuda_tensor(a_cu_seqlens_q, "cu_seqlens_q"); + check_dtype(a_cu_seqlens_q, DLDataType{kDLInt, 32, 1}, "cu_seqlens_q"); + check_cuda_tensor(a_cu_seqlens_k, "cu_seqlens_k"); + check_dtype(a_cu_seqlens_k, DLDataType{kDLInt, 32, 1}, "cu_seqlens_k"); + check_cuda_tensor(a_tile_counter, "tile_counter"); + check_dtype(a_tile_counter, DLDataType{kDLUInt, 32, 1}, "tile_counter"); + TVM_FFI_CHECK(a_num_segments >= -2147483648LL && a_num_segments <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_segments' = " << a_num_segments + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_num_heads_q >= -2147483648LL && a_num_heads_q <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_heads_q' = " << a_num_heads_q + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_group >= -2147483648LL && a_group <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'group' = " << a_group + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_total_q >= -2147483648LL && a_total_q <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'total_q' = " << a_total_q + << " is outside the i32 range [-2147483648, 2147483647]"; + check_same_device(a_K, a_Q, "K", "Q"); + check_same_device(a_V, a_Q, "V", "Q"); + check_contiguous(a_O, "O"); + check_same_device(a_O, a_Q, "O", "Q"); + check_contiguous(a_LSE, "LSE"); + check_same_device(a_LSE, a_Q, "LSE", "Q"); + check_contiguous(a_cu_seqlens_q, "cu_seqlens_q"); + check_same_device(a_cu_seqlens_q, a_Q, "cu_seqlens_q", "Q"); + check_contiguous(a_cu_seqlens_k, "cu_seqlens_k"); + check_same_device(a_cu_seqlens_k, a_Q, "cu_seqlens_k", "Q"); + check_contiguous(a_tile_counter, "tile_counter"); + check_same_device(a_tile_counter, a_Q, "tile_counter", "Q"); + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") is not a multiple of the cluster shape (2, 1, 1)"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_gqa_train_493c4a28dafa5fc6b2f1), cudaFuncAttributeMaxDynamicSharedMemorySize, 232064), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_Q = EncodeTma_Q(a_Q); + CUtensorMap v_K = EncodeTma_K(a_K); + CUtensorMap v_V = EncodeTma_V(a_V); + void* v_O = DataPtr(a_O); + void* v_LSE = DataPtr(a_LSE); + void* v_cu_seqlens_q = DataPtr(a_cu_seqlens_q); + void* v_cu_seqlens_k = DataPtr(a_cu_seqlens_k); + void* v_tile_counter = DataPtr(a_tile_counter); + int32_t v_num_segments = static_cast(a_num_segments); + int32_t v_num_heads_q = static_cast(a_num_heads_q); + int32_t v_group = static_cast(a_group); + int32_t v_total_q = static_cast(a_total_q); + float v_softmax_scale_log2 = static_cast(a_softmax_scale_log2); + void* kernel_args[] = {&v_Q, &v_K, &v_V, &v_O, &v_LSE, &v_cu_seqlens_q, &v_cu_seqlens_k, &v_tile_counter, &v_num_segments, &v_num_heads_q, &v_group, &v_total_q, &v_softmax_scale_log2}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(384, 1, 1); + config.dynamicSmemBytes = 232064; + config.stream = stream; + cudaLaunchAttribute attrs[2] = {}; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = 2; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + config.attrs = attrs; + config.numAttrs = 2; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gqa_train_493c4a28dafa5fc6b2f1), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gqa_train_493c4a28dafa5fc6b2f1 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gqa_train_493c4a28dafa5fc6b2f1::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu new file mode 100644 index 0000000..11dc829 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu @@ -0,0 +1,2998 @@ +/* + * 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 TMEM_NCOLS 512 +#define TMEM_S_OFFSET 0 +#define TMEM_P0_OFFSET 64 +#define TMEM_P1_OFFSET 192 +#define TMEM_O_HI_OFFSET 256 +#define TMEM_O_LO_OFFSET 384 +#define NUM_KV_PIPE_STAGES 5 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_S_PIPE_STAGES 2 +#define NUM_P_PIPE_STAGES 2 +#define NUM_O_PIPE_STAGES 2 +#define NUM_OF_PIPE_STAGES 1 +#define NUM_SC_PIPE_STAGES 1 +#define NUM_ST_PIPE_STAGES 1 +#define NUM_WORK_PIPE_STAGES 2 +#define SMEM_S_SCALE_OFF 230400 +#define SMEM_S_SCALE_STAGE_BYTES 512 +#define SMEM_S_SCALE_STRIDE 512 +#define SMEM_S_STATS_OFF 230912 +#define SMEM_S_STATS_STAGE_BYTES 1024 +#define SMEM_S_STATS_STRIDE 1024 +#define SMEM_SMEM_Q_HI_OFF 1024 +#define SMEM_SMEM_Q_HI_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_HI_STRIDE 32768 +#define SMEM_SMEM_Q_LO_OFF 33792 +#define SMEM_SMEM_Q_LO_STAGE_BYTES 32768 +#define SMEM_SMEM_Q_LO_STRIDE 32768 +#define SMEM_SMEM_K_HI_OFF 66560 +#define SMEM_SMEM_K_HI_STAGE_BYTES 16384 +#define SMEM_SMEM_K_HI_STRIDE 16384 +#define SMEM_SMEM_K_LO_OFF 148480 +#define SMEM_SMEM_K_LO_STAGE_BYTES 16384 +#define SMEM_SMEM_K_LO_STRIDE 16384 +#define SMEM_SMEM_V_HI_OFF 66560 +#define SMEM_SMEM_V_HI_STAGE_BYTES 16384 +#define SMEM_SMEM_V_HI_STRIDE 16384 +#define SMEM_SMEM_V_LO_OFF 148480 +#define SMEM_SMEM_V_LO_STAGE_BYTES 16384 +#define SMEM_SMEM_V_LO_STRIDE 16384 +#define SMEM_CELLS_OFF 231936 +#define SMEM_CELLS_STAGE_BYTES 8 +#define SMEM_CELLS_STRIDE 8 +#define SMEM_TOTAL 232064 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + +__device__ __forceinline__ void elect_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + + + + +__device__ __forceinline__ void tmem_st_x32_f32(int tmem_addr, const float* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x32.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8," + " %9, %10, %11, %12, %13, %14, %15, %16," + " %17, %18, %19, %20, %21, %22, %23, %24," + " %25, %26, %27, %28, %29, %30, %31, %32};" + :: "r"(tmem_addr), + "f"(src[0]), "f"(src[1]), "f"(src[2]), "f"(src[3]), + "f"(src[4]), "f"(src[5]), "f"(src[6]), "f"(src[7]), + "f"(src[8]), "f"(src[9]), "f"(src[10]), "f"(src[11]), + "f"(src[12]), "f"(src[13]), "f"(src[14]), "f"(src[15]), + "f"(src[16]), "f"(src[17]), "f"(src[18]), "f"(src[19]), + "f"(src[20]), "f"(src[21]), "f"(src[22]), "f"(src[23]), + "f"(src[24]), "f"(src[25]), "f"(src[26]), "f"(src[27]), + "f"(src[28]), "f"(src[29]), "f"(src[30]), "f"(src[31])); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float approx_rcp(float x) { + float y; + asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + + + +__device__ __forceinline__ float row_max_reduce(float2 acc) { + return max_noftz(acc.x, acc.y); +} + + +__device__ __forceinline__ void row_max_x32_accum(const float* sv, float2& acc) { + #pragma unroll + for (int j = 0; j < 16; j++) { + if (j % 2 == 0) + acc.x = max_noftz(acc.x, max_noftz(sv[j*2], sv[j*2+1])); + else + acc.y = max_noftz(acc.y, max_noftz(sv[j*2], sv[j*2+1])); + } +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void softmax_block_sum(const float* sv, float2* acc) { + const float2* sv2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) { + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "+l"(reinterpret_cast(*acc)) + : "l"(reinterpret_cast(*acc)), + "l"(reinterpret_cast(sv2[j]))); + } +} + + +__device__ __forceinline__ void fma_f32x2_inplace(float2* a, float2 b, float2 c) { + unsigned long long r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(r) + : "l"(*(unsigned long long*)a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + *(unsigned long long*)a = r; +} + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + + + +__device__ __forceinline__ void tma_4d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int w, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4, %5}], [%6];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(w), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast_leader( + int mbar_addr, uint16_t cta_mask, uint32_t leader) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "setp.ne.b32 leader, %2, 0;\n\t" + "@leader tcgen05.commit.cta_group::2.mbarrier::arrive::one.shared::cluster.multicast::cluster.b64 [%0], %1;\n\t" + "}\n" + :: "r"(mbar_addr), "h"(cta_mask), "r"(leader) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(384, LAUNCH_MIN_BLOCKS) __cluster_dims__(2,1,1) void +kernel_cake_gqa_train_493c4a28dafa5fc6b2f1(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, __nv_bfloat16* __restrict__ O, float* __restrict__ LSE, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, unsigned int* __restrict__ tile_counter, int num_segments, int num_heads_q, int group, int total_q, float softmax_scale_log2) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define q_full_addr (mbar_base + 0) + #define q_empty_addr (mbar_base + 8) + #define kv_full_addr (mbar_base + 16) + #define kv_empty_addr (mbar_base + 56) + #define s_full_addr (mbar_base + 96) + #define p_full_addr (mbar_base + 112) + #define p_full_2_addr (mbar_base + 128) + #define o_full_addr (mbar_base + 144) + #define o_free_addr (mbar_base + 160) + #define scale_full_addr (mbar_base + 168) + #define scale_empty_addr (mbar_base + 176) + #define stats_full_addr (mbar_base + 184) + #define stats_empty_addr (mbar_base + 192) + #define work_full_addr (mbar_base + 200) + #define work_empty_addr (mbar_base + 216) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + float* s_scale = reinterpret_cast(smem_raw + 230400); + const int s_scale_addr = smem + 230400; + float* s_stats = reinterpret_cast(smem_raw + 230912); + const int s_stats_addr = smem + 230912; + __nv_bfloat16* smem_q_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_q_hi_addr = smem + 1024; + __nv_bfloat16* smem_q_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int smem_q_lo_addr = smem + 33792; + __nv_bfloat16* smem_k_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_k_hi_addr = smem + 66560; + __nv_bfloat16* smem_k_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 148480); + const int smem_k_lo_addr = smem + 148480; + __nv_bfloat16* smem_v_hi = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int smem_v_hi_addr = smem + 66560; + __nv_bfloat16* smem_v_lo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 148480); + const int smem_v_lo_addr = smem + 148480; + unsigned int* cells = reinterpret_cast(smem_raw + 231936); + const int cells_addr = smem + 231936; + + // Mbarrier init (15 pipeline groups, 0 ordered-sequence groups, 29 barriers) + // Mbarriers at smem_raw[0..232) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=2 + mbarrier_init(smem + 0, 2); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 8, 1); + // --- pipeline 'kv_pipe' --- + // kv_full: 5 barriers, init_count=2 + mbarrier_init(smem + 16, 2); + mbarrier_init(smem + 24, 2); + mbarrier_init(smem + 32, 2); + mbarrier_init(smem + 40, 2); + mbarrier_init(smem + 48, 2); + // kv_empty: 5 barriers, init_count=1 + mbarrier_init(smem + 56, 1); + mbarrier_init(smem + 64, 1); + mbarrier_init(smem + 72, 1); + mbarrier_init(smem + 80, 1); + mbarrier_init(smem + 88, 1); + // --- pipeline 's_pipe' --- + // s_full: 2 barriers, init_count=1 + mbarrier_init(smem + 96, 1); + mbarrier_init(smem + 104, 1); + // --- pipeline 'p_pipe' --- + // p_full: 2 barriers, init_count=512 + mbarrier_init(smem + 112, 512); + mbarrier_init(smem + 120, 512); + // p_full_2: 2 barriers, init_count=256 + mbarrier_init(smem + 128, 256); + mbarrier_init(smem + 136, 256); + // --- pipeline 'o_pipe' --- + // o_full: 2 barriers, init_count=1 + mbarrier_init(smem + 144, 1); + mbarrier_init(smem + 152, 1); + // --- pipeline 'of_pipe' --- + // o_free: 1 barriers, init_count=256 + mbarrier_init(smem + 160, 256); + // --- pipeline 'sc_pipe' --- + // scale_full: 1 barriers, init_count=128 + mbarrier_init(smem + 168, 128); + // scale_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 176, 128); + // --- pipeline 'st_pipe' --- + // stats_full: 1 barriers, init_count=128 + mbarrier_init(smem + 184, 128); + // stats_empty: 1 barriers, init_count=128 + mbarrier_init(smem + 192, 128); + // --- pipeline 'work_pipe' --- + // work_full: 2 barriers, init_count=1 + mbarrier_init(smem + 200, 1); + mbarrier_init(smem + 208, 1); + // work_empty: 2 barriers, init_count=640 + mbarrier_init(smem + 216, 640); + mbarrier_init(smem + 224, 640); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 232); + if (warp == 0) { + int _tmem_hold = smem + 232; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_s = taddr; + const int tmem_p0 = taddr + 64; + const int tmem_p1 = taddr + 192; + const int tmem_o_hi = taddr + 256; + const int tmem_o_lo = taddr + 384; + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); + } + + // ---- Role: softmax ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 256;"); + { // softmax_main + unsigned int wk = 0; + unsigned int sq = 0; + unsigned int pq = 0; + const int tr = warp % 4 * 32 + lane; + unsigned int head_tiles = 0; + #pragma unroll 1 + for (int s = 0; s < num_segments; s++) { + head_tiles = head_tiles + ((unsigned int)(cu_seqlens_q[s + 1] - cu_seqlens_q[s] + 255) >> 8); + } + int wmax = 0; + #pragma unroll 1 + for (int s_1 = 0; s_1 < num_segments; s_1++) { + int q_len = cu_seqlens_q[s_1 + 1] - cu_seqlens_q[s_1]; + int k_len = cu_seqlens_k[s_1 + 1] - cu_seqlens_k[s_1]; + int row0 = ((q_len + 255 >> 8) - 1) * 256; + int diag = k_len - q_len; + int n_k = k_len + 127 >> 7; + int n_c = row0 + 256 + diag + 127 >> 7; + int n_seen = ((n_c < n_k) ? n_c : n_k); + int n_blocks = ((n_seen > 0) ? n_seen : 0); + int first_seen = row0 + diag + 1; + int n_first = ((first_seen > 0) ? first_seen >> 7 : 0); + int k_full = k_len >> 7; + int n_kf = ((n_first < k_full) ? n_first : k_full); + int n_full = ((n_kf < n_blocks) ? n_kf : n_blocks); + int n_masked = n_blocks - n_full; + int w = n_blocks; + wmax = ((w > wmax) ? w : wmax); + } + unsigned int t0 = 0; + unsigned int t1 = 0; + unsigned int t2 = 0; + #pragma unroll 1 + for (int s_2 = 0; s_2 < num_segments; s_2++) { + unsigned int tiles = (unsigned int)(cu_seqlens_q[s_2 + 1] - cu_seqlens_q[s_2] + 255) >> 8; + int q_len_1 = cu_seqlens_q[s_2 + 1] - cu_seqlens_q[s_2]; + int k_len_1 = cu_seqlens_k[s_2 + 1] - cu_seqlens_k[s_2]; + int row0_1 = ((q_len_1 + 255 >> 8) - 1) * 256; + int diag_1 = k_len_1 - q_len_1; + int n_k_1 = k_len_1 + 127 >> 7; + int n_c_1 = row0_1 + 256 + diag_1 + 127 >> 7; + int n_seen_1 = ((n_c_1 < n_k_1) ? n_c_1 : n_k_1); + int n_blocks_1 = ((n_seen_1 > 0) ? n_seen_1 : 0); + int first_seen_1 = row0_1 + diag_1 + 1; + int n_first_1 = ((first_seen_1 > 0) ? first_seen_1 >> 7 : 0); + int k_full_1 = k_len_1 >> 7; + int n_kf_1 = ((n_first_1 < k_full_1) ? n_first_1 : k_full_1); + int n_full_1 = ((n_kf_1 < n_blocks_1) ? n_kf_1 : n_blocks_1); + int n_masked_1 = n_blocks_1 - n_full_1; + int c = ((2 * n_blocks_1 >= wmax) ? 0 : ((4 * n_blocks_1 >= wmax) ? 1 : ((8 * n_blocks_1 >= wmax) ? 2 : 3))); + t0 = t0 + ((c == 0) ? tiles : 0); + t1 = t1 + ((c == 1) ? tiles : 0); + t2 = t2 + ((c == 2) ? tiles : 0); + } + unsigned int total = (unsigned int)num_heads_q * head_tiles; + unsigned int _phase_work_full = 0; + unsigned int _phase_s_full = 0; + unsigned int _phase_scale_empty = 1; + unsigned int _phase_stats_empty = 1; + #pragma unroll 1 + for (unsigned int it = 0; it < (unsigned int)num_heads_q * ((unsigned int)(total_q + 255 >> 8) + (unsigned int)num_segments) + 1; it++) { + mbarrier_wait_cluster_hint(work_full_addr + (wk) * 8, _phase_work_full, 10000000); + unsigned int ticket = cells[wk]; + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + (wk) * 8), "r"(0) : "memory"); + wk += 1; + if (wk == 2) { wk = 0; _phase_work_full ^= 1; } + if (ticket >= total) { + break; + } + unsigned int hq = (unsigned int)num_heads_q; + unsigned int p1 = hq * t0; + unsigned int p2 = p1 + hq * t1; + unsigned int p3 = p2 + hq * t2; + int c_1 = ((ticket < p1) ? 0 : ((ticket < p2) ? 1 : ((ticket < p3) ? 2 : 3))); + unsigned int base = ((c_1 == 0) ? 0 : ((c_1 == 1) ? p1 : ((c_1 == 2) ? p2 : p3))); + unsigned int tc = ((c_1 == 0) ? t0 : ((c_1 == 1) ? t1 : ((c_1 == 2) ? t2 : head_tiles - t0 - t1 - t2))); + unsigned int r = ticket - base; + unsigned int per_kv = tc * (unsigned int)group; + unsigned int kvh_u = r / per_kv; + unsigned int rem = r - kvh_u * per_kv; + unsigned int j = rem / (unsigned int)group; + unsigned int head_u = kvh_u * (unsigned int)group + (rem - j * (unsigned int)group); + int seg = 0; + unsigned int before = 0; + unsigned int ends = 0; + #pragma unroll 1 + for (int s_3 = 0; s_3 < num_segments; s_3++) { + int q_len_2 = cu_seqlens_q[s_3 + 1] - cu_seqlens_q[s_3]; + int k_len_2 = cu_seqlens_k[s_3 + 1] - cu_seqlens_k[s_3]; + int row0_2 = ((q_len_2 + 255 >> 8) - 1) * 256; + int diag_2 = k_len_2 - q_len_2; + int n_k_2 = k_len_2 + 127 >> 7; + int n_c_2 = row0_2 + 256 + diag_2 + 127 >> 7; + int n_seen_2 = ((n_c_2 < n_k_2) ? n_c_2 : n_k_2); + int n_blocks_2 = ((n_seen_2 > 0) ? n_seen_2 : 0); + int first_seen_2 = row0_2 + diag_2 + 1; + int n_first_2 = ((first_seen_2 > 0) ? first_seen_2 >> 7 : 0); + int k_full_2 = k_len_2 >> 7; + int n_kf_2 = ((n_first_2 < k_full_2) ? n_first_2 : k_full_2); + int n_full_2 = ((n_kf_2 < n_blocks_2) ? n_kf_2 : n_blocks_2); + int n_masked_2 = n_blocks_2 - n_full_2; + bool in_class = ((2 * n_blocks_2 >= wmax) ? 0 : ((4 * n_blocks_2 >= wmax) ? 1 : ((8 * n_blocks_2 >= wmax) ? 2 : 3))) == c_1; + ends = ends + ((in_class) ? (unsigned int)(cu_seqlens_q[s_3 + 1] - cu_seqlens_q[s_3] + 255) >> 8 : 0); + seg = seg + ((ends <= j) ? 1 : 0); + before = ((ends <= j) ? ends : before); + } + unsigned int m_rel = j - before; + int q0 = cu_seqlens_q[seg]; + int q_len_3 = cu_seqlens_q[seg + 1] - q0; + int k0 = cu_seqlens_k[seg]; + int k_len_3 = cu_seqlens_k[seg + 1] - k0; + int m = (int)(((unsigned int)(q_len_3 + 255) >> 8) - 1 - m_rel); + int head = (int)head_u; + int kvh = (int)kvh_u; + int row0_3 = m * 256; + int diag_3 = k_len_3 - q_len_3; + int n_k_3 = k_len_3 + 127 >> 7; + int n_c_3 = row0_3 + 256 + diag_3 + 127 >> 7; + int n_seen_3 = ((n_c_3 < n_k_3) ? n_c_3 : n_k_3); + int n_blocks_3 = ((n_seen_3 > 0) ? n_seen_3 : 0); + int first_seen_3 = row0_3 + diag_3 + 1; + int n_first_3 = ((first_seen_3 > 0) ? first_seen_3 >> 7 : 0); + int k_full_3 = k_len_3 >> 7; + int n_kf_3 = ((n_first_3 < k_full_3) ? n_first_3 : k_full_3); + int n_full_3 = ((n_kf_3 < n_blocks_3) ? n_kf_3 : n_blocks_3); + int n_masked_3 = n_blocks_3 - n_full_3; + int row = m * 256 + cta_rank * 128 + tr; + int seen = row + (k_len_3 - q_len_3) + 1; + int vis = ((seen < k_len_3) ? seen : k_len_3); + float row_max = -CUDART_INF_F; + float row_sum = 0.0f; + #pragma unroll 1 + for (int i = 0; i < n_masked_3; i++) { + mbarrier_wait(s_full_addr + (sq) * 8, _phase_s_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int s_addr = taddr + (unsigned int)(sq * 128) + (unsigned int)(warp % 4 * 32 << 16); + float sv[128]; + float tile_max = -CUDART_INF_F; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[0]), "=f"(sv[1]), "=f"(sv[2]), "=f"(sv[3]), "=f"(sv[4]), "=f"(sv[5]), "=f"(sv[6]), "=f"(sv[7]), "=f"(sv[8]), "=f"(sv[9]), "=f"(sv[10]), "=f"(sv[11]), "=f"(sv[12]), "=f"(sv[13]), "=f"(sv[14]), "=f"(sv[15]), "=f"(sv[16]), "=f"(sv[17]), "=f"(sv[18]), "=f"(sv[19]), "=f"(sv[20]), "=f"(sv[21]), "=f"(sv[22]), "=f"(sv[23]), "=f"(sv[24]), "=f"(sv[25]), "=f"(sv[26]), "=f"(sv[27]), "=f"(sv[28]), "=f"(sv[29]), "=f"(sv[30]), "=f"(sv[31]) + : "r"(s_addr)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[32]), "=f"(sv[33]), "=f"(sv[34]), "=f"(sv[35]), "=f"(sv[36]), "=f"(sv[37]), "=f"(sv[38]), "=f"(sv[39]), "=f"(sv[40]), "=f"(sv[41]), "=f"(sv[42]), "=f"(sv[43]), "=f"(sv[44]), "=f"(sv[45]), "=f"(sv[46]), "=f"(sv[47]), "=f"(sv[48]), "=f"(sv[49]), "=f"(sv[50]), "=f"(sv[51]), "=f"(sv[52]), "=f"(sv[53]), "=f"(sv[54]), "=f"(sv[55]), "=f"(sv[56]), "=f"(sv[57]), "=f"(sv[58]), "=f"(sv[59]), "=f"(sv[60]), "=f"(sv[61]), "=f"(sv[62]), "=f"(sv[63]) + : "r"(s_addr + 32)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[64]), "=f"(sv[65]), "=f"(sv[66]), "=f"(sv[67]), "=f"(sv[68]), "=f"(sv[69]), "=f"(sv[70]), "=f"(sv[71]), "=f"(sv[72]), "=f"(sv[73]), "=f"(sv[74]), "=f"(sv[75]), "=f"(sv[76]), "=f"(sv[77]), "=f"(sv[78]), "=f"(sv[79]), "=f"(sv[80]), "=f"(sv[81]), "=f"(sv[82]), "=f"(sv[83]), "=f"(sv[84]), "=f"(sv[85]), "=f"(sv[86]), "=f"(sv[87]), "=f"(sv[88]), "=f"(sv[89]), "=f"(sv[90]), "=f"(sv[91]), "=f"(sv[92]), "=f"(sv[93]), "=f"(sv[94]), "=f"(sv[95]) + : "r"(s_addr + 64)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv[96]), "=f"(sv[97]), "=f"(sv[98]), "=f"(sv[99]), "=f"(sv[100]), "=f"(sv[101]), "=f"(sv[102]), "=f"(sv[103]), "=f"(sv[104]), "=f"(sv[105]), "=f"(sv[106]), "=f"(sv[107]), "=f"(sv[108]), "=f"(sv[109]), "=f"(sv[110]), "=f"(sv[111]), "=f"(sv[112]), "=f"(sv[113]), "=f"(sv[114]), "=f"(sv[115]), "=f"(sv[116]), "=f"(sv[117]), "=f"(sv[118]), "=f"(sv[119]), "=f"(sv[120]), "=f"(sv[121]), "=f"(sv[122]), "=f"(sv[123]), "=f"(sv[124]), "=f"(sv[125]), "=f"(sv[126]), "=f"(sv[127]) + : "r"(s_addr + 96)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + int lim = vis - (n_blocks_3 - 1 - i) * 128; + uint32_t _thresh_mask_0; + { + int _lim_1 = lim; + if (_lim_1 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_1)); + } + } + if (!(_thresh_mask_0 & (1u << 0))) sv[0] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 1))) sv[1] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 2))) sv[2] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 3))) sv[3] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 4))) sv[4] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 5))) sv[5] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 6))) sv[6] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 7))) sv[7] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 8))) sv[8] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 9))) sv[9] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 10))) sv[10] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 11))) sv[11] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 12))) sv[12] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 13))) sv[13] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 14))) sv[14] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 15))) sv[15] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 16))) sv[16] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 17))) sv[17] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 18))) sv[18] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 19))) sv[19] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 20))) sv[20] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 21))) sv[21] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 22))) sv[22] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 23))) sv[23] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 24))) sv[24] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 25))) sv[25] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 26))) sv[26] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 27))) sv[27] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 28))) sv[28] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 29))) sv[29] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 30))) sv[30] = -CUDART_INF_F; + if (!(_thresh_mask_0 & (1u << 31))) sv[31] = -CUDART_INF_F; + uint32_t _thresh_mask_2; + { + int _lim_3 = lim - 32; + if (_lim_3 <= 0) { _thresh_mask_2 = 0u; } + else if (_lim_3 >= 32) { _thresh_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_2) : "r"(_lim_3)); + } + } + if (!(_thresh_mask_2 & (1u << 0))) sv[32] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 1))) sv[33] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 2))) sv[34] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 3))) sv[35] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 4))) sv[36] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 5))) sv[37] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 6))) sv[38] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 7))) sv[39] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 8))) sv[40] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 9))) sv[41] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 10))) sv[42] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 11))) sv[43] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 12))) sv[44] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 13))) sv[45] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 14))) sv[46] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 15))) sv[47] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 16))) sv[48] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 17))) sv[49] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 18))) sv[50] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 19))) sv[51] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 20))) sv[52] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 21))) sv[53] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 22))) sv[54] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 23))) sv[55] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 24))) sv[56] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 25))) sv[57] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 26))) sv[58] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 27))) sv[59] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 28))) sv[60] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 29))) sv[61] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 30))) sv[62] = -CUDART_INF_F; + if (!(_thresh_mask_2 & (1u << 31))) sv[63] = -CUDART_INF_F; + uint32_t _thresh_mask_4; + { + int _lim_5 = lim - 64; + if (_lim_5 <= 0) { _thresh_mask_4 = 0u; } + else if (_lim_5 >= 32) { _thresh_mask_4 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_4) : "r"(_lim_5)); + } + } + if (!(_thresh_mask_4 & (1u << 0))) sv[64] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 1))) sv[65] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 2))) sv[66] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 3))) sv[67] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 4))) sv[68] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 5))) sv[69] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 6))) sv[70] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 7))) sv[71] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 8))) sv[72] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 9))) sv[73] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 10))) sv[74] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 11))) sv[75] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 12))) sv[76] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 13))) sv[77] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 14))) sv[78] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 15))) sv[79] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 16))) sv[80] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 17))) sv[81] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 18))) sv[82] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 19))) sv[83] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 20))) sv[84] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 21))) sv[85] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 22))) sv[86] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 23))) sv[87] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 24))) sv[88] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 25))) sv[89] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 26))) sv[90] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 27))) sv[91] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 28))) sv[92] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 29))) sv[93] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 30))) sv[94] = -CUDART_INF_F; + if (!(_thresh_mask_4 & (1u << 31))) sv[95] = -CUDART_INF_F; + uint32_t _thresh_mask_6; + { + int _lim_7 = lim - 96; + if (_lim_7 <= 0) { _thresh_mask_6 = 0u; } + else if (_lim_7 >= 32) { _thresh_mask_6 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_6) : "r"(_lim_7)); + } + } + if (!(_thresh_mask_6 & (1u << 0))) sv[96] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 1))) sv[97] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 2))) sv[98] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 3))) sv[99] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 4))) sv[100] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 5))) sv[101] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 6))) sv[102] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 7))) sv[103] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 8))) sv[104] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 9))) sv[105] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 10))) sv[106] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 11))) sv[107] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 12))) sv[108] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 13))) sv[109] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 14))) sv[110] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 15))) sv[111] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 16))) sv[112] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 17))) sv[113] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 18))) sv[114] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 19))) sv[115] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 20))) sv[116] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 21))) sv[117] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 22))) sv[118] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 23))) sv[119] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 24))) sv[120] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 25))) sv[121] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 26))) sv[122] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 27))) sv[123] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 28))) sv[124] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 29))) sv[125] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 30))) sv[126] = -CUDART_INF_F; + if (!(_thresh_mask_6 & (1u << 31))) sv[127] = -CUDART_INF_F; + float2 _reg_reduce_max2_9 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&sv[0], _reg_reduce_max2_9); + row_max_x32_accum(&sv[32], _reg_reduce_max2_9); + row_max_x32_accum(&sv[64], _reg_reduce_max2_9); + row_max_x32_accum(&sv[96], _reg_reduce_max2_9); + float _reg_array_reduce_8 = row_max_reduce(_reg_reduce_max2_9); + tile_max = _reg_array_reduce_8; + float new_max; + asm("max.f32 %0, %1, %2;" : "=f"(new_max) : "f"(tile_max), "f"(row_max)); + float new_max_scaled = ((new_max == -CUDART_INF_F) ? 0.0f : new_max) * softmax_scale_log2; + float log2_ratio = __fmaf_rn(row_max, softmax_scale_log2, -new_max_scaled); + float kept_max = row_max; + float acc_scale = 1.0f; + if (log2_ratio >= -8.0f) { + new_max_scaled = ((row_max == -CUDART_INF_F) ? 0.0f : row_max) * softmax_scale_log2; + } else { + kept_max = new_max; + if (row_max > -CUDART_INF_F) { + float _exp2_10 = approx_exp2(log2_ratio); + acc_scale = _exp2_10; + } + } + mbarrier_wait(scale_empty_addr, _phase_scale_empty); + s_scale[(int)0 * 128 + tr] = acc_scale; + mbarrier_arrive(scale_full_addr); + _phase_scale_empty ^= 1; + const float2 _fma_b2_11 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_12 = {-new_max_scaled, -new_max_scaled}; + #pragma unroll + for (int _lf = 0; _lf < 64; _lf++) + fma_f32x2_inplace(&reinterpret_cast(sv)[_lf], _fma_b2_11, _fma_c2_12); + unsigned int p_addr = s_addr + 64; + unsigned int pk_0[16]; + softmax_frag_exp2_cast(&sv[0], pk_0, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 0), "r"(*reinterpret_cast(&pk_0[0])), "r"(*reinterpret_cast(&pk_0[1])), "r"(*reinterpret_cast(&pk_0[2])), "r"(*reinterpret_cast(&pk_0[3])), "r"(*reinterpret_cast(&pk_0[4])), "r"(*reinterpret_cast(&pk_0[5])), "r"(*reinterpret_cast(&pk_0[6])), "r"(*reinterpret_cast(&pk_0[7])), "r"(*reinterpret_cast(&pk_0[8])), "r"(*reinterpret_cast(&pk_0[9])), "r"(*reinterpret_cast(&pk_0[10])), "r"(*reinterpret_cast(&pk_0[11])), "r"(*reinterpret_cast(&pk_0[12])), "r"(*reinterpret_cast(&pk_0[13])), "r"(*reinterpret_cast(&pk_0[14])), "r"(*reinterpret_cast(&pk_0[15]))); + unsigned int pk_1[16]; + softmax_frag_exp2_cast(&sv[32], pk_1, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 16), "r"(*reinterpret_cast(&pk_1[0])), "r"(*reinterpret_cast(&pk_1[1])), "r"(*reinterpret_cast(&pk_1[2])), "r"(*reinterpret_cast(&pk_1[3])), "r"(*reinterpret_cast(&pk_1[4])), "r"(*reinterpret_cast(&pk_1[5])), "r"(*reinterpret_cast(&pk_1[6])), "r"(*reinterpret_cast(&pk_1[7])), "r"(*reinterpret_cast(&pk_1[8])), "r"(*reinterpret_cast(&pk_1[9])), "r"(*reinterpret_cast(&pk_1[10])), "r"(*reinterpret_cast(&pk_1[11])), "r"(*reinterpret_cast(&pk_1[12])), "r"(*reinterpret_cast(&pk_1[13])), "r"(*reinterpret_cast(&pk_1[14])), "r"(*reinterpret_cast(&pk_1[15]))); + unsigned int pk_2[16]; + softmax_frag_exp2_cast(&sv[64], pk_2, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 32), "r"(*reinterpret_cast(&pk_2[0])), "r"(*reinterpret_cast(&pk_2[1])), "r"(*reinterpret_cast(&pk_2[2])), "r"(*reinterpret_cast(&pk_2[3])), "r"(*reinterpret_cast(&pk_2[4])), "r"(*reinterpret_cast(&pk_2[5])), "r"(*reinterpret_cast(&pk_2[6])), "r"(*reinterpret_cast(&pk_2[7])), "r"(*reinterpret_cast(&pk_2[8])), "r"(*reinterpret_cast(&pk_2[9])), "r"(*reinterpret_cast(&pk_2[10])), "r"(*reinterpret_cast(&pk_2[11])), "r"(*reinterpret_cast(&pk_2[12])), "r"(*reinterpret_cast(&pk_2[13])), "r"(*reinterpret_cast(&pk_2[14])), "r"(*reinterpret_cast(&pk_2[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((p_full_addr + (pq) * 8) & 0xFEFFFFFF) : "memory"); + unsigned int pk_3[16]; + softmax_frag_exp2_cast(&sv[96], pk_3, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr + 48), "r"(*reinterpret_cast(&pk_3[0])), "r"(*reinterpret_cast(&pk_3[1])), "r"(*reinterpret_cast(&pk_3[2])), "r"(*reinterpret_cast(&pk_3[3])), "r"(*reinterpret_cast(&pk_3[4])), "r"(*reinterpret_cast(&pk_3[5])), "r"(*reinterpret_cast(&pk_3[6])), "r"(*reinterpret_cast(&pk_3[7])), "r"(*reinterpret_cast(&pk_3[8])), "r"(*reinterpret_cast(&pk_3[9])), "r"(*reinterpret_cast(&pk_3[10])), "r"(*reinterpret_cast(&pk_3[11])), "r"(*reinterpret_cast(&pk_3[12])), "r"(*reinterpret_cast(&pk_3[13])), "r"(*reinterpret_cast(&pk_3[14])), "r"(*reinterpret_cast(&pk_3[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((p_full_2_addr + (pq) * 8) & 0xFEFFFFFF) : "memory"); + pq += 1; + if (pq == 2) { pq = 0; } + sq += 1; + if (sq == 2) { sq = 0; _phase_s_full ^= 1; } + row_max = kept_max; + float2 _reg_reduce_sum2_14 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv[0], &_reg_reduce_sum2_14); + softmax_block_sum(&sv[32], &_reg_reduce_sum2_14); + softmax_block_sum(&sv[64], &_reg_reduce_sum2_14); + softmax_block_sum(&sv[96], &_reg_reduce_sum2_14); + float _reg_array_reduce_13 = _reg_reduce_sum2_14.x + _reg_reduce_sum2_14.y; + row_sum = __fmaf_rn(row_sum, acc_scale, _reg_array_reduce_13); + } + #pragma unroll 1 + for (int i_1 = n_masked_3; i_1 < n_blocks_3; i_1++) { + mbarrier_wait(s_full_addr + (sq) * 8, _phase_s_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + unsigned int s_addr_1 = taddr + (unsigned int)(sq * 128) + (unsigned int)(warp % 4 * 32 << 16); + float sv_1[128]; + float tile_max_1 = -CUDART_INF_F; +#if __CUDA_ARCH__ == 1030 + float lo_max = -CUDART_INF_F; + float hi_max = -CUDART_INF_F; + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif +#endif + asm volatile( +#if __CUDA_ARCH__ == 1030 + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv_1[0]), "=f"(sv_1[1]), "=f"(sv_1[2]), "=f"(sv_1[3]), "=f"(sv_1[4]), "=f"(sv_1[5]), "=f"(sv_1[6]), "=f"(sv_1[7]), "=f"(sv_1[8]), "=f"(sv_1[9]), "=f"(sv_1[10]), "=f"(sv_1[11]), "=f"(sv_1[12]), "=f"(sv_1[13]), "=f"(sv_1[14]), "=f"(sv_1[15]), "=f"(sv_1[16]), "=f"(sv_1[17]), "=f"(sv_1[18]), "=f"(sv_1[19]), "=f"(sv_1[20]), "=f"(sv_1[21]), "=f"(sv_1[22]), "=f"(sv_1[23]), "=f"(sv_1[24]), "=f"(sv_1[25]), "=f"(sv_1[26]), "=f"(sv_1[27]), "=f"(sv_1[28]), "=f"(sv_1[29]), "=f"(sv_1[30]), "=f"(sv_1[31]), "=f"(sv_1[32]), "=f"(sv_1[33]), "=f"(sv_1[34]), "=f"(sv_1[35]), "=f"(sv_1[36]), "=f"(sv_1[37]), "=f"(sv_1[38]), "=f"(sv_1[39]), "=f"(sv_1[40]), "=f"(sv_1[41]), "=f"(sv_1[42]), "=f"(sv_1[43]), "=f"(sv_1[44]), "=f"(sv_1[45]), "=f"(sv_1[46]), "=f"(sv_1[47]), "=f"(sv_1[48]), "=f"(sv_1[49]), "=f"(sv_1[50]), "=f"(sv_1[51]), "=f"(sv_1[52]), "=f"(sv_1[53]), "=f"(sv_1[54]), "=f"(sv_1[55]), "=f"(sv_1[56]), "=f"(sv_1[57]), "=f"(sv_1[58]), "=f"(sv_1[59]), "=f"(sv_1[60]), "=f"(sv_1[61]), "=f"(sv_1[62]), "=f"(sv_1[63]), "=f"(lo_max) +#else + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_1[0]), "=f"(sv_1[1]), "=f"(sv_1[2]), "=f"(sv_1[3]), "=f"(sv_1[4]), "=f"(sv_1[5]), "=f"(sv_1[6]), "=f"(sv_1[7]), "=f"(sv_1[8]), "=f"(sv_1[9]), "=f"(sv_1[10]), "=f"(sv_1[11]), "=f"(sv_1[12]), "=f"(sv_1[13]), "=f"(sv_1[14]), "=f"(sv_1[15]), "=f"(sv_1[16]), "=f"(sv_1[17]), "=f"(sv_1[18]), "=f"(sv_1[19]), "=f"(sv_1[20]), "=f"(sv_1[21]), "=f"(sv_1[22]), "=f"(sv_1[23]), "=f"(sv_1[24]), "=f"(sv_1[25]), "=f"(sv_1[26]), "=f"(sv_1[27]), "=f"(sv_1[28]), "=f"(sv_1[29]), "=f"(sv_1[30]), "=f"(sv_1[31]) +#endif + : "r"(s_addr_1)); +#if __CUDA_ARCH__ == 1030 + #if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ == 1000) + #error "TmemLoadRed requires tcgen05.ld.red support (sm_103/sm_101-sm_110 family), not sm_100" + #endif +#endif + asm volatile( +#if __CUDA_ARCH__ == 1030 + "tcgen05.ld.red.sync.aligned.32x32b.x64.max.f32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31, %32, %33, %34, %35, %36, %37, %38, %39, %40, %41, %42, %43, %44, %45, %46, %47, %48, %49, %50, %51, %52, %53, %54, %55, %56, %57, %58, %59, %60, %61, %62, %63}, %64, [%65];" + : "=f"(sv_1[64]), "=f"(sv_1[65]), "=f"(sv_1[66]), "=f"(sv_1[67]), "=f"(sv_1[68]), "=f"(sv_1[69]), "=f"(sv_1[70]), "=f"(sv_1[71]), "=f"(sv_1[72]), "=f"(sv_1[73]), "=f"(sv_1[74]), "=f"(sv_1[75]), "=f"(sv_1[76]), "=f"(sv_1[77]), "=f"(sv_1[78]), "=f"(sv_1[79]), "=f"(sv_1[80]), "=f"(sv_1[81]), "=f"(sv_1[82]), "=f"(sv_1[83]), "=f"(sv_1[84]), "=f"(sv_1[85]), "=f"(sv_1[86]), "=f"(sv_1[87]), "=f"(sv_1[88]), "=f"(sv_1[89]), "=f"(sv_1[90]), "=f"(sv_1[91]), "=f"(sv_1[92]), "=f"(sv_1[93]), "=f"(sv_1[94]), "=f"(sv_1[95]), "=f"(sv_1[96]), "=f"(sv_1[97]), "=f"(sv_1[98]), "=f"(sv_1[99]), "=f"(sv_1[100]), "=f"(sv_1[101]), "=f"(sv_1[102]), "=f"(sv_1[103]), "=f"(sv_1[104]), "=f"(sv_1[105]), "=f"(sv_1[106]), "=f"(sv_1[107]), "=f"(sv_1[108]), "=f"(sv_1[109]), "=f"(sv_1[110]), "=f"(sv_1[111]), "=f"(sv_1[112]), "=f"(sv_1[113]), "=f"(sv_1[114]), "=f"(sv_1[115]), "=f"(sv_1[116]), "=f"(sv_1[117]), "=f"(sv_1[118]), "=f"(sv_1[119]), "=f"(sv_1[120]), "=f"(sv_1[121]), "=f"(sv_1[122]), "=f"(sv_1[123]), "=f"(sv_1[124]), "=f"(sv_1[125]), "=f"(sv_1[126]), "=f"(sv_1[127]), "=f"(hi_max) +#else + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_1[32]), "=f"(sv_1[33]), "=f"(sv_1[34]), "=f"(sv_1[35]), "=f"(sv_1[36]), "=f"(sv_1[37]), "=f"(sv_1[38]), "=f"(sv_1[39]), "=f"(sv_1[40]), "=f"(sv_1[41]), "=f"(sv_1[42]), "=f"(sv_1[43]), "=f"(sv_1[44]), "=f"(sv_1[45]), "=f"(sv_1[46]), "=f"(sv_1[47]), "=f"(sv_1[48]), "=f"(sv_1[49]), "=f"(sv_1[50]), "=f"(sv_1[51]), "=f"(sv_1[52]), "=f"(sv_1[53]), "=f"(sv_1[54]), "=f"(sv_1[55]), "=f"(sv_1[56]), "=f"(sv_1[57]), "=f"(sv_1[58]), "=f"(sv_1[59]), "=f"(sv_1[60]), "=f"(sv_1[61]), "=f"(sv_1[62]), "=f"(sv_1[63]) + : "r"(s_addr_1 + 32)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_1[64]), "=f"(sv_1[65]), "=f"(sv_1[66]), "=f"(sv_1[67]), "=f"(sv_1[68]), "=f"(sv_1[69]), "=f"(sv_1[70]), "=f"(sv_1[71]), "=f"(sv_1[72]), "=f"(sv_1[73]), "=f"(sv_1[74]), "=f"(sv_1[75]), "=f"(sv_1[76]), "=f"(sv_1[77]), "=f"(sv_1[78]), "=f"(sv_1[79]), "=f"(sv_1[80]), "=f"(sv_1[81]), "=f"(sv_1[82]), "=f"(sv_1[83]), "=f"(sv_1[84]), "=f"(sv_1[85]), "=f"(sv_1[86]), "=f"(sv_1[87]), "=f"(sv_1[88]), "=f"(sv_1[89]), "=f"(sv_1[90]), "=f"(sv_1[91]), "=f"(sv_1[92]), "=f"(sv_1[93]), "=f"(sv_1[94]), "=f"(sv_1[95]) +#endif + : "r"(s_addr_1 + 64)); +#if __CUDA_ARCH__ == 1030 +#else + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(sv_1[96]), "=f"(sv_1[97]), "=f"(sv_1[98]), "=f"(sv_1[99]), "=f"(sv_1[100]), "=f"(sv_1[101]), "=f"(sv_1[102]), "=f"(sv_1[103]), "=f"(sv_1[104]), "=f"(sv_1[105]), "=f"(sv_1[106]), "=f"(sv_1[107]), "=f"(sv_1[108]), "=f"(sv_1[109]), "=f"(sv_1[110]), "=f"(sv_1[111]), "=f"(sv_1[112]), "=f"(sv_1[113]), "=f"(sv_1[114]), "=f"(sv_1[115]), "=f"(sv_1[116]), "=f"(sv_1[117]), "=f"(sv_1[118]), "=f"(sv_1[119]), "=f"(sv_1[120]), "=f"(sv_1[121]), "=f"(sv_1[122]), "=f"(sv_1[123]), "=f"(sv_1[124]), "=f"(sv_1[125]), "=f"(sv_1[126]), "=f"(sv_1[127]) + : "r"(s_addr_1 + 96)); +#endif + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); +#if __CUDA_ARCH__ == 1030 + tile_max_1 = max_noftz(lo_max, hi_max); +#else + float2 _reg_reduce_max2_16 = {-CAKE_INF, -CAKE_INF}; + row_max_x32_accum(&sv_1[0], _reg_reduce_max2_16); + row_max_x32_accum(&sv_1[32], _reg_reduce_max2_16); + row_max_x32_accum(&sv_1[64], _reg_reduce_max2_16); + row_max_x32_accum(&sv_1[96], _reg_reduce_max2_16); + float _reg_array_reduce_15 = row_max_reduce(_reg_reduce_max2_16); + tile_max_1 = _reg_array_reduce_15; +#endif + float new_max_1; + asm("max.f32 %0, %1, %2;" : "=f"(new_max_1) : "f"(tile_max_1), "f"(row_max)); + float new_max_scaled_1 = ((new_max_1 == -CUDART_INF_F) ? 0.0f : new_max_1) * softmax_scale_log2; + float log2_ratio_1 = __fmaf_rn(row_max, softmax_scale_log2, -new_max_scaled_1); + float kept_max_1 = row_max; + float acc_scale_1 = 1.0f; + if (log2_ratio_1 >= -8.0f) { + new_max_scaled_1 = ((row_max == -CUDART_INF_F) ? 0.0f : row_max) * softmax_scale_log2; + } else { + kept_max_1 = new_max_1; + if (row_max > -CUDART_INF_F) { +#if __CUDA_ARCH__ == 1030 + float _exp2_15 = approx_exp2(log2_ratio_1); + acc_scale_1 = _exp2_15; +#else + float _exp2_17 = approx_exp2(log2_ratio_1); + acc_scale_1 = _exp2_17; +#endif + } + } + mbarrier_wait(scale_empty_addr, _phase_scale_empty); + s_scale[(int)0 * 128 + tr] = acc_scale_1; + mbarrier_arrive(scale_full_addr); + _phase_scale_empty ^= 1; +#if __CUDA_ARCH__ == 1030 + const float2 _fma_b2_16 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_17 = {-new_max_scaled_1, -new_max_scaled_1}; +#else + const float2 _fma_b2_18 = {softmax_scale_log2, softmax_scale_log2}; + const float2 _fma_c2_19 = {-new_max_scaled_1, -new_max_scaled_1}; +#endif + #pragma unroll + for (int _lf = 0; _lf < 64; _lf++) +#if __CUDA_ARCH__ == 1030 + fma_f32x2_inplace(&reinterpret_cast(sv_1)[_lf], _fma_b2_16, _fma_c2_17); +#else + fma_f32x2_inplace(&reinterpret_cast(sv_1)[_lf], _fma_b2_18, _fma_c2_19); +#endif + unsigned int p_addr_1 = s_addr_1 + 64; + unsigned int pk_4[16]; + softmax_frag_exp2_cast(&sv_1[0], pk_4, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_1 + 0), "r"(*reinterpret_cast(&pk_4[0])), "r"(*reinterpret_cast(&pk_4[1])), "r"(*reinterpret_cast(&pk_4[2])), "r"(*reinterpret_cast(&pk_4[3])), "r"(*reinterpret_cast(&pk_4[4])), "r"(*reinterpret_cast(&pk_4[5])), "r"(*reinterpret_cast(&pk_4[6])), "r"(*reinterpret_cast(&pk_4[7])), "r"(*reinterpret_cast(&pk_4[8])), "r"(*reinterpret_cast(&pk_4[9])), "r"(*reinterpret_cast(&pk_4[10])), "r"(*reinterpret_cast(&pk_4[11])), "r"(*reinterpret_cast(&pk_4[12])), "r"(*reinterpret_cast(&pk_4[13])), "r"(*reinterpret_cast(&pk_4[14])), "r"(*reinterpret_cast(&pk_4[15]))); + unsigned int pk_5[16]; + softmax_frag_exp2_cast(&sv_1[32], pk_5, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_1 + 16), "r"(*reinterpret_cast(&pk_5[0])), "r"(*reinterpret_cast(&pk_5[1])), "r"(*reinterpret_cast(&pk_5[2])), "r"(*reinterpret_cast(&pk_5[3])), "r"(*reinterpret_cast(&pk_5[4])), "r"(*reinterpret_cast(&pk_5[5])), "r"(*reinterpret_cast(&pk_5[6])), "r"(*reinterpret_cast(&pk_5[7])), "r"(*reinterpret_cast(&pk_5[8])), "r"(*reinterpret_cast(&pk_5[9])), "r"(*reinterpret_cast(&pk_5[10])), "r"(*reinterpret_cast(&pk_5[11])), "r"(*reinterpret_cast(&pk_5[12])), "r"(*reinterpret_cast(&pk_5[13])), "r"(*reinterpret_cast(&pk_5[14])), "r"(*reinterpret_cast(&pk_5[15]))); + unsigned int pk_6[16]; + softmax_frag_exp2_cast(&sv_1[64], pk_6, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_1 + 32), "r"(*reinterpret_cast(&pk_6[0])), "r"(*reinterpret_cast(&pk_6[1])), "r"(*reinterpret_cast(&pk_6[2])), "r"(*reinterpret_cast(&pk_6[3])), "r"(*reinterpret_cast(&pk_6[4])), "r"(*reinterpret_cast(&pk_6[5])), "r"(*reinterpret_cast(&pk_6[6])), "r"(*reinterpret_cast(&pk_6[7])), "r"(*reinterpret_cast(&pk_6[8])), "r"(*reinterpret_cast(&pk_6[9])), "r"(*reinterpret_cast(&pk_6[10])), "r"(*reinterpret_cast(&pk_6[11])), "r"(*reinterpret_cast(&pk_6[12])), "r"(*reinterpret_cast(&pk_6[13])), "r"(*reinterpret_cast(&pk_6[14])), "r"(*reinterpret_cast(&pk_6[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((p_full_addr + (pq) * 8) & 0xFEFFFFFF) : "memory"); + unsigned int pk_7[16]; + softmax_frag_exp2_cast(&sv_1[96], pk_7, 0); + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x16.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(p_addr_1 + 48), "r"(*reinterpret_cast(&pk_7[0])), "r"(*reinterpret_cast(&pk_7[1])), "r"(*reinterpret_cast(&pk_7[2])), "r"(*reinterpret_cast(&pk_7[3])), "r"(*reinterpret_cast(&pk_7[4])), "r"(*reinterpret_cast(&pk_7[5])), "r"(*reinterpret_cast(&pk_7[6])), "r"(*reinterpret_cast(&pk_7[7])), "r"(*reinterpret_cast(&pk_7[8])), "r"(*reinterpret_cast(&pk_7[9])), "r"(*reinterpret_cast(&pk_7[10])), "r"(*reinterpret_cast(&pk_7[11])), "r"(*reinterpret_cast(&pk_7[12])), "r"(*reinterpret_cast(&pk_7[13])), "r"(*reinterpret_cast(&pk_7[14])), "r"(*reinterpret_cast(&pk_7[15]))); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((p_full_2_addr + (pq) * 8) & 0xFEFFFFFF) : "memory"); + pq += 1; + if (pq == 2) { pq = 0; } + sq += 1; + if (sq == 2) { sq = 0; _phase_s_full ^= 1; } + row_max = kept_max_1; +#if __CUDA_ARCH__ == 1030 + float2 _reg_reduce_sum2_19 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv_1[0], &_reg_reduce_sum2_19); + softmax_block_sum(&sv_1[32], &_reg_reduce_sum2_19); + softmax_block_sum(&sv_1[64], &_reg_reduce_sum2_19); + softmax_block_sum(&sv_1[96], &_reg_reduce_sum2_19); + float _reg_array_reduce_18 = _reg_reduce_sum2_19.x + _reg_reduce_sum2_19.y; + row_sum = __fmaf_rn(row_sum, acc_scale_1, _reg_array_reduce_18); +#else + float2 _reg_reduce_sum2_21 = make_float2(0.0f, 0.0f); + softmax_block_sum(&sv_1[0], &_reg_reduce_sum2_21); + softmax_block_sum(&sv_1[32], &_reg_reduce_sum2_21); + softmax_block_sum(&sv_1[64], &_reg_reduce_sum2_21); + softmax_block_sum(&sv_1[96], &_reg_reduce_sum2_21); + float _reg_array_reduce_20 = _reg_reduce_sum2_21.x + _reg_reduce_sum2_21.y; + row_sum = __fmaf_rn(row_sum, acc_scale_1, _reg_array_reduce_20); +#endif + } + mbarrier_wait(stats_empty_addr, _phase_stats_empty); + s_stats[tr] = row_sum; + s_stats[128 + tr] = row_max; + mbarrier_arrive(stats_full_addr); + _phase_stats_empty ^= 1; + } + } + } + // ---- Role: correction ---- + if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 160;"); + { // correction_main + unsigned int wk_1 = 0; + unsigned int pq_1 = 0; + unsigned int oq = 0; + const int tr_1 = warp % 4 * 32 + lane; + unsigned int head_tiles_1 = 0; + #pragma unroll 1 + for (int s_4 = 0; s_4 < num_segments; s_4++) { + head_tiles_1 = head_tiles_1 + ((unsigned int)(cu_seqlens_q[s_4 + 1] - cu_seqlens_q[s_4] + 255) >> 8); + } + int wmax_1 = 0; + #pragma unroll 1 + for (int s_5 = 0; s_5 < num_segments; s_5++) { + int q_len_4 = cu_seqlens_q[s_5 + 1] - cu_seqlens_q[s_5]; + int k_len_4 = cu_seqlens_k[s_5 + 1] - cu_seqlens_k[s_5]; + int row0_4 = ((q_len_4 + 255 >> 8) - 1) * 256; + int diag_4 = k_len_4 - q_len_4; + int n_k_4 = k_len_4 + 127 >> 7; + int n_c_4 = row0_4 + 256 + diag_4 + 127 >> 7; + int n_seen_4 = ((n_c_4 < n_k_4) ? n_c_4 : n_k_4); + int n_blocks_4 = ((n_seen_4 > 0) ? n_seen_4 : 0); + int first_seen_4 = row0_4 + diag_4 + 1; + int n_first_4 = ((first_seen_4 > 0) ? first_seen_4 >> 7 : 0); + int k_full_4 = k_len_4 >> 7; + int n_kf_4 = ((n_first_4 < k_full_4) ? n_first_4 : k_full_4); + int n_full_4 = ((n_kf_4 < n_blocks_4) ? n_kf_4 : n_blocks_4); + int n_masked_4 = n_blocks_4 - n_full_4; + int w_1 = n_blocks_4; + wmax_1 = ((w_1 > wmax_1) ? w_1 : wmax_1); + } + unsigned int t0_1 = 0; + unsigned int t1_1 = 0; + unsigned int t2_1 = 0; + #pragma unroll 1 + for (int s_6 = 0; s_6 < num_segments; s_6++) { + unsigned int tiles_1 = (unsigned int)(cu_seqlens_q[s_6 + 1] - cu_seqlens_q[s_6] + 255) >> 8; + int q_len_5 = cu_seqlens_q[s_6 + 1] - cu_seqlens_q[s_6]; + int k_len_5 = cu_seqlens_k[s_6 + 1] - cu_seqlens_k[s_6]; + int row0_5 = ((q_len_5 + 255 >> 8) - 1) * 256; + int diag_5 = k_len_5 - q_len_5; + int n_k_5 = k_len_5 + 127 >> 7; + int n_c_5 = row0_5 + 256 + diag_5 + 127 >> 7; + int n_seen_5 = ((n_c_5 < n_k_5) ? n_c_5 : n_k_5); + int n_blocks_5 = ((n_seen_5 > 0) ? n_seen_5 : 0); + int first_seen_5 = row0_5 + diag_5 + 1; + int n_first_5 = ((first_seen_5 > 0) ? first_seen_5 >> 7 : 0); + int k_full_5 = k_len_5 >> 7; + int n_kf_5 = ((n_first_5 < k_full_5) ? n_first_5 : k_full_5); + int n_full_5 = ((n_kf_5 < n_blocks_5) ? n_kf_5 : n_blocks_5); + int n_masked_5 = n_blocks_5 - n_full_5; + int c_2 = ((2 * n_blocks_5 >= wmax_1) ? 0 : ((4 * n_blocks_5 >= wmax_1) ? 1 : ((8 * n_blocks_5 >= wmax_1) ? 2 : 3))); + t0_1 = t0_1 + ((c_2 == 0) ? tiles_1 : 0); + t1_1 = t1_1 + ((c_2 == 1) ? tiles_1 : 0); + t2_1 = t2_1 + ((c_2 == 2) ? tiles_1 : 0); + } + unsigned int total_1 = (unsigned int)num_heads_q * head_tiles_1; + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((o_free_addr) & 0xFEFFFFFF) : "memory"); + unsigned int _phase_work_full_1 = 0; + unsigned int _phase_scale_full = 0; + unsigned int _phase_o_full = 0; + unsigned int _phase_stats_full = 0; + #pragma unroll 1 + for (unsigned int it_1 = 0; it_1 < (unsigned int)num_heads_q * ((unsigned int)(total_q + 255 >> 8) + (unsigned int)num_segments) + 1; it_1++) { + mbarrier_wait_cluster_hint(work_full_addr + (wk_1) * 8, _phase_work_full_1, 10000000); + unsigned int ticket_1 = cells[wk_1]; + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + (wk_1) * 8), "r"(0) : "memory"); + wk_1 += 1; + if (wk_1 == 2) { wk_1 = 0; _phase_work_full_1 ^= 1; } + if (ticket_1 >= total_1) { + break; + } + unsigned int hq_1 = (unsigned int)num_heads_q; + unsigned int p1_1 = hq_1 * t0_1; + unsigned int p2_1 = p1_1 + hq_1 * t1_1; + unsigned int p3_1 = p2_1 + hq_1 * t2_1; + int c_3 = ((ticket_1 < p1_1) ? 0 : ((ticket_1 < p2_1) ? 1 : ((ticket_1 < p3_1) ? 2 : 3))); + unsigned int base_1 = ((c_3 == 0) ? 0 : ((c_3 == 1) ? p1_1 : ((c_3 == 2) ? p2_1 : p3_1))); + unsigned int tc_1 = ((c_3 == 0) ? t0_1 : ((c_3 == 1) ? t1_1 : ((c_3 == 2) ? t2_1 : head_tiles_1 - t0_1 - t1_1 - t2_1))); + unsigned int r_1 = ticket_1 - base_1; + unsigned int per_kv_1 = tc_1 * (unsigned int)group; + unsigned int kvh_u_1 = r_1 / per_kv_1; + unsigned int rem_1 = r_1 - kvh_u_1 * per_kv_1; + unsigned int j_1 = rem_1 / (unsigned int)group; + unsigned int head_u_1 = kvh_u_1 * (unsigned int)group + (rem_1 - j_1 * (unsigned int)group); + int seg_1 = 0; + unsigned int before_1 = 0; + unsigned int ends_1 = 0; + #pragma unroll 1 + for (int s_7 = 0; s_7 < num_segments; s_7++) { + int q_len_6 = cu_seqlens_q[s_7 + 1] - cu_seqlens_q[s_7]; + int k_len_6 = cu_seqlens_k[s_7 + 1] - cu_seqlens_k[s_7]; + int row0_6 = ((q_len_6 + 255 >> 8) - 1) * 256; + int diag_6 = k_len_6 - q_len_6; + int n_k_6 = k_len_6 + 127 >> 7; + int n_c_6 = row0_6 + 256 + diag_6 + 127 >> 7; + int n_seen_6 = ((n_c_6 < n_k_6) ? n_c_6 : n_k_6); + int n_blocks_6 = ((n_seen_6 > 0) ? n_seen_6 : 0); + int first_seen_6 = row0_6 + diag_6 + 1; + int n_first_6 = ((first_seen_6 > 0) ? first_seen_6 >> 7 : 0); + int k_full_6 = k_len_6 >> 7; + int n_kf_6 = ((n_first_6 < k_full_6) ? n_first_6 : k_full_6); + int n_full_6 = ((n_kf_6 < n_blocks_6) ? n_kf_6 : n_blocks_6); + int n_masked_6 = n_blocks_6 - n_full_6; + bool in_class_1 = ((2 * n_blocks_6 >= wmax_1) ? 0 : ((4 * n_blocks_6 >= wmax_1) ? 1 : ((8 * n_blocks_6 >= wmax_1) ? 2 : 3))) == c_3; + ends_1 = ends_1 + ((in_class_1) ? (unsigned int)(cu_seqlens_q[s_7 + 1] - cu_seqlens_q[s_7] + 255) >> 8 : 0); + seg_1 = seg_1 + ((ends_1 <= j_1) ? 1 : 0); + before_1 = ((ends_1 <= j_1) ? ends_1 : before_1); + } + unsigned int m_rel_1 = j_1 - before_1; + int q0_1 = cu_seqlens_q[seg_1]; + int q_len_7 = cu_seqlens_q[seg_1 + 1] - q0_1; + int k0_1 = cu_seqlens_k[seg_1]; + int k_len_7 = cu_seqlens_k[seg_1 + 1] - k0_1; + int m_1 = (int)(((unsigned int)(q_len_7 + 255) >> 8) - 1 - m_rel_1); + int head_1 = (int)head_u_1; + int kvh_1 = (int)kvh_u_1; + int row0_7 = m_1 * 256; + int diag_7 = k_len_7 - q_len_7; + int n_k_7 = k_len_7 + 127 >> 7; + int n_c_7 = row0_7 + 256 + diag_7 + 127 >> 7; + int n_seen_7 = ((n_c_7 < n_k_7) ? n_c_7 : n_k_7); + int n_blocks_7 = ((n_seen_7 > 0) ? n_seen_7 : 0); + int first_seen_7 = row0_7 + diag_7 + 1; + int n_first_7 = ((first_seen_7 > 0) ? first_seen_7 >> 7 : 0); + int k_full_7 = k_len_7 >> 7; + int n_kf_7 = ((n_first_7 < k_full_7) ? n_first_7 : k_full_7); + int n_full_7 = ((n_kf_7 < n_blocks_7) ? n_kf_7 : n_blocks_7); + int n_masked_7 = n_blocks_7 - n_full_7; + int row_1 = m_1 * 256 + cta_rank * 128 + tr_1; + bool row_valid = row_1 < q_len_7; + long long o_base = ((long long)(q0_1 + row_1) * (long long)num_heads_q + (long long)head_1) * 256; + long long lse_index = (long long)head_1 * (long long)total_q + (long long)(q0_1 + row_1); + int has_block = ((n_blocks_7 > 0) ? 1 : 0); + #pragma unroll 1 + for (int _ = 0; _ < has_block; _++) { + mbarrier_wait(scale_full_addr, _phase_scale_full); + mbarrier_arrive(scale_empty_addr); + _phase_scale_full ^= 1; + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((p_full_addr + (pq_1) * 8) & 0xFEFFFFFF) : "memory"); + pq_1 += 1; + if (pq_1 == 2) { pq_1 = 0; } + } + #pragma unroll 1 + for (int i_2 = 1; i_2 < n_blocks_7; i_2++) { + mbarrier_wait(scale_full_addr, _phase_scale_full); + float acc_scale_2 = s_scale[(int)0 * 128 + tr_1]; + mbarrier_arrive(scale_empty_addr); + _phase_scale_full ^= 1; + mbarrier_wait(o_full_addr + (oq) * 8, _phase_o_full); + oq += 1; + if (oq == 2) { oq = 0; _phase_o_full ^= 1; } + asm volatile("tcgen05.fence::after_thread_sync;"); + int _warp_vote_0 = __any_sync(0xFFFFFFFF, acc_scale_2 < 1.0f); + if (_warp_vote_0 != 0) { + unsigned int addr_0 = taddr + 256 + (unsigned int)(warp % 4 * 32 << 16); + float ov_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_0[0]), "=f"(ov_0[1]), "=f"(ov_0[2]), "=f"(ov_0[3]), "=f"(ov_0[4]), "=f"(ov_0[5]), "=f"(ov_0[6]), "=f"(ov_0[7]), "=f"(ov_0[8]), "=f"(ov_0[9]), "=f"(ov_0[10]), "=f"(ov_0[11]), "=f"(ov_0[12]), "=f"(ov_0[13]), "=f"(ov_0[14]), "=f"(ov_0[15]), "=f"(ov_0[16]), "=f"(ov_0[17]), "=f"(ov_0[18]), "=f"(ov_0[19]), "=f"(ov_0[20]), "=f"(ov_0[21]), "=f"(ov_0[22]), "=f"(ov_0[23]), "=f"(ov_0[24]), "=f"(ov_0[25]), "=f"(ov_0[26]), "=f"(ov_0[27]), "=f"(ov_0[28]), "=f"(ov_0[29]), "=f"(ov_0[30]), "=f"(ov_0[31]) + : "r"(addr_0)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_1 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_0)[_ls], _scale2_1); + tmem_st_x32_f32(addr_0, ov_0); + unsigned int addr_1 = taddr + 288 + (unsigned int)(warp % 4 * 32 << 16); + float ov_1[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_1[0]), "=f"(ov_1[1]), "=f"(ov_1[2]), "=f"(ov_1[3]), "=f"(ov_1[4]), "=f"(ov_1[5]), "=f"(ov_1[6]), "=f"(ov_1[7]), "=f"(ov_1[8]), "=f"(ov_1[9]), "=f"(ov_1[10]), "=f"(ov_1[11]), "=f"(ov_1[12]), "=f"(ov_1[13]), "=f"(ov_1[14]), "=f"(ov_1[15]), "=f"(ov_1[16]), "=f"(ov_1[17]), "=f"(ov_1[18]), "=f"(ov_1[19]), "=f"(ov_1[20]), "=f"(ov_1[21]), "=f"(ov_1[22]), "=f"(ov_1[23]), "=f"(ov_1[24]), "=f"(ov_1[25]), "=f"(ov_1[26]), "=f"(ov_1[27]), "=f"(ov_1[28]), "=f"(ov_1[29]), "=f"(ov_1[30]), "=f"(ov_1[31]) + : "r"(addr_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_2 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_1)[_ls], _scale2_2); + tmem_st_x32_f32(addr_1, ov_1); + unsigned int addr_2 = taddr + 320 + (unsigned int)(warp % 4 * 32 << 16); + float ov_2[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_2[0]), "=f"(ov_2[1]), "=f"(ov_2[2]), "=f"(ov_2[3]), "=f"(ov_2[4]), "=f"(ov_2[5]), "=f"(ov_2[6]), "=f"(ov_2[7]), "=f"(ov_2[8]), "=f"(ov_2[9]), "=f"(ov_2[10]), "=f"(ov_2[11]), "=f"(ov_2[12]), "=f"(ov_2[13]), "=f"(ov_2[14]), "=f"(ov_2[15]), "=f"(ov_2[16]), "=f"(ov_2[17]), "=f"(ov_2[18]), "=f"(ov_2[19]), "=f"(ov_2[20]), "=f"(ov_2[21]), "=f"(ov_2[22]), "=f"(ov_2[23]), "=f"(ov_2[24]), "=f"(ov_2[25]), "=f"(ov_2[26]), "=f"(ov_2[27]), "=f"(ov_2[28]), "=f"(ov_2[29]), "=f"(ov_2[30]), "=f"(ov_2[31]) + : "r"(addr_2)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_3 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_2)[_ls], _scale2_3); + tmem_st_x32_f32(addr_2, ov_2); + unsigned int addr_3 = taddr + 352 + (unsigned int)(warp % 4 * 32 << 16); + float ov_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_3[0]), "=f"(ov_3[1]), "=f"(ov_3[2]), "=f"(ov_3[3]), "=f"(ov_3[4]), "=f"(ov_3[5]), "=f"(ov_3[6]), "=f"(ov_3[7]), "=f"(ov_3[8]), "=f"(ov_3[9]), "=f"(ov_3[10]), "=f"(ov_3[11]), "=f"(ov_3[12]), "=f"(ov_3[13]), "=f"(ov_3[14]), "=f"(ov_3[15]), "=f"(ov_3[16]), "=f"(ov_3[17]), "=f"(ov_3[18]), "=f"(ov_3[19]), "=f"(ov_3[20]), "=f"(ov_3[21]), "=f"(ov_3[22]), "=f"(ov_3[23]), "=f"(ov_3[24]), "=f"(ov_3[25]), "=f"(ov_3[26]), "=f"(ov_3[27]), "=f"(ov_3[28]), "=f"(ov_3[29]), "=f"(ov_3[30]), "=f"(ov_3[31]) + : "r"(addr_3)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_4 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_3)[_ls], _scale2_4); + tmem_st_x32_f32(addr_3, ov_3); + unsigned int addr_4 = taddr + 384 + (unsigned int)(warp % 4 * 32 << 16); + float ov_4[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_4[0]), "=f"(ov_4[1]), "=f"(ov_4[2]), "=f"(ov_4[3]), "=f"(ov_4[4]), "=f"(ov_4[5]), "=f"(ov_4[6]), "=f"(ov_4[7]), "=f"(ov_4[8]), "=f"(ov_4[9]), "=f"(ov_4[10]), "=f"(ov_4[11]), "=f"(ov_4[12]), "=f"(ov_4[13]), "=f"(ov_4[14]), "=f"(ov_4[15]), "=f"(ov_4[16]), "=f"(ov_4[17]), "=f"(ov_4[18]), "=f"(ov_4[19]), "=f"(ov_4[20]), "=f"(ov_4[21]), "=f"(ov_4[22]), "=f"(ov_4[23]), "=f"(ov_4[24]), "=f"(ov_4[25]), "=f"(ov_4[26]), "=f"(ov_4[27]), "=f"(ov_4[28]), "=f"(ov_4[29]), "=f"(ov_4[30]), "=f"(ov_4[31]) + : "r"(addr_4)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_5 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_4)[_ls], _scale2_5); + tmem_st_x32_f32(addr_4, ov_4); + unsigned int addr_5 = taddr + 416 + (unsigned int)(warp % 4 * 32 << 16); + float ov_5[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_5[0]), "=f"(ov_5[1]), "=f"(ov_5[2]), "=f"(ov_5[3]), "=f"(ov_5[4]), "=f"(ov_5[5]), "=f"(ov_5[6]), "=f"(ov_5[7]), "=f"(ov_5[8]), "=f"(ov_5[9]), "=f"(ov_5[10]), "=f"(ov_5[11]), "=f"(ov_5[12]), "=f"(ov_5[13]), "=f"(ov_5[14]), "=f"(ov_5[15]), "=f"(ov_5[16]), "=f"(ov_5[17]), "=f"(ov_5[18]), "=f"(ov_5[19]), "=f"(ov_5[20]), "=f"(ov_5[21]), "=f"(ov_5[22]), "=f"(ov_5[23]), "=f"(ov_5[24]), "=f"(ov_5[25]), "=f"(ov_5[26]), "=f"(ov_5[27]), "=f"(ov_5[28]), "=f"(ov_5[29]), "=f"(ov_5[30]), "=f"(ov_5[31]) + : "r"(addr_5)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_6 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_5)[_ls], _scale2_6); + tmem_st_x32_f32(addr_5, ov_5); + unsigned int addr_6 = taddr + 448 + (unsigned int)(warp % 4 * 32 << 16); + float ov_6[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_6[0]), "=f"(ov_6[1]), "=f"(ov_6[2]), "=f"(ov_6[3]), "=f"(ov_6[4]), "=f"(ov_6[5]), "=f"(ov_6[6]), "=f"(ov_6[7]), "=f"(ov_6[8]), "=f"(ov_6[9]), "=f"(ov_6[10]), "=f"(ov_6[11]), "=f"(ov_6[12]), "=f"(ov_6[13]), "=f"(ov_6[14]), "=f"(ov_6[15]), "=f"(ov_6[16]), "=f"(ov_6[17]), "=f"(ov_6[18]), "=f"(ov_6[19]), "=f"(ov_6[20]), "=f"(ov_6[21]), "=f"(ov_6[22]), "=f"(ov_6[23]), "=f"(ov_6[24]), "=f"(ov_6[25]), "=f"(ov_6[26]), "=f"(ov_6[27]), "=f"(ov_6[28]), "=f"(ov_6[29]), "=f"(ov_6[30]), "=f"(ov_6[31]) + : "r"(addr_6)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_7 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_6)[_ls], _scale2_7); + tmem_st_x32_f32(addr_6, ov_6); + unsigned int addr_7 = taddr + 480 + (unsigned int)(warp % 4 * 32 << 16); + float ov_7[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_7[0]), "=f"(ov_7[1]), "=f"(ov_7[2]), "=f"(ov_7[3]), "=f"(ov_7[4]), "=f"(ov_7[5]), "=f"(ov_7[6]), "=f"(ov_7[7]), "=f"(ov_7[8]), "=f"(ov_7[9]), "=f"(ov_7[10]), "=f"(ov_7[11]), "=f"(ov_7[12]), "=f"(ov_7[13]), "=f"(ov_7[14]), "=f"(ov_7[15]), "=f"(ov_7[16]), "=f"(ov_7[17]), "=f"(ov_7[18]), "=f"(ov_7[19]), "=f"(ov_7[20]), "=f"(ov_7[21]), "=f"(ov_7[22]), "=f"(ov_7[23]), "=f"(ov_7[24]), "=f"(ov_7[25]), "=f"(ov_7[26]), "=f"(ov_7[27]), "=f"(ov_7[28]), "=f"(ov_7[29]), "=f"(ov_7[30]), "=f"(ov_7[31]) + : "r"(addr_7)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_8 = {acc_scale_2, acc_scale_2}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_7)[_ls], _scale2_8); + tmem_st_x32_f32(addr_7, ov_7); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + } + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((p_full_addr + (pq_1) * 8) & 0xFEFFFFFF) : "memory"); + pq_1 += 1; + if (pq_1 == 2) { pq_1 = 0; } + } + mbarrier_wait(stats_full_addr, _phase_stats_full); + float final_sum = s_stats[tr_1]; + float final_max = s_stats[128 + tr_1]; + mbarrier_arrive(stats_empty_addr); + _phase_stats_full ^= 1; + mbarrier_wait(o_full_addr + (oq) * 8, _phase_o_full); + oq += 1; + if (oq == 2) { oq = 0; _phase_o_full ^= 1; } + asm volatile("tcgen05.fence::after_thread_sync;"); + asm volatile( + "mbarrier.arrive.release.cta.shared::cluster.b64 _, [%0];" + :: "r"((o_free_addr) & 0xFEFFFFFF) : "memory"); + float final_scale = 0.0f; + float lse = -CUDART_INF_F; + if (final_sum > 0.0f) { + float _rcp_9 = approx_rcp(final_sum); + final_scale = _rcp_9; + float _log2_10; + asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(_log2_10) : "f"(final_sum)); + lse = (final_max * softmax_scale_log2 + _log2_10) * 0.6931472f; + } + if (n_blocks_7 > 0) { + unsigned int addr_8 = taddr + 256 + (unsigned int)(warp % 4 * 32 << 16); + float ov_8[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_8[0]), "=f"(ov_8[1]), "=f"(ov_8[2]), "=f"(ov_8[3]), "=f"(ov_8[4]), "=f"(ov_8[5]), "=f"(ov_8[6]), "=f"(ov_8[7]), "=f"(ov_8[8]), "=f"(ov_8[9]), "=f"(ov_8[10]), "=f"(ov_8[11]), "=f"(ov_8[12]), "=f"(ov_8[13]), "=f"(ov_8[14]), "=f"(ov_8[15]), "=f"(ov_8[16]), "=f"(ov_8[17]), "=f"(ov_8[18]), "=f"(ov_8[19]), "=f"(ov_8[20]), "=f"(ov_8[21]), "=f"(ov_8[22]), "=f"(ov_8[23]), "=f"(ov_8[24]), "=f"(ov_8[25]), "=f"(ov_8[26]), "=f"(ov_8[27]), "=f"(ov_8[28]), "=f"(ov_8[29]), "=f"(ov_8[30]), "=f"(ov_8[31]) + : "r"(addr_8)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_11 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_8)[_ls], _scale2_11); + unsigned int pk_8[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_8[_lp*2 + 0], ov_8[_lp*2+1 + 0])); + pk_8[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 0)), "r"(pk_8[(0) + 0]), "r"(pk_8[(0) + 1]), "r"(pk_8[(0) + 2]), "r"(pk_8[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_9[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_8[_lp*2 + 8], ov_8[_lp*2+1 + 8])); + pk_9[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 8)), "r"(pk_9[(0) + 0]), "r"(pk_9[(0) + 1]), "r"(pk_9[(0) + 2]), "r"(pk_9[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_10[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_8[_lp*2 + 16], ov_8[_lp*2+1 + 16])); + pk_10[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 16)), "r"(pk_10[(0) + 0]), "r"(pk_10[(0) + 1]), "r"(pk_10[(0) + 2]), "r"(pk_10[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_11[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_8[_lp*2 + 24], ov_8[_lp*2+1 + 24])); + pk_11[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 24)), "r"(pk_11[(0) + 0]), "r"(pk_11[(0) + 1]), "r"(pk_11[(0) + 2]), "r"(pk_11[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_9 = taddr + 288 + (unsigned int)(warp % 4 * 32 << 16); + float ov_9[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_9[0]), "=f"(ov_9[1]), "=f"(ov_9[2]), "=f"(ov_9[3]), "=f"(ov_9[4]), "=f"(ov_9[5]), "=f"(ov_9[6]), "=f"(ov_9[7]), "=f"(ov_9[8]), "=f"(ov_9[9]), "=f"(ov_9[10]), "=f"(ov_9[11]), "=f"(ov_9[12]), "=f"(ov_9[13]), "=f"(ov_9[14]), "=f"(ov_9[15]), "=f"(ov_9[16]), "=f"(ov_9[17]), "=f"(ov_9[18]), "=f"(ov_9[19]), "=f"(ov_9[20]), "=f"(ov_9[21]), "=f"(ov_9[22]), "=f"(ov_9[23]), "=f"(ov_9[24]), "=f"(ov_9[25]), "=f"(ov_9[26]), "=f"(ov_9[27]), "=f"(ov_9[28]), "=f"(ov_9[29]), "=f"(ov_9[30]), "=f"(ov_9[31]) + : "r"(addr_9)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_12 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_9)[_ls], _scale2_12); + unsigned int pk_12[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_9[_lp*2 + 0], ov_9[_lp*2+1 + 0])); + pk_12[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 32)), "r"(pk_12[(0) + 0]), "r"(pk_12[(0) + 1]), "r"(pk_12[(0) + 2]), "r"(pk_12[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_13[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_9[_lp*2 + 8], ov_9[_lp*2+1 + 8])); + pk_13[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 40)), "r"(pk_13[(0) + 0]), "r"(pk_13[(0) + 1]), "r"(pk_13[(0) + 2]), "r"(pk_13[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_14[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_9[_lp*2 + 16], ov_9[_lp*2+1 + 16])); + pk_14[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 48)), "r"(pk_14[(0) + 0]), "r"(pk_14[(0) + 1]), "r"(pk_14[(0) + 2]), "r"(pk_14[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_15[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_9[_lp*2 + 24], ov_9[_lp*2+1 + 24])); + pk_15[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 56)), "r"(pk_15[(0) + 0]), "r"(pk_15[(0) + 1]), "r"(pk_15[(0) + 2]), "r"(pk_15[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_10 = taddr + 320 + (unsigned int)(warp % 4 * 32 << 16); + float ov_10[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_10[0]), "=f"(ov_10[1]), "=f"(ov_10[2]), "=f"(ov_10[3]), "=f"(ov_10[4]), "=f"(ov_10[5]), "=f"(ov_10[6]), "=f"(ov_10[7]), "=f"(ov_10[8]), "=f"(ov_10[9]), "=f"(ov_10[10]), "=f"(ov_10[11]), "=f"(ov_10[12]), "=f"(ov_10[13]), "=f"(ov_10[14]), "=f"(ov_10[15]), "=f"(ov_10[16]), "=f"(ov_10[17]), "=f"(ov_10[18]), "=f"(ov_10[19]), "=f"(ov_10[20]), "=f"(ov_10[21]), "=f"(ov_10[22]), "=f"(ov_10[23]), "=f"(ov_10[24]), "=f"(ov_10[25]), "=f"(ov_10[26]), "=f"(ov_10[27]), "=f"(ov_10[28]), "=f"(ov_10[29]), "=f"(ov_10[30]), "=f"(ov_10[31]) + : "r"(addr_10)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_13 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_10)[_ls], _scale2_13); + unsigned int pk_16[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_10[_lp*2 + 0], ov_10[_lp*2+1 + 0])); + pk_16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 64)), "r"(pk_16[(0) + 0]), "r"(pk_16[(0) + 1]), "r"(pk_16[(0) + 2]), "r"(pk_16[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_17[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_10[_lp*2 + 8], ov_10[_lp*2+1 + 8])); + pk_17[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 72)), "r"(pk_17[(0) + 0]), "r"(pk_17[(0) + 1]), "r"(pk_17[(0) + 2]), "r"(pk_17[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_18[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_10[_lp*2 + 16], ov_10[_lp*2+1 + 16])); + pk_18[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 80)), "r"(pk_18[(0) + 0]), "r"(pk_18[(0) + 1]), "r"(pk_18[(0) + 2]), "r"(pk_18[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_19[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_10[_lp*2 + 24], ov_10[_lp*2+1 + 24])); + pk_19[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 88)), "r"(pk_19[(0) + 0]), "r"(pk_19[(0) + 1]), "r"(pk_19[(0) + 2]), "r"(pk_19[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_11 = taddr + 352 + (unsigned int)(warp % 4 * 32 << 16); + float ov_11[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_11[0]), "=f"(ov_11[1]), "=f"(ov_11[2]), "=f"(ov_11[3]), "=f"(ov_11[4]), "=f"(ov_11[5]), "=f"(ov_11[6]), "=f"(ov_11[7]), "=f"(ov_11[8]), "=f"(ov_11[9]), "=f"(ov_11[10]), "=f"(ov_11[11]), "=f"(ov_11[12]), "=f"(ov_11[13]), "=f"(ov_11[14]), "=f"(ov_11[15]), "=f"(ov_11[16]), "=f"(ov_11[17]), "=f"(ov_11[18]), "=f"(ov_11[19]), "=f"(ov_11[20]), "=f"(ov_11[21]), "=f"(ov_11[22]), "=f"(ov_11[23]), "=f"(ov_11[24]), "=f"(ov_11[25]), "=f"(ov_11[26]), "=f"(ov_11[27]), "=f"(ov_11[28]), "=f"(ov_11[29]), "=f"(ov_11[30]), "=f"(ov_11[31]) + : "r"(addr_11)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_14 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_11)[_ls], _scale2_14); + unsigned int pk_20[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_11[_lp*2 + 0], ov_11[_lp*2+1 + 0])); + pk_20[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 96)), "r"(pk_20[(0) + 0]), "r"(pk_20[(0) + 1]), "r"(pk_20[(0) + 2]), "r"(pk_20[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_21[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_11[_lp*2 + 8], ov_11[_lp*2+1 + 8])); + pk_21[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 104)), "r"(pk_21[(0) + 0]), "r"(pk_21[(0) + 1]), "r"(pk_21[(0) + 2]), "r"(pk_21[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_22[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_11[_lp*2 + 16], ov_11[_lp*2+1 + 16])); + pk_22[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 112)), "r"(pk_22[(0) + 0]), "r"(pk_22[(0) + 1]), "r"(pk_22[(0) + 2]), "r"(pk_22[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_23[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_11[_lp*2 + 24], ov_11[_lp*2+1 + 24])); + pk_23[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 120)), "r"(pk_23[(0) + 0]), "r"(pk_23[(0) + 1]), "r"(pk_23[(0) + 2]), "r"(pk_23[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_12 = taddr + 384 + (unsigned int)(warp % 4 * 32 << 16); + float ov_12[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_12[0]), "=f"(ov_12[1]), "=f"(ov_12[2]), "=f"(ov_12[3]), "=f"(ov_12[4]), "=f"(ov_12[5]), "=f"(ov_12[6]), "=f"(ov_12[7]), "=f"(ov_12[8]), "=f"(ov_12[9]), "=f"(ov_12[10]), "=f"(ov_12[11]), "=f"(ov_12[12]), "=f"(ov_12[13]), "=f"(ov_12[14]), "=f"(ov_12[15]), "=f"(ov_12[16]), "=f"(ov_12[17]), "=f"(ov_12[18]), "=f"(ov_12[19]), "=f"(ov_12[20]), "=f"(ov_12[21]), "=f"(ov_12[22]), "=f"(ov_12[23]), "=f"(ov_12[24]), "=f"(ov_12[25]), "=f"(ov_12[26]), "=f"(ov_12[27]), "=f"(ov_12[28]), "=f"(ov_12[29]), "=f"(ov_12[30]), "=f"(ov_12[31]) + : "r"(addr_12)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_15 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_12)[_ls], _scale2_15); + unsigned int pk_24[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_12[_lp*2 + 0], ov_12[_lp*2+1 + 0])); + pk_24[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 128)), "r"(pk_24[(0) + 0]), "r"(pk_24[(0) + 1]), "r"(pk_24[(0) + 2]), "r"(pk_24[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_25[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_12[_lp*2 + 8], ov_12[_lp*2+1 + 8])); + pk_25[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 136)), "r"(pk_25[(0) + 0]), "r"(pk_25[(0) + 1]), "r"(pk_25[(0) + 2]), "r"(pk_25[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_26[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_12[_lp*2 + 16], ov_12[_lp*2+1 + 16])); + pk_26[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 144)), "r"(pk_26[(0) + 0]), "r"(pk_26[(0) + 1]), "r"(pk_26[(0) + 2]), "r"(pk_26[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_27[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_12[_lp*2 + 24], ov_12[_lp*2+1 + 24])); + pk_27[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 152)), "r"(pk_27[(0) + 0]), "r"(pk_27[(0) + 1]), "r"(pk_27[(0) + 2]), "r"(pk_27[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_13 = taddr + 416 + (unsigned int)(warp % 4 * 32 << 16); + float ov_13[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_13[0]), "=f"(ov_13[1]), "=f"(ov_13[2]), "=f"(ov_13[3]), "=f"(ov_13[4]), "=f"(ov_13[5]), "=f"(ov_13[6]), "=f"(ov_13[7]), "=f"(ov_13[8]), "=f"(ov_13[9]), "=f"(ov_13[10]), "=f"(ov_13[11]), "=f"(ov_13[12]), "=f"(ov_13[13]), "=f"(ov_13[14]), "=f"(ov_13[15]), "=f"(ov_13[16]), "=f"(ov_13[17]), "=f"(ov_13[18]), "=f"(ov_13[19]), "=f"(ov_13[20]), "=f"(ov_13[21]), "=f"(ov_13[22]), "=f"(ov_13[23]), "=f"(ov_13[24]), "=f"(ov_13[25]), "=f"(ov_13[26]), "=f"(ov_13[27]), "=f"(ov_13[28]), "=f"(ov_13[29]), "=f"(ov_13[30]), "=f"(ov_13[31]) + : "r"(addr_13)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_16 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_13)[_ls], _scale2_16); + unsigned int pk_28[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_13[_lp*2 + 0], ov_13[_lp*2+1 + 0])); + pk_28[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 160)), "r"(pk_28[(0) + 0]), "r"(pk_28[(0) + 1]), "r"(pk_28[(0) + 2]), "r"(pk_28[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_29[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_13[_lp*2 + 8], ov_13[_lp*2+1 + 8])); + pk_29[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 168)), "r"(pk_29[(0) + 0]), "r"(pk_29[(0) + 1]), "r"(pk_29[(0) + 2]), "r"(pk_29[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_30[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_13[_lp*2 + 16], ov_13[_lp*2+1 + 16])); + pk_30[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 176)), "r"(pk_30[(0) + 0]), "r"(pk_30[(0) + 1]), "r"(pk_30[(0) + 2]), "r"(pk_30[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_31[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_13[_lp*2 + 24], ov_13[_lp*2+1 + 24])); + pk_31[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 184)), "r"(pk_31[(0) + 0]), "r"(pk_31[(0) + 1]), "r"(pk_31[(0) + 2]), "r"(pk_31[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_14 = taddr + 448 + (unsigned int)(warp % 4 * 32 << 16); + float ov_14[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_14[0]), "=f"(ov_14[1]), "=f"(ov_14[2]), "=f"(ov_14[3]), "=f"(ov_14[4]), "=f"(ov_14[5]), "=f"(ov_14[6]), "=f"(ov_14[7]), "=f"(ov_14[8]), "=f"(ov_14[9]), "=f"(ov_14[10]), "=f"(ov_14[11]), "=f"(ov_14[12]), "=f"(ov_14[13]), "=f"(ov_14[14]), "=f"(ov_14[15]), "=f"(ov_14[16]), "=f"(ov_14[17]), "=f"(ov_14[18]), "=f"(ov_14[19]), "=f"(ov_14[20]), "=f"(ov_14[21]), "=f"(ov_14[22]), "=f"(ov_14[23]), "=f"(ov_14[24]), "=f"(ov_14[25]), "=f"(ov_14[26]), "=f"(ov_14[27]), "=f"(ov_14[28]), "=f"(ov_14[29]), "=f"(ov_14[30]), "=f"(ov_14[31]) + : "r"(addr_14)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_17 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_14)[_ls], _scale2_17); + unsigned int pk_32[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_14[_lp*2 + 0], ov_14[_lp*2+1 + 0])); + pk_32[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 192)), "r"(pk_32[(0) + 0]), "r"(pk_32[(0) + 1]), "r"(pk_32[(0) + 2]), "r"(pk_32[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_33[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_14[_lp*2 + 8], ov_14[_lp*2+1 + 8])); + pk_33[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 200)), "r"(pk_33[(0) + 0]), "r"(pk_33[(0) + 1]), "r"(pk_33[(0) + 2]), "r"(pk_33[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_34[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_14[_lp*2 + 16], ov_14[_lp*2+1 + 16])); + pk_34[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 208)), "r"(pk_34[(0) + 0]), "r"(pk_34[(0) + 1]), "r"(pk_34[(0) + 2]), "r"(pk_34[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_35[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_14[_lp*2 + 24], ov_14[_lp*2+1 + 24])); + pk_35[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 216)), "r"(pk_35[(0) + 0]), "r"(pk_35[(0) + 1]), "r"(pk_35[(0) + 2]), "r"(pk_35[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int addr_15 = taddr + 480 + (unsigned int)(warp % 4 * 32 << 16); + float ov_15[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(ov_15[0]), "=f"(ov_15[1]), "=f"(ov_15[2]), "=f"(ov_15[3]), "=f"(ov_15[4]), "=f"(ov_15[5]), "=f"(ov_15[6]), "=f"(ov_15[7]), "=f"(ov_15[8]), "=f"(ov_15[9]), "=f"(ov_15[10]), "=f"(ov_15[11]), "=f"(ov_15[12]), "=f"(ov_15[13]), "=f"(ov_15[14]), "=f"(ov_15[15]), "=f"(ov_15[16]), "=f"(ov_15[17]), "=f"(ov_15[18]), "=f"(ov_15[19]), "=f"(ov_15[20]), "=f"(ov_15[21]), "=f"(ov_15[22]), "=f"(ov_15[23]), "=f"(ov_15[24]), "=f"(ov_15[25]), "=f"(ov_15[26]), "=f"(ov_15[27]), "=f"(ov_15[28]), "=f"(ov_15[29]), "=f"(ov_15[30]), "=f"(ov_15[31]) + : "r"(addr_15)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + const float2 _scale2_18 = {final_scale, final_scale}; + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(ov_15)[_ls], _scale2_18); + unsigned int pk_36[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_15[_lp*2 + 0], ov_15[_lp*2+1 + 0])); + pk_36[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 224)), "r"(pk_36[(0) + 0]), "r"(pk_36[(0) + 1]), "r"(pk_36[(0) + 2]), "r"(pk_36[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_37[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_15[_lp*2 + 8], ov_15[_lp*2+1 + 8])); + pk_37[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 232)), "r"(pk_37[(0) + 0]), "r"(pk_37[(0) + 1]), "r"(pk_37[(0) + 2]), "r"(pk_37[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_38[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_15[_lp*2 + 16], ov_15[_lp*2+1 + 16])); + pk_38[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 240)), "r"(pk_38[(0) + 0]), "r"(pk_38[(0) + 1]), "r"(pk_38[(0) + 2]), "r"(pk_38[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + unsigned int pk_39[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ov_15[_lp*2 + 24], ov_15[_lp*2+1 + 24])); + pk_39[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 248)), "r"(pk_39[(0) + 0]), "r"(pk_39[(0) + 1]), "r"(pk_39[(0) + 2]), "r"(pk_39[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + } else { + unsigned int zero[4]; + zero[0] = 0; + zero[1] = 0; + zero[2] = 0; + zero[3] = 0; + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 0)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 8)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 16)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 24)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 32)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 40)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 48)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 56)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 64)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 72)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 80)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 88)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 96)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 104)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 112)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 120)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 128)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 136)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 144)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 152)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 160)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 168)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 176)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 184)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 192)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 200)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 208)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 216)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 224)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 232)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 240)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %5, 0; @p_ st.global.v4.b32 [%0], {%1, %2, %3, %4}; }" :: "l"(O + (o_base + 248)), "r"(zero[(0) + 0]), "r"(zero[(0) + 1]), "r"(zero[(0) + 2]), "r"(zero[(0) + 3]), "r"((unsigned int)(row_valid)) : "memory"); + } + asm volatile("{ .reg .pred p_; setp.ne.b32 p_, %2, 0; @p_ st.global.f32 [%0], %1; }" :: "l"(reinterpret_cast(LSE + lse_index)), "f"(lse), "r"((unsigned int)(row_valid)) : "memory"); + } + } + } + // ---- Role: mma ---- + if (warp == 8) { + { // mma_main + unsigned int wk_2 = 0; + unsigned int kv = 0; + unsigned int sq_1 = 0; + unsigned int pq_2 = 0; + unsigned int oq_1 = 0; + unsigned int head_tiles_2 = 0; + #pragma unroll 1 + for (int s_8 = 0; s_8 < num_segments; s_8++) { + head_tiles_2 = head_tiles_2 + ((unsigned int)(cu_seqlens_q[s_8 + 1] - cu_seqlens_q[s_8] + 255) >> 8); + } + int wmax_2 = 0; + #pragma unroll 1 + for (int s_9 = 0; s_9 < num_segments; s_9++) { + int q_len_8 = cu_seqlens_q[s_9 + 1] - cu_seqlens_q[s_9]; + int k_len_8 = cu_seqlens_k[s_9 + 1] - cu_seqlens_k[s_9]; + int row0_8 = ((q_len_8 + 255 >> 8) - 1) * 256; + int diag_8 = k_len_8 - q_len_8; + int n_k_8 = k_len_8 + 127 >> 7; + int n_c_8 = row0_8 + 256 + diag_8 + 127 >> 7; + int n_seen_8 = ((n_c_8 < n_k_8) ? n_c_8 : n_k_8); + int n_blocks_8 = ((n_seen_8 > 0) ? n_seen_8 : 0); + int first_seen_8 = row0_8 + diag_8 + 1; + int n_first_8 = ((first_seen_8 > 0) ? first_seen_8 >> 7 : 0); + int k_full_8 = k_len_8 >> 7; + int n_kf_8 = ((n_first_8 < k_full_8) ? n_first_8 : k_full_8); + int n_full_8 = ((n_kf_8 < n_blocks_8) ? n_kf_8 : n_blocks_8); + int n_masked_8 = n_blocks_8 - n_full_8; + int w_2 = n_blocks_8; + wmax_2 = ((w_2 > wmax_2) ? w_2 : wmax_2); + } + unsigned int t0_2 = 0; + unsigned int t1_2 = 0; + unsigned int t2_2 = 0; + #pragma unroll 1 + for (int s_10 = 0; s_10 < num_segments; s_10++) { + unsigned int tiles_2 = (unsigned int)(cu_seqlens_q[s_10 + 1] - cu_seqlens_q[s_10] + 255) >> 8; + int q_len_9 = cu_seqlens_q[s_10 + 1] - cu_seqlens_q[s_10]; + int k_len_9 = cu_seqlens_k[s_10 + 1] - cu_seqlens_k[s_10]; + int row0_9 = ((q_len_9 + 255 >> 8) - 1) * 256; + int diag_9 = k_len_9 - q_len_9; + int n_k_9 = k_len_9 + 127 >> 7; + int n_c_9 = row0_9 + 256 + diag_9 + 127 >> 7; + int n_seen_9 = ((n_c_9 < n_k_9) ? n_c_9 : n_k_9); + int n_blocks_9 = ((n_seen_9 > 0) ? n_seen_9 : 0); + int first_seen_9 = row0_9 + diag_9 + 1; + int n_first_9 = ((first_seen_9 > 0) ? first_seen_9 >> 7 : 0); + int k_full_9 = k_len_9 >> 7; + int n_kf_9 = ((n_first_9 < k_full_9) ? n_first_9 : k_full_9); + int n_full_9 = ((n_kf_9 < n_blocks_9) ? n_kf_9 : n_blocks_9); + int n_masked_9 = n_blocks_9 - n_full_9; + int c_4 = ((2 * n_blocks_9 >= wmax_2) ? 0 : ((4 * n_blocks_9 >= wmax_2) ? 1 : ((8 * n_blocks_9 >= wmax_2) ? 2 : 3))); + t0_2 = t0_2 + ((c_4 == 0) ? tiles_2 : 0); + t1_2 = t1_2 + ((c_4 == 1) ? tiles_2 : 0); + t2_2 = t2_2 + ((c_4 == 2) ? tiles_2 : 0); + } + unsigned int total_2 = (unsigned int)num_heads_q * head_tiles_2; + unsigned int _phase_work_full_2 = 0; + unsigned int _phase_q_full = 0; + unsigned int _phase_kv_full = 0; + unsigned int _phase_p_full = 0; + unsigned int _phase_o_free = 0; + #pragma unroll 1 + for (unsigned int it_2 = 0; it_2 < (unsigned int)num_heads_q * ((unsigned int)(total_q + 255 >> 8) + (unsigned int)num_segments) + 1; it_2++) { + mbarrier_wait_cluster_hint(work_full_addr + (wk_2) * 8, _phase_work_full_2, 10000000); + unsigned int ticket_2 = cells[wk_2]; + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + (wk_2) * 8), "r"(0) : "memory"); + wk_2 += 1; + if (wk_2 == 2) { wk_2 = 0; _phase_work_full_2 ^= 1; } + if (ticket_2 >= total_2) { + break; + } + unsigned int hq_2 = (unsigned int)num_heads_q; + unsigned int p1_2 = hq_2 * t0_2; + unsigned int p2_2 = p1_2 + hq_2 * t1_2; + unsigned int p3_2 = p2_2 + hq_2 * t2_2; + int c_5 = ((ticket_2 < p1_2) ? 0 : ((ticket_2 < p2_2) ? 1 : ((ticket_2 < p3_2) ? 2 : 3))); + unsigned int base_2 = ((c_5 == 0) ? 0 : ((c_5 == 1) ? p1_2 : ((c_5 == 2) ? p2_2 : p3_2))); + unsigned int tc_2 = ((c_5 == 0) ? t0_2 : ((c_5 == 1) ? t1_2 : ((c_5 == 2) ? t2_2 : head_tiles_2 - t0_2 - t1_2 - t2_2))); + unsigned int r_2 = ticket_2 - base_2; + unsigned int per_kv_2 = tc_2 * (unsigned int)group; + unsigned int kvh_u_2 = r_2 / per_kv_2; + unsigned int rem_2 = r_2 - kvh_u_2 * per_kv_2; + unsigned int j_2 = rem_2 / (unsigned int)group; + unsigned int head_u_2 = kvh_u_2 * (unsigned int)group + (rem_2 - j_2 * (unsigned int)group); + int seg_2 = 0; + unsigned int before_2 = 0; + unsigned int ends_2 = 0; + #pragma unroll 1 + for (int s_11 = 0; s_11 < num_segments; s_11++) { + int q_len_10 = cu_seqlens_q[s_11 + 1] - cu_seqlens_q[s_11]; + int k_len_10 = cu_seqlens_k[s_11 + 1] - cu_seqlens_k[s_11]; + int row0_10 = ((q_len_10 + 255 >> 8) - 1) * 256; + int diag_10 = k_len_10 - q_len_10; + int n_k_10 = k_len_10 + 127 >> 7; + int n_c_10 = row0_10 + 256 + diag_10 + 127 >> 7; + int n_seen_10 = ((n_c_10 < n_k_10) ? n_c_10 : n_k_10); + int n_blocks_10 = ((n_seen_10 > 0) ? n_seen_10 : 0); + int first_seen_10 = row0_10 + diag_10 + 1; + int n_first_10 = ((first_seen_10 > 0) ? first_seen_10 >> 7 : 0); + int k_full_10 = k_len_10 >> 7; + int n_kf_10 = ((n_first_10 < k_full_10) ? n_first_10 : k_full_10); + int n_full_10 = ((n_kf_10 < n_blocks_10) ? n_kf_10 : n_blocks_10); + int n_masked_10 = n_blocks_10 - n_full_10; + bool in_class_2 = ((2 * n_blocks_10 >= wmax_2) ? 0 : ((4 * n_blocks_10 >= wmax_2) ? 1 : ((8 * n_blocks_10 >= wmax_2) ? 2 : 3))) == c_5; + ends_2 = ends_2 + ((in_class_2) ? (unsigned int)(cu_seqlens_q[s_11 + 1] - cu_seqlens_q[s_11] + 255) >> 8 : 0); + seg_2 = seg_2 + ((ends_2 <= j_2) ? 1 : 0); + before_2 = ((ends_2 <= j_2) ? ends_2 : before_2); + } + unsigned int m_rel_2 = j_2 - before_2; + int q0_2 = cu_seqlens_q[seg_2]; + int q_len_11 = cu_seqlens_q[seg_2 + 1] - q0_2; + int k0_2 = cu_seqlens_k[seg_2]; + int k_len_11 = cu_seqlens_k[seg_2 + 1] - k0_2; + int m_2 = (int)(((unsigned int)(q_len_11 + 255) >> 8) - 1 - m_rel_2); + int head_2 = (int)head_u_2; + int kvh_2 = (int)kvh_u_2; + int row0_11 = m_2 * 256; + int diag_11 = k_len_11 - q_len_11; + int n_k_11 = k_len_11 + 127 >> 7; + int n_c_11 = row0_11 + 256 + diag_11 + 127 >> 7; + int n_seen_11 = ((n_c_11 < n_k_11) ? n_c_11 : n_k_11); + int n_blocks_11 = ((n_seen_11 > 0) ? n_seen_11 : 0); + int first_seen_11 = row0_11 + diag_11 + 1; + int n_first_11 = ((first_seen_11 > 0) ? first_seen_11 >> 7 : 0); + int k_full_11 = k_len_11 >> 7; + int n_kf_11 = ((n_first_11 < k_full_11) ? n_first_11 : k_full_11); + int n_full_11 = ((n_kf_11 < n_blocks_11) ? n_kf_11 : n_blocks_11); + int n_masked_11 = n_blocks_11 - n_full_11; + if (cta_rank == 0) { + int has_block_1 = ((n_blocks_11 > 0) ? 1 : 0); + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + #pragma unroll 1 + for (int __1 = 0; __1 < has_block_1; __1++) { + mbarrier_wait(kv_full_addr + (kv) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 + int a_hi = ((smem_q_hi_addr) >> 4) & 0x3FFF; + int a_lo = ((smem_q_lo_addr) >> 4) & 0x3FFF; + int b_hi = (((smem_k_hi_addr) >> 4) & 0x3FFF) + (kv) * 1024; + int b_lo = (((smem_k_lo_addr) >> 4) & 0x3FFF) + (kv) * 1024; +#else + int a_hi = ((smem_q_hi_addr) >> 4) & 0x7FFF; + int a_lo = ((smem_q_lo_addr) >> 4) & 0x7FFF; + int b_hi = (((smem_k_hi_addr) >> 4) & 0x7FFF) + (kv) * 1024; + int b_lo = (((smem_k_lo_addr) >> 4) & 0x7FFF) + (kv) * 1024; +#endif + uint32_t _elected_0 = elect_sync(); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "setp.ne.b32 leader, %4, 0;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 270533776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 1018;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(a_hi), "r"(b_hi), "r"((tmem_s + (sq_1 * 128))), "r"(0), "r"(_elected_0) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "setp.ne.b32 leader, %4, 0;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 270533776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 1018;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(a_lo), "r"(b_lo), "r"((tmem_s + (sq_1 * 128))), "r"(1), "r"(_elected_0) : "memory"); + tcgen05_commit_cg2_multicast_leader(s_full_addr + (sq_1) * 8, (uint16_t)(3), _elected_0); + tcgen05_commit_cg2_multicast_leader(kv_empty_addr + (kv) * 8, (uint16_t)(3), _elected_0); + kv += 1; + if (kv == 5) { kv = 0; _phase_kv_full ^= 1; } + sq_1 += 1; + if (sq_1 == 2) { sq_1 = 0; } + } + #pragma unroll 1 + for (int i_3 = 1; i_3 < n_blocks_11; i_3++) { + mbarrier_wait(kv_full_addr + (kv) * 8, _phase_kv_full); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 + int a_hi_1 = ((smem_q_hi_addr) >> 4) & 0x3FFF; + int a_lo_1 = ((smem_q_lo_addr) >> 4) & 0x3FFF; + int b_hi_1 = (((smem_k_hi_addr) >> 4) & 0x3FFF) + (kv) * 1024; + int b_lo_1 = (((smem_k_lo_addr) >> 4) & 0x3FFF) + (kv) * 1024; +#else + int a_hi_1 = ((smem_q_hi_addr) >> 4) & 0x7FFF; + int a_lo_1 = ((smem_q_lo_addr) >> 4) & 0x7FFF; + int b_hi_1 = (((smem_k_hi_addr) >> 4) & 0x7FFF) + (kv) * 1024; + int b_lo_1 = (((smem_k_lo_addr) >> 4) & 0x7FFF) + (kv) * 1024; +#endif + uint32_t _elected_1 = elect_sync(); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "setp.ne.b32 leader, %4, 0;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 270533776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 1018;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(a_hi_1), "r"(b_hi_1), "r"((tmem_s + (sq_1 * 128))), "r"(0), "r"(_elected_1) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "setp.ne.b32 leader, %4, 0;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 270533776;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 1018;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "@leader tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(a_lo_1), "r"(b_lo_1), "r"((tmem_s + (sq_1 * 128))), "r"(1), "r"(_elected_1) : "memory"); + tcgen05_commit_cg2_multicast_leader(s_full_addr + (sq_1) * 8, (uint16_t)(3), _elected_1); + tcgen05_commit_cg2_multicast_leader(kv_empty_addr + (kv) * 8, (uint16_t)(3), _elected_1); + kv += 1; + if (kv == 5) { kv = 0; _phase_kv_full ^= 1; } + sq_1 += 1; + if (sq_1 == 2) { sq_1 = 0; } + int first = ((i_3 == 1) ? 1 : 0); + mbarrier_wait(kv_full_addr + (kv) * 8, _phase_kv_full); + mbarrier_wait(p_full_addr + (pq_2) * 8, _phase_p_full); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 + int b_hi_0 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv) * 1024; + int b_lo_0 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv) * 1024; +#else + int b_hi_0 = ((((smem_v_hi_addr) >> 4) & 0x7FFF) | 0x4000000) + (kv) * 1024; + int b_lo_0 = ((((smem_v_lo_addr) >> 4) & 0x7FFF) | 0x4000000) + (kv) * 1024; +#endif + if (elect_sync()) { + if (pq_2 == 0) { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_0), "r"(tmem_p0), "r"(((first) ? 0 : 1)) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_0), "r"(tmem_p0), "r"(((first) ? 0 : 1)) : "memory"); + } else { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_0), "r"(tmem_p1), "r"(((first) ? 0 : 1)) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_0), "r"(tmem_p1), "r"(((first) ? 0 : 1)) : "memory"); + } + } + mbarrier_wait(p_full_2_addr + (pq_2) * 8, _phase_p_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + if (elect_sync()) { + if (pq_2 == 0) { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_0), "r"(tmem_p0), "r"(1) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_0), "r"(tmem_p0), "r"(1) : "memory"); + } else { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_0), "r"(tmem_p1), "r"(1) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_0), "r"(tmem_p1), "r"(1) : "memory"); + } + tcgen05_commit_cg2_multicast(o_full_addr + (oq_1) * 8, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(kv_empty_addr + (kv) * 8, (uint16_t)(3)); + } + kv += 1; + if (kv == 5) { kv = 0; _phase_kv_full ^= 1; } + pq_2 += 1; + if (pq_2 == 2) { pq_2 = 0; _phase_p_full ^= 1; } + oq_1 += 1; + if (oq_1 == 2) { oq_1 = 0; } + } + elect_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + int last_first = ((n_blocks_11 == 1) ? 1 : 0); + #pragma unroll 1 + for (int __2 = 0; __2 < has_block_1; __2++) { + mbarrier_wait(kv_full_addr + (kv) * 8, _phase_kv_full); + mbarrier_wait(p_full_addr + (pq_2) * 8, _phase_p_full); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1000 || __CUDA_ARCH__ == 1030 + int b_hi_2 = ((((smem_v_hi_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv) * 1024; + int b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x3FFF) | 0x4000000) + (kv) * 1024; +#else + int b_hi_2 = ((((smem_v_hi_addr) >> 4) & 0x7FFF) | 0x4000000) + (kv) * 1024; + int b_lo_2 = ((((smem_v_lo_addr) >> 4) & 0x7FFF) | 0x4000000) + (kv) * 1024; +#endif + if (elect_sync()) { + if (pq_2 == 0) { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_2), "r"(tmem_p0), "r"(((last_first) ? 0 : 1)) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_2), "r"(tmem_p0), "r"(((last_first) ? 0 : 1)) : "memory"); + } else { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_2), "r"(tmem_p1), "r"(((last_first) ? 0 : 1)) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "mov.b32 ta, %2;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_2), "r"(tmem_p1), "r"(((last_first) ? 0 : 1)) : "memory"); + } + } + mbarrier_wait(p_full_2_addr + (pq_2) * 8, _phase_p_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + if (elect_sync()) { + if (pq_2 == 0) { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_2), "r"(tmem_p0), "r"(1) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_2), "r"(tmem_p0), "r"(1) : "memory"); + } else { + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_hi), "r"(b_hi_2), "r"(tmem_p1), "r"(1) : "memory"); + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 dhi, blo, ta, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 270599312;\n\t" + "add.u32 ta, %2, 48;\n\t" + "add.u32 blo, %1, 768;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 ta, ta, 8;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%0], [ta], db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(tmem_o_lo), "r"(b_lo_2), "r"(tmem_p1), "r"(1) : "memory"); + } + tcgen05_commit_cg2_multicast(kv_empty_addr + (kv) * 8, (uint16_t)(3)); + } + kv += 1; + if (kv == 5) { kv = 0; _phase_kv_full ^= 1; } + pq_2 += 1; + if (pq_2 == 2) { pq_2 = 0; _phase_p_full ^= 1; } + } + mbarrier_wait(o_free_addr, _phase_o_free); + elect_commit_cg2_multicast(o_full_addr + (oq_1) * 8, (uint16_t)(3)); + oq_1 += 1; + if (oq_1 == 2) { oq_1 = 0; } + _phase_o_free ^= 1; + } + } + if (cta_rank == 0) { + mbarrier_wait(o_free_addr, _phase_o_free); + } + } + } + // ---- Role: load ---- + if (warp == 9) { + { // load_main + unsigned int wk_3 = 0; + unsigned int kv_1 = 0; + unsigned int head_tiles_3 = 0; + #pragma unroll 1 + for (int s_12 = 0; s_12 < num_segments; s_12++) { + head_tiles_3 = head_tiles_3 + ((unsigned int)(cu_seqlens_q[s_12 + 1] - cu_seqlens_q[s_12] + 255) >> 8); + } + int wmax_3 = 0; + #pragma unroll 1 + for (int s_13 = 0; s_13 < num_segments; s_13++) { + int q_len_12 = cu_seqlens_q[s_13 + 1] - cu_seqlens_q[s_13]; + int k_len_12 = cu_seqlens_k[s_13 + 1] - cu_seqlens_k[s_13]; + int row0_12 = ((q_len_12 + 255 >> 8) - 1) * 256; + int diag_12 = k_len_12 - q_len_12; + int n_k_12 = k_len_12 + 127 >> 7; + int n_c_12 = row0_12 + 256 + diag_12 + 127 >> 7; + int n_seen_12 = ((n_c_12 < n_k_12) ? n_c_12 : n_k_12); + int n_blocks_12 = ((n_seen_12 > 0) ? n_seen_12 : 0); + int first_seen_12 = row0_12 + diag_12 + 1; + int n_first_12 = ((first_seen_12 > 0) ? first_seen_12 >> 7 : 0); + int k_full_12 = k_len_12 >> 7; + int n_kf_12 = ((n_first_12 < k_full_12) ? n_first_12 : k_full_12); + int n_full_12 = ((n_kf_12 < n_blocks_12) ? n_kf_12 : n_blocks_12); + int n_masked_12 = n_blocks_12 - n_full_12; + int w_3 = n_blocks_12; + wmax_3 = ((w_3 > wmax_3) ? w_3 : wmax_3); + } + unsigned int t0_3 = 0; + unsigned int t1_3 = 0; + unsigned int t2_3 = 0; + #pragma unroll 1 + for (int s_14 = 0; s_14 < num_segments; s_14++) { + unsigned int tiles_3 = (unsigned int)(cu_seqlens_q[s_14 + 1] - cu_seqlens_q[s_14] + 255) >> 8; + int q_len_13 = cu_seqlens_q[s_14 + 1] - cu_seqlens_q[s_14]; + int k_len_13 = cu_seqlens_k[s_14 + 1] - cu_seqlens_k[s_14]; + int row0_13 = ((q_len_13 + 255 >> 8) - 1) * 256; + int diag_13 = k_len_13 - q_len_13; + int n_k_13 = k_len_13 + 127 >> 7; + int n_c_13 = row0_13 + 256 + diag_13 + 127 >> 7; + int n_seen_13 = ((n_c_13 < n_k_13) ? n_c_13 : n_k_13); + int n_blocks_13 = ((n_seen_13 > 0) ? n_seen_13 : 0); + int first_seen_13 = row0_13 + diag_13 + 1; + int n_first_13 = ((first_seen_13 > 0) ? first_seen_13 >> 7 : 0); + int k_full_13 = k_len_13 >> 7; + int n_kf_13 = ((n_first_13 < k_full_13) ? n_first_13 : k_full_13); + int n_full_13 = ((n_kf_13 < n_blocks_13) ? n_kf_13 : n_blocks_13); + int n_masked_13 = n_blocks_13 - n_full_13; + int c_6 = ((2 * n_blocks_13 >= wmax_3) ? 0 : ((4 * n_blocks_13 >= wmax_3) ? 1 : ((8 * n_blocks_13 >= wmax_3) ? 2 : 3))); + t0_3 = t0_3 + ((c_6 == 0) ? tiles_3 : 0); + t1_3 = t1_3 + ((c_6 == 1) ? tiles_3 : 0); + t2_3 = t2_3 + ((c_6 == 2) ? tiles_3 : 0); + } + unsigned int total_3 = (unsigned int)num_heads_q * head_tiles_3; + unsigned int _phase_work_full_3 = 0; + unsigned int _phase_q_empty = 1; + unsigned int _phase_kv_empty = 1; + #pragma unroll 1 + for (unsigned int it_3 = 0; it_3 < (unsigned int)num_heads_q * ((unsigned int)(total_q + 255 >> 8) + (unsigned int)num_segments) + 1; it_3++) { + mbarrier_wait_cluster_hint(work_full_addr + (wk_3) * 8, _phase_work_full_3, 10000000); + unsigned int ticket_3 = cells[wk_3]; + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(work_empty_addr + (wk_3) * 8), "r"(0) : "memory"); + wk_3 += 1; + if (wk_3 == 2) { wk_3 = 0; _phase_work_full_3 ^= 1; } + if (ticket_3 >= total_3) { + break; + } + unsigned int hq_3 = (unsigned int)num_heads_q; + unsigned int p1_3 = hq_3 * t0_3; + unsigned int p2_3 = p1_3 + hq_3 * t1_3; + unsigned int p3_3 = p2_3 + hq_3 * t2_3; + int c_7 = ((ticket_3 < p1_3) ? 0 : ((ticket_3 < p2_3) ? 1 : ((ticket_3 < p3_3) ? 2 : 3))); + unsigned int base_3 = ((c_7 == 0) ? 0 : ((c_7 == 1) ? p1_3 : ((c_7 == 2) ? p2_3 : p3_3))); + unsigned int tc_3 = ((c_7 == 0) ? t0_3 : ((c_7 == 1) ? t1_3 : ((c_7 == 2) ? t2_3 : head_tiles_3 - t0_3 - t1_3 - t2_3))); + unsigned int r_3 = ticket_3 - base_3; + unsigned int per_kv_3 = tc_3 * (unsigned int)group; + unsigned int kvh_u_3 = r_3 / per_kv_3; + unsigned int rem_3 = r_3 - kvh_u_3 * per_kv_3; + unsigned int j_3 = rem_3 / (unsigned int)group; + unsigned int head_u_3 = kvh_u_3 * (unsigned int)group + (rem_3 - j_3 * (unsigned int)group); + int seg_3 = 0; + unsigned int before_3 = 0; + unsigned int ends_3 = 0; + #pragma unroll 1 + for (int s_15 = 0; s_15 < num_segments; s_15++) { + int q_len_14 = cu_seqlens_q[s_15 + 1] - cu_seqlens_q[s_15]; + int k_len_14 = cu_seqlens_k[s_15 + 1] - cu_seqlens_k[s_15]; + int row0_14 = ((q_len_14 + 255 >> 8) - 1) * 256; + int diag_14 = k_len_14 - q_len_14; + int n_k_14 = k_len_14 + 127 >> 7; + int n_c_14 = row0_14 + 256 + diag_14 + 127 >> 7; + int n_seen_14 = ((n_c_14 < n_k_14) ? n_c_14 : n_k_14); + int n_blocks_14 = ((n_seen_14 > 0) ? n_seen_14 : 0); + int first_seen_14 = row0_14 + diag_14 + 1; + int n_first_14 = ((first_seen_14 > 0) ? first_seen_14 >> 7 : 0); + int k_full_14 = k_len_14 >> 7; + int n_kf_14 = ((n_first_14 < k_full_14) ? n_first_14 : k_full_14); + int n_full_14 = ((n_kf_14 < n_blocks_14) ? n_kf_14 : n_blocks_14); + int n_masked_14 = n_blocks_14 - n_full_14; + bool in_class_3 = ((2 * n_blocks_14 >= wmax_3) ? 0 : ((4 * n_blocks_14 >= wmax_3) ? 1 : ((8 * n_blocks_14 >= wmax_3) ? 2 : 3))) == c_7; + ends_3 = ends_3 + ((in_class_3) ? (unsigned int)(cu_seqlens_q[s_15 + 1] - cu_seqlens_q[s_15] + 255) >> 8 : 0); + seg_3 = seg_3 + ((ends_3 <= j_3) ? 1 : 0); + before_3 = ((ends_3 <= j_3) ? ends_3 : before_3); + } + unsigned int m_rel_3 = j_3 - before_3; + int q0_3 = cu_seqlens_q[seg_3]; + int q_len_15 = cu_seqlens_q[seg_3 + 1] - q0_3; + int k0_3 = cu_seqlens_k[seg_3]; + int k_len_15 = cu_seqlens_k[seg_3 + 1] - k0_3; + int m_3 = (int)(((unsigned int)(q_len_15 + 255) >> 8) - 1 - m_rel_3); + int head_3 = (int)head_u_3; + int kvh_3 = (int)kvh_u_3; + int row0_15 = m_3 * 256; + int diag_15 = k_len_15 - q_len_15; + int n_k_15 = k_len_15 + 127 >> 7; + int n_c_15 = row0_15 + 256 + diag_15 + 127 >> 7; + int n_seen_15 = ((n_c_15 < n_k_15) ? n_c_15 : n_k_15); + int n_blocks_15 = ((n_seen_15 > 0) ? n_seen_15 : 0); + int first_seen_15 = row0_15 + diag_15 + 1; + int n_first_15 = ((first_seen_15 > 0) ? first_seen_15 >> 7 : 0); + int k_full_15 = k_len_15 >> 7; + int n_kf_15 = ((n_first_15 < k_full_15) ? n_first_15 : k_full_15); + int n_full_15 = ((n_kf_15 < n_blocks_15) ? n_kf_15 : n_blocks_15); + int n_masked_15 = n_blocks_15 - n_full_15; + int has_block_2 = ((n_blocks_15 > 0) ? 1 : 0); + int q_row = q0_3 + m_3 * 256 + cta_rank * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty); + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(65536)) : "memory"); + tma_4d_gmem2smem_cta2(smem_q_hi_addr, (&Q), 0, q_row, head_3, 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(smem_q_lo_addr, (&Q), 0, q_row, head_3, 2, ((q_full_addr) & 0xFEFFFFFF)); + } + _phase_q_empty ^= 1; + int k_row = k0_3 + cta_rank * 64; + #pragma unroll 1 + for (int __3 = 0; __3 < has_block_2; __3++) { + mbarrier_wait(kv_empty_addr + (kv_1) * 8, _phase_kv_empty); + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_4d_gmem2smem_cta2(smem_k_hi_addr + kv_1 * 16384, (&K), 0, k_row + (n_blocks_15 - 1) * 128, kvh_3, 0, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(smem_k_lo_addr + kv_1 * 16384, (&K), 0, k_row + (n_blocks_15 - 1) * 128, kvh_3, 2, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + } + kv_1 += 1; + if (kv_1 == 5) { kv_1 = 0; _phase_kv_empty ^= 1; } + } + #pragma unroll 1 + for (int i_4 = 1; i_4 < n_blocks_15; i_4++) { + mbarrier_wait(kv_empty_addr + (kv_1) * 8, _phase_kv_empty); + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_4d_gmem2smem_cta2(smem_k_hi_addr + kv_1 * 16384, (&K), 0, k_row + (n_blocks_15 - 1 - i_4) * 128, kvh_3, 0, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(smem_k_lo_addr + kv_1 * 16384, (&K), 0, k_row + (n_blocks_15 - 1 - i_4) * 128, kvh_3, 2, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + } + kv_1 += 1; + if (kv_1 == 5) { kv_1 = 0; _phase_kv_empty ^= 1; } + mbarrier_wait(kv_empty_addr + (kv_1) * 8, _phase_kv_empty); + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_4d_gmem2smem_cta2(smem_v_hi_addr + kv_1 * 16384, (&V), 0, k0_3 + (n_blocks_15 - i_4) * 128, cta_rank, kvh_3, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(smem_v_lo_addr + kv_1 * 16384, (&V), 0, k0_3 + (n_blocks_15 - i_4) * 128, cta_rank + 2, kvh_3, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + } + kv_1 += 1; + if (kv_1 == 5) { kv_1 = 0; _phase_kv_empty ^= 1; } + } + #pragma unroll 1 + for (int __4 = 0; __4 < has_block_2; __4++) { + mbarrier_wait(kv_empty_addr + (kv_1) * 8, _phase_kv_empty); + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_4d_gmem2smem_cta2(smem_v_hi_addr + kv_1 * 16384, (&V), 0, k0_3, cta_rank, kvh_3, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + tma_4d_gmem2smem_cta2(smem_v_lo_addr + kv_1 * 16384, (&V), 0, k0_3, cta_rank + 2, kvh_3, ((kv_full_addr + (kv_1) * 8) & 0xFEFFFFFF)); + } + kv_1 += 1; + if (kv_1 == 5) { kv_1 = 0; _phase_kv_empty ^= 1; } + } + } + #pragma unroll 1 + for (int __5 = 0; __5 < 5; __5++) { + mbarrier_wait(kv_empty_addr + (kv_1) * 8, _phase_kv_empty); + kv_1 += 1; + if (kv_1 == 5) { kv_1 = 0; _phase_kv_empty ^= 1; } + } + mbarrier_wait(q_empty_addr, _phase_q_empty); + } + } + // ---- Role: sched ---- + if (warp == 10) { + { // sched_main + unsigned int wk_4 = 0; + unsigned int head_tiles_4 = 0; + #pragma unroll 1 + for (int s_16 = 0; s_16 < num_segments; s_16++) { + head_tiles_4 = head_tiles_4 + ((unsigned int)(cu_seqlens_q[s_16 + 1] - cu_seqlens_q[s_16] + 255) >> 8); + } + unsigned int total_4 = (unsigned int)num_heads_q * head_tiles_4; + unsigned int _phase_work_empty = 1; + if (cta_rank == 0) { + if (elect_sync()) { + #pragma unroll 1 + for (unsigned int it_4 = 0; it_4 < (unsigned int)num_heads_q * ((unsigned int)(total_q + 255 >> 8) + (unsigned int)num_segments) + 1; it_4++) { + mbarrier_wait_cluster_hint(work_empty_addr + (wk_4) * 8, _phase_work_empty, 10000000); + unsigned int ticket_4 = atomicAdd(tile_counter, 1); + cells[wk_4] = ticket_4; + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.release.cluster.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(work_full_addr + (wk_4) * 8), "r"(1), "r"((uint32_t)(4)) : "memory"); + uint32_t _cluster_mapa_0; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_cluster_mapa_0) : "r"(cells_addr + wk_4 * 4), "r"(1)); + uint32_t _mbar_peer_1; + asm( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mbar_peer_1) : "r"(work_full_addr + (wk_4) * 8), "r"(1)); + asm volatile( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 [%0], %1, [%2];" + :: "r"(_cluster_mapa_0), "r"(ticket_4), "r"(_mbar_peer_1) : "memory"); + mbarrier_arrive(work_full_addr + (wk_4) * 8); + wk_4 += 1; + if (wk_4 == 2) { wk_4 = 0; _phase_work_empty ^= 1; } + if (ticket_4 >= total_4) { + if (ticket_4 == total_4 + (unsigned int)num_clusters - 1) { + *(reinterpret_cast(tile_counter) + (0)) = 0; + } + break; + } + } + } + } + } + } + + // Cleanup + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + if (warp == 0) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + } +} + +} // extern "C" diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu new file mode 100644 index 0000000..39577c1 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu @@ -0,0 +1,349 @@ +/* + * 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 Cake: the tvm-ffi host binding of kernel_cake_gqa_train_6fa13db592acdb66a6c9, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" +#include + +#include + +extern "C" __global__ void kernel_cake_gqa_train_6fa13db592acdb66a6c9(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap dO, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, float* __restrict__ lse_log2, float* __restrict__ dpsum, float* __restrict__ dq_accum, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, float* __restrict__ dk_out, float* __restrict__ dv_out, int num_segments, int h_kv, int group, int rows_pad, float softmax_scale, float softmax_scale_log2, int* __restrict__ row_map); + +namespace cake_binding_kernel_cake_gqa_train_6fa13db592acdb66a6c9 { + +using tvm::ffi::Optional; +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gqa_train_6fa13db592acdb66a6c9"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void CheckCu(CUresult result, const char* what) { + if (result != CUDA_SUCCESS) { + const char* name = nullptr; + const char* message = nullptr; + cuGetErrorName(result, &name); + cuGetErrorString(result, &message); + TVM_FFI_THROW(RuntimeError) << KERNEL_NAME << ": " << what << " failed with " + << (name ? name : "an unknown CUresult") << " (" << static_cast(result) + << "): " << (message ? message : ""); + } +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void CheckUnitInnerStride(const TensorView& t, const char* name) { + TVM_FFI_CHECK(t.stride(-1) == 1, ValueError) + << KERNEL_NAME << ": '" << name << "' needs a unit innermost stride, got " << t.stride(-1); +} + +[[noreturn]] void NoInt64(int64_t a, const char* op, int64_t b) { + TVM_FFI_THROW(ValueError) << KERNEL_NAME << ": " << a << " " << op << " " << b + << " has no int64 value in a host check"; +} + +int64_t CMod(int64_t a, int64_t b) { + if (b == 0) NoInt64(a, "%", b); + return b == -1 ? 0 : a % b; +} + +cuuint64_t Extent(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value > 0, ValueError) + << KERNEL_NAME << ": an extent of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +cuuint64_t Stride(int64_t value, const char* descriptor) { + TVM_FFI_CHECK(value >= 0, ValueError) + << KERNEL_NAME << ": a stride of descriptor '" << descriptor << "' resolves to " << value; + return static_cast(value); +} + +CUtensorMap EncodeTma_Q(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'Q' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "Q"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'Q' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "Q"), Extent(d2, "Q"), Extent(d3, "Q")}; + cuuint64_t global_strides[2] = {Stride(s2, "Q"), Stride(s3, "Q")}; + cuuint32_t box_dim[3] = {64u, 2u, 32u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'Q' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'Q' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'Q' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of Q"); + return map; +} + +CUtensorMap EncodeTma_dO(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'dO' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "dO"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'dO' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "dO"), Extent(d2, "dO"), Extent(d3, "dO")}; + cuuint64_t global_strides[2] = {Stride(s2, "dO"), Stride(s3, "dO")}; + cuuint32_t box_dim[3] = {64u, 2u, 32u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'dO' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'dO' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'dO' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of dO"); + return map; +} + +CUtensorMap EncodeTma_K(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'K' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "K"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'K' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "K"), Extent(d2, "K"), Extent(d3, "K")}; + cuuint64_t global_strides[2] = {Stride(s2, "K"), Stride(s3, "K")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'K' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'K' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'K' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of K"); + return map; +} + +CUtensorMap EncodeTma_V(const TensorView& t) { + TVM_FFI_CHECK(t.ndim() >= 3, ValueError) + << KERNEL_NAME << ": descriptor 'V' takes a source of at least 3 dims, got " << t.ndim(); + CheckUnitInnerStride(t, "V"); + int64_t d1 = t.size(-1); + int64_t d2 = t.size(-2); + int64_t d3 = t.size(-3); + int64_t s2 = t.stride(-2) * 2; + int64_t s3 = t.stride(-3) * 2; + TVM_FFI_CHECK((d1 == 256LL), ValueError) + << KERNEL_NAME << ": descriptor 'V' requires (d1 == 256)"; + cuuint64_t global_dim[3] = {Extent(d1, "V"), Extent(d2, "V"), Extent(d3, "V")}; + cuuint64_t global_strides[2] = {Stride(s2, "V"), Stride(s3, "V")}; + cuuint32_t box_dim[3] = {64u, 1u, 64u}; + cuuint32_t elem_strides[3] = {1u, 1u, 1u}; + const bool allow_oob_box[3] = {false, false, true}; + for (int axis = 0; axis < 3; ++axis) { + TVM_FFI_CHECK(allow_oob_box[axis] || box_dim[axis] <= global_dim[axis], ValueError) + << KERNEL_NAME << ": descriptor 'V' axis " << axis << " (0 is the innermost) has a TMA box of " + << box_dim[axis] << " elements but the source extent along it is " << global_dim[axis] + << ". A box may not exceed its extent unless the axis is listed in the descriptor's allow_oob_box (CK.describe_tma(..., allow_oob_box=(axis,))). List it when the box is a tile that may run past the end of the axis, a tail tile: a load fills the part past the extent with zeros (its bytes still count toward the barrier's transaction) and a store or reduce skips it, so that is safe when the kernel ignores or masks those rows and columns. If the box was meant to fit, box_dim or the extent is wrong: check box_dim, global_dim and the tensor's shape"; + } + TVM_FFI_CHECK(global_strides[0] % 16 == 0 && global_strides[1] % 16 == 0, ValueError) + << KERNEL_NAME << ": the strides of descriptor 'V' must be multiples of 16 bytes"; + void* address = DataPtr(t); + TVM_FFI_CHECK(reinterpret_cast(address) % 16 == 0, ValueError) + << KERNEL_NAME << ": the address of descriptor 'V' must be 16-byte aligned"; + CUtensorMap map{}; + CheckCu(cuTensorMapEncodeTiled(&map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, address, global_dim, global_strides, + box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_128B, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE), + "cuTensorMapEncodeTiled of V"); + return map; +} + +void Run(TensorView a_Q, TensorView a_dO, TensorView a_K, TensorView a_V, TensorView a_lse_log2, TensorView a_dpsum, TensorView a_dq_accum, TensorView a_cu_seqlens_q, TensorView a_cu_seqlens_k, TensorView a_dk_out, TensorView a_dv_out, int64_t a_num_segments, int64_t a_h_kv, int64_t a_group, int64_t a_rows_pad, double a_softmax_scale, double a_softmax_scale_log2, Optional a_row_map, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_Q.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_Q, "Q"); + check_dtype(a_Q, DLDataType{kDLBfloat, 16, 1}, "Q"); + check_cuda_tensor(a_dO, "dO"); + check_dtype(a_dO, DLDataType{kDLBfloat, 16, 1}, "dO"); + check_cuda_tensor(a_K, "K"); + check_dtype(a_K, DLDataType{kDLBfloat, 16, 1}, "K"); + check_cuda_tensor(a_V, "V"); + check_dtype(a_V, DLDataType{kDLBfloat, 16, 1}, "V"); + check_cuda_tensor(a_lse_log2, "lse_log2"); + check_dtype(a_lse_log2, DLDataType{kDLFloat, 32, 1}, "lse_log2"); + check_cuda_tensor(a_dpsum, "dpsum"); + check_dtype(a_dpsum, DLDataType{kDLFloat, 32, 1}, "dpsum"); + check_cuda_tensor(a_dq_accum, "dq_accum"); + check_dtype(a_dq_accum, DLDataType{kDLFloat, 32, 1}, "dq_accum"); + check_cuda_tensor(a_cu_seqlens_q, "cu_seqlens_q"); + check_dtype(a_cu_seqlens_q, DLDataType{kDLInt, 32, 1}, "cu_seqlens_q"); + check_cuda_tensor(a_cu_seqlens_k, "cu_seqlens_k"); + check_dtype(a_cu_seqlens_k, DLDataType{kDLInt, 32, 1}, "cu_seqlens_k"); + check_cuda_tensor(a_dk_out, "dk_out"); + check_dtype(a_dk_out, DLDataType{kDLFloat, 32, 1}, "dk_out"); + check_cuda_tensor(a_dv_out, "dv_out"); + check_dtype(a_dv_out, DLDataType{kDLFloat, 32, 1}, "dv_out"); + TVM_FFI_CHECK(a_num_segments >= -2147483648LL && a_num_segments <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_segments' = " << a_num_segments + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_h_kv >= -2147483648LL && a_h_kv <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'h_kv' = " << a_h_kv + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_group >= -2147483648LL && a_group <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'group' = " << a_group + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_pad >= -2147483648LL && a_rows_pad <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_pad' = " << a_rows_pad + << " is outside the i32 range [-2147483648, 2147483647]"; + if (a_row_map.has_value()) check_cuda_tensor(a_row_map.value(), "row_map"); + if (a_row_map.has_value()) check_dtype(a_row_map.value(), DLDataType{kDLInt, 32, 1}, "row_map"); + check_same_device(a_dO, a_Q, "dO", "Q"); + check_same_device(a_K, a_Q, "K", "Q"); + check_same_device(a_V, a_Q, "V", "Q"); + check_contiguous(a_lse_log2, "lse_log2"); + check_same_device(a_lse_log2, a_Q, "lse_log2", "Q"); + check_contiguous(a_dpsum, "dpsum"); + check_same_device(a_dpsum, a_Q, "dpsum", "Q"); + check_contiguous(a_dq_accum, "dq_accum"); + check_same_device(a_dq_accum, a_Q, "dq_accum", "Q"); + check_contiguous(a_cu_seqlens_q, "cu_seqlens_q"); + check_same_device(a_cu_seqlens_q, a_Q, "cu_seqlens_q", "Q"); + check_contiguous(a_cu_seqlens_k, "cu_seqlens_k"); + check_same_device(a_cu_seqlens_k, a_Q, "cu_seqlens_k", "Q"); + check_contiguous(a_dk_out, "dk_out"); + check_same_device(a_dk_out, a_Q, "dk_out", "Q"); + check_contiguous(a_dv_out, "dv_out"); + check_same_device(a_dv_out, a_Q, "dv_out", "Q"); + if (a_row_map.has_value()) check_contiguous(a_row_map.value(), "row_map"); + if (a_row_map.has_value()) check_same_device(a_row_map.value(), a_Q, "row_map", "Q"); + TVM_FFI_CHECK((a_group >= 2LL), ValueError) + << KERNEL_NAME << ": requires (group >= 2)"; + TVM_FFI_CHECK((a_group <= 64LL), ValueError) + << KERNEL_NAME << ": requires (group <= 64)"; + TVM_FFI_CHECK((CMod(a_group, 2LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(group, 2) == 0)"; + TVM_FFI_CHECK((CMod(a_rows_pad, 128LL) == 0LL), ValueError) + << KERNEL_NAME << ": requires (CMod(rows_pad, 128) == 0)"; + TVM_FFI_CHECK((a_num_segments >= 1LL), ValueError) + << KERNEL_NAME << ": requires (num_segments >= 1)"; + TVM_FFI_CHECK((a_h_kv >= 1LL), ValueError) + << KERNEL_NAME << ": requires (h_kv >= 1)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + TVM_FFI_CHECK(grid_x % 2 == 0 && grid_y % 1 == 0 && grid_z % 1 == 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z + << ") is not a multiple of the cluster shape (2, 1, 1)"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + { + // Once per thread per device: the kernel's function attributes hold per device; the tensor-map encoders need this thread's context of the device bound. + static thread_local uint64_t prepared_devices = 0; + const uint64_t device_bit = dev.device_id < 64 ? uint64_t{1} << dev.device_id : uint64_t{0}; + if (device_bit == 0 || (prepared_devices & device_bit) == 0) { + CheckCuda(cudaSetDevice(dev.device_id), "cudaSetDevice"); + CheckCuda(cudaFuncSetAttribute(reinterpret_cast(kernel_cake_gqa_train_6fa13db592acdb66a6c9), cudaFuncAttributeMaxDynamicSharedMemorySize, 231424), "cudaFuncSetAttribute"); + prepared_devices |= device_bit; + } + } + CUtensorMap v_Q = EncodeTma_Q(a_Q); + CUtensorMap v_dO = EncodeTma_dO(a_dO); + CUtensorMap v_K = EncodeTma_K(a_K); + CUtensorMap v_V = EncodeTma_V(a_V); + void* v_lse_log2 = DataPtr(a_lse_log2); + void* v_dpsum = DataPtr(a_dpsum); + void* v_dq_accum = DataPtr(a_dq_accum); + void* v_cu_seqlens_q = DataPtr(a_cu_seqlens_q); + void* v_cu_seqlens_k = DataPtr(a_cu_seqlens_k); + void* v_dk_out = DataPtr(a_dk_out); + void* v_dv_out = DataPtr(a_dv_out); + int32_t v_num_segments = static_cast(a_num_segments); + int32_t v_h_kv = static_cast(a_h_kv); + int32_t v_group = static_cast(a_group); + int32_t v_rows_pad = static_cast(a_rows_pad); + float v_softmax_scale = static_cast(a_softmax_scale); + float v_softmax_scale_log2 = static_cast(a_softmax_scale_log2); + void* v_row_map = a_row_map.has_value() ? DataPtr(a_row_map.value()) : nullptr; + void* kernel_args[] = {&v_Q, &v_dO, &v_K, &v_V, &v_lse_log2, &v_dpsum, &v_dq_accum, &v_cu_seqlens_q, &v_cu_seqlens_k, &v_dk_out, &v_dv_out, &v_num_segments, &v_h_kv, &v_group, &v_rows_pad, &v_softmax_scale, &v_softmax_scale_log2, &v_row_map}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(512, 1, 1); + config.dynamicSmemBytes = 231424; + config.stream = stream; + cudaLaunchAttribute attrs[2] = {}; + attrs[0].id = cudaLaunchAttributeClusterDimension; + attrs[0].val.clusterDim.x = 2; + attrs[0].val.clusterDim.y = 1; + attrs[0].val.clusterDim.z = 1; + attrs[1].id = cudaLaunchAttributeClusterSchedulingPolicyPreference; + attrs[1].val.clusterSchedulingPolicyPreference = cudaClusterSchedulingPolicySpread; + config.attrs = attrs; + config.numAttrs = 2; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gqa_train_6fa13db592acdb66a6c9), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gqa_train_6fa13db592acdb66a6c9 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gqa_train_6fa13db592acdb66a6c9::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu new file mode 100644 index 0000000..d690061 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu @@ -0,0 +1,1603 @@ +/* + * 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 TMEM_NCOLS 512 +#define TMEM_DV_OFFSET 0 +#define TMEM_DK_OFFSET 128 +#define TMEM_S_OFFSET 256 +#define TMEM_DQ_OFFSET 384 +#define TMEM_DP_OFFSET 320 +#define NUM_Q_PIPE_STAGES 1 +#define NUM_DO_PIPE_STAGES 1 +#define SMEM_SQ_OFF 1024 +#define SMEM_SQ_STAGE_BYTES 32768 +#define SMEM_SQ_STRIDE 32768 +#define SMEM_SQT_OFF 1024 +#define SMEM_SQT_STAGE_BYTES 32768 +#define SMEM_SQT_STRIDE 32768 +#define SMEM_SDO_OFF 33792 +#define SMEM_SDO_STAGE_BYTES 32768 +#define SMEM_SDO_STRIDE 32768 +#define SMEM_SDOT_OFF 33792 +#define SMEM_SDOT_STAGE_BYTES 32768 +#define SMEM_SDOT_STRIDE 32768 +#define SMEM_SK_OFF 66560 +#define SMEM_SK_STAGE_BYTES 32768 +#define SMEM_SK_STRIDE 32768 +#define SMEM_SV_OFF 99328 +#define SMEM_SV_STAGE_BYTES 32768 +#define SMEM_SV_STRIDE 32768 +#define SMEM_SKT_OFF 132096 +#define SMEM_SKT_STAGE_BYTES 32768 +#define SMEM_SKT_STRIDE 32768 +#define SMEM_SP_OFF 164864 +#define SMEM_SP_STAGE_BYTES 16384 +#define SMEM_SP_STRIDE 16384 +#define SMEM_SDST_OFF 181248 +#define SMEM_SDST_STAGE_BYTES 16384 +#define SMEM_SDST_STRIDE 16384 +#define SMEM_SDS_OFF 197632 +#define SMEM_SDS_STAGE_BYTES 16384 +#define SMEM_SDS_STRIDE 16384 +#define SMEM_SDQ_OFF 214016 +#define SMEM_SDQ_STAGE_BYTES 8192 +#define SMEM_SDQ_STRIDE 8192 +#define SMEM_SLSE_OFF 230400 +#define SMEM_SLSE_STAGE_BYTES 512 +#define SMEM_SLSE_STRIDE 512 +#define SMEM_SDPS_OFF 230912 +#define SMEM_SDPS_STAGE_BYTES 512 +#define SMEM_SDPS_STRIDE 512 +#define SMEM_TOTAL 231424 +#define LAUNCH_MIN_BLOCKS 1 + +#include + +__device__ __forceinline__ uint32_t elect_sync() { + uint32_t pred = 0; + asm volatile( + "{\n\t" + ".reg .pred %%px;\n\t" + "elect.sync _|%%px, %1;\n\t" + "@%%px mov.s32 %0, 1;\n\t" + "}\n" + : "+r"(pred) + : "r"(0xFFFFFFFF)); + return pred; +} + + +__device__ __forceinline__ void mbarrier_init(int mbar_addr, int count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" + :: "r"(mbar_addr), "r"(count) : "memory"); +} + + + +// CTA-local pipelines have short, resident producer/consumer edges. Omitting +// suspendTimeHint keeps a miss on the lightweight TRYWAIT retry path; the +// explicit loop still makes this helper blocking until acquire succeeds. +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT:\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%0], %1;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster_hint( + int mbar_addr, int phase, uint32_t suspend_time_hint) { + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER_HINT:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER_HINT;\n\t" + "bra.uni LAB_WAIT_CLUSTER_HINT;\n\t" + "DONE_CLUSTER_HINT:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(suspend_time_hint) : "memory"); +} + + + + +union MmaSmemDesc { + uint64_t u64; + uint32_t u32[2]; +}; + + + +__device__ __forceinline__ void mbarrier_arrive(int mbar_addr) { + asm volatile( + "mbarrier.arrive.release.cta.shared::cta.b64 _, [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void mbarrier_arrive_expect_tx(int mbar_addr, uint32_t bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;" + :: "r"(mbar_addr), "r"(bytes) : "memory"); +} + + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__device__ __forceinline__ float max_noftz(float a, float b) { + float c; + asm("max.f32 %0, %1, %2;" : "=f"(c) : "f"(a), "f"(b)); + return c; +} + + +__device__ __forceinline__ float2 ex2_emulation_f32x2_value(float2 value) { + const float c0 = 1.0f, c1 = 0.695146143436431884765625f; + const float c2 = 0.227564394474029541015625f, c3 = 0.077119089663028717041015625f; + const float magic = 12582912.0f; + float x0 = max_noftz(value.x, -127.0f), x1 = max_noftz(value.y, -127.0f); + float2 xc2 = make_float2(x0, x1), magic2 = make_float2(magic, magic); + float2 xr2; + asm("add.rm.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xr2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&magic2)); + float2 c3_2 = make_float2(c3, c3), c2_2 = make_float2(c2, c2); + float2 c1_2 = make_float2(c1, c1), c0_2 = make_float2(c0, c0); + float2 xrb2, xfrac2; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xrb2) + : "l"(*(unsigned long long*)&xr2), "l"(*(unsigned long long*)&magic2)); + asm("sub.rn.ftz.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&xfrac2) + : "l"(*(unsigned long long*)&xc2), "l"(*(unsigned long long*)&xrb2)); + float2 poly2; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&c3_2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c2_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c1_2)); + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" : "=l"(*(unsigned long long*)&poly2) + : "l"(*(unsigned long long*)&poly2), "l"(*(unsigned long long*)&xfrac2), "l"(*(unsigned long long*)&c0_2)); + int x0r_i, x1r_i, p0_i, p1_i; + asm("mov.b64 {%0, %1}, %2;" : "=r"(x0r_i), "=r"(x1r_i) : "l"(*(unsigned long long*)&xr2)); + asm("mov.b64 {%0, %1}, %2;" : "=r"(p0_i), "=r"(p1_i) : "l"(*(unsigned long long*)&poly2)); + float r0, r1; + asm("mov.b32 %0, %1;" : "=f"(r0) : "r"((x0r_i << 23) + p0_i)); + asm("mov.b32 %0, %1;" : "=f"(r1) : "r"((x1r_i << 23) + p1_i)); + return make_float2(r0, r1); +} + +__device__ __forceinline__ void ex2_emulation_f32x2(float* x0_ptr, float* x1_ptr) { + float2 result = ex2_emulation_f32x2_value(make_float2(*x0_ptr, *x1_ptr)); + *x0_ptr = result.x; *x1_ptr = result.y; +} + +__device__ __forceinline__ void softmax_frag_exp2_cast( + float* sv, uint32_t* pv, int use_emu) +{ + #pragma unroll + for (int j = 0; j < 16; j++) { + if (use_emu && j >= 12) + ex2_emulation_f32x2(&sv[j*2], &sv[j*2+1]); + else { + sv[j*2] = approx_exp2(sv[j*2]); + sv[j*2+1] = approx_exp2(sv[j*2+1]); + } + } + #pragma unroll + for (int j = 0; j < 16; j++) { + __nv_bfloat162 bf = __float22bfloat162_rn({sv[j*2], sv[j*2+1]}); + pv[j] = reinterpret_cast(bf); + } +} + + + +__device__ __forceinline__ void mul_f32x2_noftz_inplace(float2* a, float2 b) { + asm("mul.rn.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void sub_f32x2_inplace(float2* a, float2 b) { + asm("sub.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + +__device__ __forceinline__ float2 fma_sub_f32x2_rn_ftz(float2 a, float2 b, float2 c) { + float2 r; + asm volatile("{\n\t" + ".reg .f32 _c0, _c1;\n\t" + ".reg .b64 _neg_c;\n\t" + "mov.b64 {_c0, _c1}, %3;\n\t" + "neg.f32 _c0, _c0;\n\t" + "neg.f32 _c1, _c1;\n\t" + "mov.b64 _neg_c, {_c0, _c1};\n\t" + "fma.rn.ftz.f32x2 %0, %1, %2, _neg_c;\n\t" + "}\n" + : "=l"(*(unsigned long long*)&r) + : "l"(*(const unsigned long long*)&a), + "l"(*(const unsigned long long*)&b), + "l"(*(const unsigned long long*)&c)); + return r; +} + + + + +__device__ __forceinline__ void tma_3d_gmem2smem_cta2( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cluster.global" + ".mbarrier::complete_tx::bytes.cta_group::2" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void cp_async_bulk_gmem2smem( + unsigned smem_addr, const void* gmem_ptr, unsigned bytes, int mbar_addr) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(smem_addr), "l"(gmem_ptr), "r"(bytes), "r"(mbar_addr) + : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit_cg2_multicast(int mbar_addr, uint16_t cta_mask) { + asm volatile( + "{\n\t" + ".reg .b16 lo, hi;\n\t" + "mov.b32 {lo, hi}, %1;\n\t" + "tcgen05.commit.cta_group::2.mbarrier::arrive::one" + ".shared::cluster.multicast::cluster.b64 [%0], lo;\n\t" + "}\n" + :: "r"(mbar_addr), "r"((uint32_t)cta_mask) : "memory"); +} + +extern "C" { + +__global__ __launch_bounds__(512, LAUNCH_MIN_BLOCKS) __cluster_dims__(2,1,1) void +kernel_cake_gqa_train_6fa13db592acdb66a6c9(const __grid_constant__ CUtensorMap Q, const __grid_constant__ CUtensorMap dO, const __grid_constant__ CUtensorMap K, const __grid_constant__ CUtensorMap V, float* __restrict__ lse_log2, float* __restrict__ dpsum, float* __restrict__ dq_accum, int* __restrict__ cu_seqlens_q, int* __restrict__ cu_seqlens_k, float* __restrict__ dk_out, float* __restrict__ dv_out, int num_segments, int h_kv, int group, int rows_pad, float softmax_scale, float softmax_scale_log2, int* __restrict__ row_map) +{ + const int tid = threadIdx.x; + const int warp = make_warp_uniform(tid / 32); + const int lane = tid % 32; + + extern __shared__ __align__(1024) char smem_raw[]; + int smem; +#if __CUDA_ARCH__ == 1000 + asm volatile("{ .reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, smem_ptr; }" : "=r"(smem) : "l"(smem_raw)); + smem = make_warp_uniform(smem); +#else + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); +#endif + + const int mbar_base = smem; + #define kv_full_addr (mbar_base + 0) + #define q_full_addr (mbar_base + 8) + #define q_empty_addr (mbar_base + 16) + #define do_full_addr (mbar_base + 24) + #define do_empty_addr (mbar_base + 32) + #define lse_full_addr (mbar_base + 40) + #define lse_empty_addr (mbar_base + 48) + #define dpsum_full_addr (mbar_base + 56) + #define dpsum_empty_addr (mbar_base + 64) + #define s_full_addr (mbar_base + 72) + #define dp_full_addr (mbar_base + 80) + #define ds_full_addr (mbar_base + 88) + #define ds_empty_addr (mbar_base + 96) + #define xchg_full_addr (mbar_base + 104) + #define xchg_done_addr (mbar_base + 112) + #define dq_full_addr (mbar_base + 120) + #define dq_empty_addr (mbar_base + 128) + #define dqlo_empty_addr (mbar_base + 136) + #define dkv_full_addr (mbar_base + 144) + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + const unsigned int clusters_x = gridDim.x / 2; + const unsigned int cluster_id = ((blockIdx.z * gridDim.y + blockIdx.y) * clusters_x) + blockIdx.x / 2; + const unsigned int num_clusters = clusters_x * gridDim.y * gridDim.z; + + int cta_rank; + asm volatile("mov.b32 %0, %%cluster_ctarank;" : "=r"(cta_rank)); + + // Kernel setup ops + __nv_bfloat16* sq = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sq_addr = smem + 1024; + __nv_bfloat16* sqt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int sqt_addr = smem + 1024; + __nv_bfloat16* sdo = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int sdo_addr = smem + 33792; + __nv_bfloat16* sdot = reinterpret_cast<__nv_bfloat16*>(smem_raw + 33792); + const int sdot_addr = smem + 33792; + __nv_bfloat16* sk = reinterpret_cast<__nv_bfloat16*>(smem_raw + 66560); + const int sk_addr = smem + 66560; + __nv_bfloat16* sv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 99328); + const int sv_addr = smem + 99328; + __nv_bfloat16* skt = reinterpret_cast<__nv_bfloat16*>(smem_raw + 132096); + const int skt_addr = smem + 132096; + __nv_bfloat16* sp = reinterpret_cast<__nv_bfloat16*>(smem_raw + 164864); + const int sp_addr = smem + 164864; + __nv_bfloat16* sdst = reinterpret_cast<__nv_bfloat16*>(smem_raw + 181248); + const int sdst_addr = smem + 181248; + __nv_bfloat16* sds = reinterpret_cast<__nv_bfloat16*>(smem_raw + 197632); + const int sds_addr = smem + 197632; + float* sdq = reinterpret_cast(smem_raw + 214016); + const int sdq_addr = smem + 214016; + float* slse = reinterpret_cast(smem_raw + 230400); + const int slse_addr = smem + 230400; + float* sdps = reinterpret_cast(smem_raw + 230912); + const int sdps_addr = smem + 230912; + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&Q))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&dO))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&K))) : "memory"); } + if (warp == 13) { asm volatile("prefetch.tensormap [%0];" :: "l"((uint64_t)((&V))) : "memory"); } + + // Mbarrier init (19 pipeline groups, 0 ordered-sequence groups, 19 barriers) + // Mbarriers at smem_raw[0..152) + + if (warp == 0) { + uint32_t leader = elect_sync(); + if (leader) { + // kv_full: 1 barriers, init_count=2 + mbarrier_init(smem + 0, 2); + // --- pipeline 'q_pipe' --- + // q_full: 1 barriers, init_count=2 + mbarrier_init(smem + 8, 2); + // q_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 16, 1); + // --- pipeline 'do_pipe' --- + // do_full: 1 barriers, init_count=2 + mbarrier_init(smem + 24, 2); + // do_empty: 1 barriers, init_count=1 + mbarrier_init(smem + 32, 1); + // lse_full: 1 barriers, init_count=1 + mbarrier_init(smem + 40, 1); + // lse_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 48, 256); + // dpsum_full: 1 barriers, init_count=1 + mbarrier_init(smem + 56, 1); + // dpsum_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 64, 256); + // s_full: 1 barriers, init_count=1 + mbarrier_init(smem + 72, 1); + // dp_full: 1 barriers, init_count=1 + mbarrier_init(smem + 80, 1); + // ds_full: 1 barriers, init_count=512 + mbarrier_init(smem + 88, 512); + // ds_empty: 1 barriers, init_count=2 + mbarrier_init(smem + 96, 2); + // xchg_full: 1 barriers, init_count=1 + mbarrier_init(smem + 104, 1); + // xchg_done: 1 barriers, init_count=2 + mbarrier_init(smem + 112, 2); + // dq_full: 1 barriers, init_count=1 + mbarrier_init(smem + 120, 1); + // dq_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 128, 256); + // dqlo_empty: 1 barriers, init_count=256 + mbarrier_init(smem + 136, 256); + // dkv_full: 1 barriers, init_count=1 + mbarrier_init(smem + 144, 1); + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + } + } + + __syncwarp(); + asm volatile("barrier.cluster.arrive.release.aligned;" ::: "memory"); + asm volatile("barrier.cluster.wait.acquire.aligned;" ::: "memory"); + + // TMEM alloc (512 columns, 512 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 152); + if (warp == 12) { + int _tmem_hold = smem + 152; + asm volatile("tcgen05.alloc.cta_group::2.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(512) : "memory"); + __syncwarp(); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::2.sync.aligned;"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;" ::: "memory"); + + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_dv = taddr; + const int tmem_dk = taddr + 128; + const int tmem_s = taddr + 256; + const int tmem_dq = taddr + 384; + const int tmem_dp = taddr + 320; + if (warp >= 12 && warp <= 15) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 88;"); + } + + // ---- Role: reduce ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 152;"); + { // reduce_main + int seg = blockIdx.z; + int n_blk = blockIdx.x >> 1; + int q0 = cu_seqlens_q[seg]; + int q_len = cu_seqlens_q[seg + 1] - q0; + int k0 = cu_seqlens_k[seg]; + int k_len = cu_seqlens_k[seg + 1] - k0; + int found = ((n_blk * 128 < k_len) ? 1 : 0); + int off = k_len - q_len; + int first_key = n_blk * 128 - off; + int m_min = ((first_key > 0) ? first_key : 0) >> 6; + int m_max = q_len + 63 >> 6; + int num_m = ((found != 0) ? m_max - m_min : 0); + int row_off = q0 * 2 + seg * 128 >> 7 << 7; + int rank = cta_rank & 1; + int h = blockIdx.y; + int n_groups = group >> 1; + const int quarter = warp % 4; + const int tl = quarter * 32 + lane; + const unsigned int lane_base = (unsigned int)(quarter * 32 << 16); + int t = 0; + #pragma unroll 1 + for (int g = 0; g < n_groups; g++) { + long long head_rows = (long long)((h * n_groups + g) * rows_pad + row_off); + #pragma unroll 1 + for (int mi = 0; mi < num_m; mi++) { + long long tile = (head_rows + (long long)((m_min + mi) * 128)) * 256 + (long long)(rank * 8 * 2048); + int _parity_dq_full = (int)((unsigned int)t % 2); + mbarrier_wait(dq_full_addr, _parity_dq_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + float dq_0[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_0[0]), "=f"(dq_0[1]), "=f"(dq_0[2]), "=f"(dq_0[3]), "=f"(dq_0[4]), "=f"(dq_0[5]), "=f"(dq_0[6]), "=f"(dq_0[7]), "=f"(dq_0[8]), "=f"(dq_0[9]), "=f"(dq_0[10]), "=f"(dq_0[11]), "=f"(dq_0[12]), "=f"(dq_0[13]), "=f"(dq_0[14]), "=f"(dq_0[15]), "=f"(dq_0[16]), "=f"(dq_0[17]), "=f"(dq_0[18]), "=f"(dq_0[19]), "=f"(dq_0[20]), "=f"(dq_0[21]), "=f"(dq_0[22]), "=f"(dq_0[23]), "=f"(dq_0[24]), "=f"(dq_0[25]), "=f"(dq_0[26]), "=f"(dq_0[27]), "=f"(dq_0[28]), "=f"(dq_0[29]), "=f"(dq_0[30]), "=f"(dq_0[31]) + : "r"(taddr + 448 + lane_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_0[32]), "=f"(dq_0[33]), "=f"(dq_0[34]), "=f"(dq_0[35]), "=f"(dq_0[36]), "=f"(dq_0[37]), "=f"(dq_0[38]), "=f"(dq_0[39]), "=f"(dq_0[40]), "=f"(dq_0[41]), "=f"(dq_0[42]), "=f"(dq_0[43]), "=f"(dq_0[44]), "=f"(dq_0[45]), "=f"(dq_0[46]), "=f"(dq_0[47]), "=f"(dq_0[48]), "=f"(dq_0[49]), "=f"(dq_0[50]), "=f"(dq_0[51]), "=f"(dq_0[52]), "=f"(dq_0[53]), "=f"(dq_0[54]), "=f"(dq_0[55]), "=f"(dq_0[56]), "=f"(dq_0[57]), "=f"(dq_0[58]), "=f"(dq_0[59]), "=f"(dq_0[60]), "=f"(dq_0[61]), "=f"(dq_0[62]), "=f"(dq_0[63]) + : "r"(taddr + 480 + lane_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((dq_empty_addr) & 0xFEFFFFFF) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_0[0]), "f"(dq_0[1]), "f"(dq_0[2]), "f"(dq_0[3]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_0[4]), "f"(dq_0[5]), "f"(dq_0[6]), "f"(dq_0[7]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_0[8]), "f"(dq_0[9]), "f"(dq_0[10]), "f"(dq_0[11]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_0[12]), "f"(dq_0[13]), "f"(dq_0[14]), "f"(dq_0[15]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_0 = reinterpret_cast(dq_accum + (tile + 8192)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_0), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_0[16]), "f"(dq_0[17]), "f"(dq_0[18]), "f"(dq_0[19]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_0[20]), "f"(dq_0[21]), "f"(dq_0[22]), "f"(dq_0[23]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_0[24]), "f"(dq_0[25]), "f"(dq_0[26]), "f"(dq_0[27]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_0[28]), "f"(dq_0[29]), "f"(dq_0[30]), "f"(dq_0[31]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_1 = reinterpret_cast(dq_accum + (tile + 10240)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_1), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_0[32]), "f"(dq_0[33]), "f"(dq_0[34]), "f"(dq_0[35]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_0[36]), "f"(dq_0[37]), "f"(dq_0[38]), "f"(dq_0[39]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_0[40]), "f"(dq_0[41]), "f"(dq_0[42]), "f"(dq_0[43]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_0[44]), "f"(dq_0[45]), "f"(dq_0[46]), "f"(dq_0[47]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_2 = reinterpret_cast(dq_accum + (tile + 12288)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_2), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_0[48]), "f"(dq_0[49]), "f"(dq_0[50]), "f"(dq_0[51]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_0[52]), "f"(dq_0[53]), "f"(dq_0[54]), "f"(dq_0[55]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_0[56]), "f"(dq_0[57]), "f"(dq_0[58]), "f"(dq_0[59]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_0[60]), "f"(dq_0[61]), "f"(dq_0[62]), "f"(dq_0[63]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_3 = reinterpret_cast(dq_accum + (tile + 14336)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_3), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + float dq_1[64]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_1[0]), "=f"(dq_1[1]), "=f"(dq_1[2]), "=f"(dq_1[3]), "=f"(dq_1[4]), "=f"(dq_1[5]), "=f"(dq_1[6]), "=f"(dq_1[7]), "=f"(dq_1[8]), "=f"(dq_1[9]), "=f"(dq_1[10]), "=f"(dq_1[11]), "=f"(dq_1[12]), "=f"(dq_1[13]), "=f"(dq_1[14]), "=f"(dq_1[15]), "=f"(dq_1[16]), "=f"(dq_1[17]), "=f"(dq_1[18]), "=f"(dq_1[19]), "=f"(dq_1[20]), "=f"(dq_1[21]), "=f"(dq_1[22]), "=f"(dq_1[23]), "=f"(dq_1[24]), "=f"(dq_1[25]), "=f"(dq_1[26]), "=f"(dq_1[27]), "=f"(dq_1[28]), "=f"(dq_1[29]), "=f"(dq_1[30]), "=f"(dq_1[31]) + : "r"(taddr + 384 + lane_base)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dq_1[32]), "=f"(dq_1[33]), "=f"(dq_1[34]), "=f"(dq_1[35]), "=f"(dq_1[36]), "=f"(dq_1[37]), "=f"(dq_1[38]), "=f"(dq_1[39]), "=f"(dq_1[40]), "=f"(dq_1[41]), "=f"(dq_1[42]), "=f"(dq_1[43]), "=f"(dq_1[44]), "=f"(dq_1[45]), "=f"(dq_1[46]), "=f"(dq_1[47]), "=f"(dq_1[48]), "=f"(dq_1[49]), "=f"(dq_1[50]), "=f"(dq_1[51]), "=f"(dq_1[52]), "=f"(dq_1[53]), "=f"(dq_1[54]), "=f"(dq_1[55]), "=f"(dq_1[56]), "=f"(dq_1[57]), "=f"(dq_1[58]), "=f"(dq_1[59]), "=f"(dq_1[60]), "=f"(dq_1[61]), "=f"(dq_1[62]), "=f"(dq_1[63]) + : "r"(taddr + 416 + lane_base)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((dqlo_empty_addr) & 0xFEFFFFFF) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_1[0]), "f"(dq_1[1]), "f"(dq_1[2]), "f"(dq_1[3]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_1[4]), "f"(dq_1[5]), "f"(dq_1[6]), "f"(dq_1[7]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_1[8]), "f"(dq_1[9]), "f"(dq_1[10]), "f"(dq_1[11]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_1[12]), "f"(dq_1[13]), "f"(dq_1[14]), "f"(dq_1[15]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_4 = reinterpret_cast(dq_accum + (tile + 0)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_4), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_1[16]), "f"(dq_1[17]), "f"(dq_1[18]), "f"(dq_1[19]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_1[20]), "f"(dq_1[21]), "f"(dq_1[22]), "f"(dq_1[23]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_1[24]), "f"(dq_1[25]), "f"(dq_1[26]), "f"(dq_1[27]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_1[28]), "f"(dq_1[29]), "f"(dq_1[30]), "f"(dq_1[31]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_5 = reinterpret_cast(dq_accum + (tile + 2048)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_5), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 0 + (unsigned int)(tl * 16)), "f"(dq_1[32]), "f"(dq_1[33]), "f"(dq_1[34]), "f"(dq_1[35]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 2048 + (unsigned int)(tl * 16)), "f"(dq_1[36]), "f"(dq_1[37]), "f"(dq_1[38]), "f"(dq_1[39]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 4096 + (unsigned int)(tl * 16)), "f"(dq_1[40]), "f"(dq_1[41]), "f"(dq_1[42]), "f"(dq_1[43]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 6144 + (unsigned int)(tl * 16)), "f"(dq_1[44]), "f"(dq_1[45]), "f"(dq_1[46]), "f"(dq_1[47]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_6 = reinterpret_cast(dq_accum + (tile + 4096)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_6), "r"(sdq_addr + 0), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 8192 + (unsigned int)(tl * 16)), "f"(dq_1[48]), "f"(dq_1[49]), "f"(dq_1[50]), "f"(dq_1[51]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 10240 + (unsigned int)(tl * 16)), "f"(dq_1[52]), "f"(dq_1[53]), "f"(dq_1[54]), "f"(dq_1[55]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 12288 + (unsigned int)(tl * 16)), "f"(dq_1[56]), "f"(dq_1[57]), "f"(dq_1[58]), "f"(dq_1[59]) : "memory"); + asm volatile("st.shared.v4.f32 [%0], {%1,%2,%3,%4};" :: "r"(sdq_addr + 14336 + (unsigned int)(tl * 16)), "f"(dq_1[60]), "f"(dq_1[61]), "f"(dq_1[62]), "f"(dq_1[63]) : "memory"); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 2, 128;" ::: "memory"); + if (warp == 0) { + if (elect_sync()) { + { + void* _cpred_dst_7 = reinterpret_cast(dq_accum + (tile + 6144)); + asm volatile("cp.reduce.async.bulk.global.shared::cta.bulk_group.add.f32" + " [%0], [%1], %2;" + :: "l"(_cpred_dst_7), "r"(sdq_addr + 8192), "r"((uint32_t)(8192)) + : "memory"); + } + asm volatile("cp.async.bulk.commit_group;"); + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 2, 128;" ::: "memory"); + t += 1; + } + } + if (warp == 0) { + if (elect_sync()) { + asm volatile("cp.async.bulk.wait_group.read 0;"); + } + } + } + } + // ---- Role: compute ---- + if (warp >= 4 && warp <= 11) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 136;"); + { // compute_main + int seg_1 = blockIdx.z; + int n_blk_1 = blockIdx.x >> 1; + int q0_1 = cu_seqlens_q[seg_1]; + int q_len_1 = cu_seqlens_q[seg_1 + 1] - q0_1; + int k0_1 = cu_seqlens_k[seg_1]; + int k_len_1 = cu_seqlens_k[seg_1 + 1] - k0_1; + int found_1 = ((n_blk_1 * 128 < k_len_1) ? 1 : 0); + int off_1 = k_len_1 - q_len_1; + int first_key_1 = n_blk_1 * 128 - off_1; + int m_min_1 = ((first_key_1 > 0) ? first_key_1 : 0) >> 6; + int m_max_1 = q_len_1 + 63 >> 6; + int num_m_1 = ((found_1 != 0) ? m_max_1 - m_min_1 : 0); + int row_off_1 = q0_1 * 2 + seg_1 * 128 >> 7 << 7; + int rank_1 = cta_rank & 1; + int h_1 = blockIdx.y; + int n_groups_1 = group >> 1; + const int quarter_1 = warp % 4; + const int wg = warp - 4 >> 2; + const int row = (quarter_1 & 1) * 32 + lane; + const int half = quarter_1 >> 1; + const int col0 = half * 64 + wg * 32; + const unsigned int lane_base_1 = (unsigned int)(quarter_1 * 32 << 16); + int key = n_blk_1 * 128 + rank_1 * 64 + row; + int prow = half * 64 + row; + int t_1 = 0; + #pragma unroll 1 + for (int g_1 = 0; g_1 < n_groups_1; g_1++) { + #pragma unroll 1 + for (int mi_1 = 0; mi_1 < num_m_1; mi_1++) { + int m = m_min_1 + mi_1; + int _parity_lse_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(lse_full_addr, _parity_lse_full); + int _parity_s_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(s_full_addr, _parity_s_full); + int _parity_ds_empty = (int)((unsigned int)t_1 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); + float s[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(s[0]), "=f"(s[1]), "=f"(s[2]), "=f"(s[3]), "=f"(s[4]), "=f"(s[5]), "=f"(s[6]), "=f"(s[7]), "=f"(s[8]), "=f"(s[9]), "=f"(s[10]), "=f"(s[11]), "=f"(s[12]), "=f"(s[13]), "=f"(s[14]), "=f"(s[15]), "=f"(s[16]), "=f"(s[17]), "=f"(s[18]), "=f"(s[19]), "=f"(s[20]), "=f"(s[21]), "=f"(s[22]), "=f"(s[23]), "=f"(s[24]), "=f"(s[25]), "=f"(s[26]), "=f"(s[27]), "=f"(s[28]), "=f"(s[29]), "=f"(s[30]), "=f"(s[31]) + : "r"(taddr + (unsigned int)(256 + wg * 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + int lo = (key - off_1 - m * 64) * 2 - col0; + int hi = (q_len_1 - m * 64) * 2 - col0; + int hi_k = ((key < k_len_1) ? hi : 0); + uint32_t _thresh_mask_0; + { + int _lim_1 = hi_k; + if (_lim_1 <= 0) { _thresh_mask_0 = 0u; } + else if (_lim_1 >= 32) { _thresh_mask_0 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_0) : "r"(_lim_1)); + } + } + uint32_t _thresh_mask_2; + { + int _lim_3 = lo; + if (_lim_3 <= 0) { _thresh_mask_2 = 0u; } + else if (_lim_3 >= 32) { _thresh_mask_2 = 0xFFFFFFFFu; } + else { + asm volatile("{" + ".reg .u32 t;\n\t" + "shl.b32 t, 1, %1;\n\t" + "add.u32 %0, t, -1;\n\t" + "}" : "=r"(_thresh_mask_2) : "r"(_lim_3)); + } + } + unsigned int keep = _thresh_mask_0 & ~_thresh_mask_2; + if (!(keep & (1u << 0))) s[0] = -CUDART_INF_F; + if (!(keep & (1u << 1))) s[1] = -CUDART_INF_F; + if (!(keep & (1u << 2))) s[2] = -CUDART_INF_F; + if (!(keep & (1u << 3))) s[3] = -CUDART_INF_F; + if (!(keep & (1u << 4))) s[4] = -CUDART_INF_F; + if (!(keep & (1u << 5))) s[5] = -CUDART_INF_F; + if (!(keep & (1u << 6))) s[6] = -CUDART_INF_F; + if (!(keep & (1u << 7))) s[7] = -CUDART_INF_F; + if (!(keep & (1u << 8))) s[8] = -CUDART_INF_F; + if (!(keep & (1u << 9))) s[9] = -CUDART_INF_F; + if (!(keep & (1u << 10))) s[10] = -CUDART_INF_F; + if (!(keep & (1u << 11))) s[11] = -CUDART_INF_F; + if (!(keep & (1u << 12))) s[12] = -CUDART_INF_F; + if (!(keep & (1u << 13))) s[13] = -CUDART_INF_F; + if (!(keep & (1u << 14))) s[14] = -CUDART_INF_F; + if (!(keep & (1u << 15))) s[15] = -CUDART_INF_F; + if (!(keep & (1u << 16))) s[16] = -CUDART_INF_F; + if (!(keep & (1u << 17))) s[17] = -CUDART_INF_F; + if (!(keep & (1u << 18))) s[18] = -CUDART_INF_F; + if (!(keep & (1u << 19))) s[19] = -CUDART_INF_F; + if (!(keep & (1u << 20))) s[20] = -CUDART_INF_F; + if (!(keep & (1u << 21))) s[21] = -CUDART_INF_F; + if (!(keep & (1u << 22))) s[22] = -CUDART_INF_F; + if (!(keep & (1u << 23))) s[23] = -CUDART_INF_F; + if (!(keep & (1u << 24))) s[24] = -CUDART_INF_F; + if (!(keep & (1u << 25))) s[25] = -CUDART_INF_F; + if (!(keep & (1u << 26))) s[26] = -CUDART_INF_F; + if (!(keep & (1u << 27))) s[27] = -CUDART_INF_F; + if (!(keep & (1u << 28))) s[28] = -CUDART_INF_F; + if (!(keep & (1u << 29))) s[29] = -CUDART_INF_F; + if (!(keep & (1u << 30))) s[30] = -CUDART_INF_F; + if (!(keep & (1u << 31))) s[31] = -CUDART_INF_F; + float lse[32]; + #pragma unroll + for (int j = 0; j < 32; j++) { + lse[j] = slse[col0 + j]; + } + float2 scale2 = make_float2(softmax_scale_log2, softmax_scale_log2); + #pragma unroll + for (int pair = 0; pair < 16; pair++) { + float2 x = fma_sub_f32x2_rn_ftz(make_float2(s[pair * 2], s[pair * 2 + 1]), scale2, make_float2(lse[pair * 2], lse[pair * 2 + 1])); + s[pair * 2] = x.x; + s[pair * 2 + 1] = x.y; + } + unsigned int p16[16]; + softmax_frag_exp2_cast(&s[0], p16, 0); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 0 ^ prow & 7) * 16))), "r"(p16[0]), "r"(p16[1]), "r"(p16[2]), "r"(p16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 1 ^ prow & 7) * 16))), "r"(p16[4]), "r"(p16[5]), "r"(p16[6]), "r"(p16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 2 ^ prow & 7) * 16))), "r"(p16[8]), "r"(p16[9]), "r"(p16[10]), "r"(p16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sp_addr + (unsigned int)(prow * 128 + (wg * 4 + 3 ^ prow & 7) * 16))), "r"(p16[12]), "r"(p16[13]), "r"(p16[14]), "r"(p16[15]) : "memory"); + mbarrier_arrive(lse_empty_addr); + int _parity_dpsum_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(dpsum_full_addr, _parity_dpsum_full); + int _parity_dp_full = (int)((unsigned int)t_1 % 2); + mbarrier_wait(dp_full_addr, _parity_dp_full); + asm volatile("tcgen05.fence::after_thread_sync;"); + float dp[32]; + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(dp[0]), "=f"(dp[1]), "=f"(dp[2]), "=f"(dp[3]), "=f"(dp[4]), "=f"(dp[5]), "=f"(dp[6]), "=f"(dp[7]), "=f"(dp[8]), "=f"(dp[9]), "=f"(dp[10]), "=f"(dp[11]), "=f"(dp[12]), "=f"(dp[13]), "=f"(dp[14]), "=f"(dp[15]), "=f"(dp[16]), "=f"(dp[17]), "=f"(dp[18]), "=f"(dp[19]), "=f"(dp[20]), "=f"(dp[21]), "=f"(dp[22]), "=f"(dp[23]), "=f"(dp[24]), "=f"(dp[25]), "=f"(dp[26]), "=f"(dp[27]), "=f"(dp[28]), "=f"(dp[29]), "=f"(dp[30]), "=f"(dp[31]) + : "r"(taddr + (unsigned int)(320 + wg * 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;"); + float dps[32]; + #pragma unroll + for (int j_1 = 0; j_1 < 32; j_1++) { + dps[j_1] = sdps[col0 + j_1]; + } + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + sub_f32x2_inplace(&reinterpret_cast(dp)[_ls], reinterpret_cast(dps)[_ls]); + #pragma unroll + for (int _ls = 0; _ls < 16; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(dp)[_ls], reinterpret_cast(s)[_ls]); + unsigned int ds16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dp[_lp*2 + 0], dp[_lp*2+1 + 0])); + ds16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 0 ^ prow & 7) * 16))), "r"(ds16[0]), "r"(ds16[1]), "r"(ds16[2]), "r"(ds16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 1 ^ prow & 7) * 16))), "r"(ds16[4]), "r"(ds16[5]), "r"(ds16[6]), "r"(ds16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 2 ^ prow & 7) * 16))), "r"(ds16[8]), "r"(ds16[9]), "r"(ds16[10]), "r"(ds16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sdst_addr + (unsigned int)(prow * 128 + (wg * 4 + 3 ^ prow & 7) * 16))), "r"(ds16[12]), "r"(ds16[13]), "r"(ds16[14]), "r"(ds16[15]) : "memory"); + if (half == rank_1) { + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 0 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[0]), "r"(ds16[1]), "r"(ds16[2]), "r"(ds16[3]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 1 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[4]), "r"(ds16[5]), "r"(ds16[6]), "r"(ds16[7]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 2 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[8]), "r"(ds16[9]), "r"(ds16[10]), "r"(ds16[11]) : "memory"); + asm volatile("st.shared.v4.b32 [%0], {%1,%2,%3,%4};" :: "r"((sds_addr + (unsigned int)((rank_1 * 64 + row) * 128 + (wg * 4 + 3 ^ rank_1 * 64 + row & 7) * 16))), "r"(ds16[12]), "r"(ds16[13]), "r"(ds16[14]), "r"(ds16[15]) : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("tcgen05.fence::before_thread_sync;"); + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((ds_full_addr) & 0xFEFFFFFF) : "memory"); + mbarrier_arrive(dpsum_empty_addr); + asm volatile("barrier.sync 1, 256;" ::: "memory"); + if (warp == 4) { + if (elect_sync()) { + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.expect_tx.relaxed.cluster.shared::cluster.b64 _, [remAddr32], %2;\n\t" + "}" + :: "r"(xchg_full_addr), "r"(1 - rank_1), "r"((uint32_t)(8192)) : "memory"); + uint32_t _mbar_peer_4; + asm( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_mbar_peer_4) : "r"(xchg_full_addr), "r"(1 - rank_1)); + uint32_t _cluster_mapa_5; + asm volatile( + "mapa.shared::cluster.u32 %0, %1, %2;" + : "=r"(_cluster_mapa_5) : "r"(sds_addr + (unsigned int)(rank_1 * 8192)), "r"(1 - rank_1)); + asm volatile( + "cp.async.bulk.shared::cluster.shared::cta.mbarrier::complete_tx::bytes" + " [%0], [%1], %2, [%3];" + :: "r"(_cluster_mapa_5), "r"(sdst_addr + (unsigned int)((1 - rank_1) * 8192)), "r"((uint32_t)(8192)), "r"(_mbar_peer_4) + : "memory"); + } + } + t_1 += 1; + } + } + int _parity_ds_empty_1 = (int)((unsigned int)t_1 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty_1, 10000000); + mbarrier_wait(dkv_full_addr, 0); + asm volatile("tcgen05.fence::after_thread_sync;"); + int hd0 = half * 128 + wg * 64; + int store_row = ((key < k_len_1) ? 1 : 0); + if (found_1 == 0) { + store_row = 0; + } + long long out_row = (long long)(k0_1 + key); + float acc_0[64]; + if (num_m_1 > 0) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_0[0]), "=f"(acc_0[1]), "=f"(acc_0[2]), "=f"(acc_0[3]), "=f"(acc_0[4]), "=f"(acc_0[5]), "=f"(acc_0[6]), "=f"(acc_0[7]), "=f"(acc_0[8]), "=f"(acc_0[9]), "=f"(acc_0[10]), "=f"(acc_0[11]), "=f"(acc_0[12]), "=f"(acc_0[13]), "=f"(acc_0[14]), "=f"(acc_0[15]), "=f"(acc_0[16]), "=f"(acc_0[17]), "=f"(acc_0[18]), "=f"(acc_0[19]), "=f"(acc_0[20]), "=f"(acc_0[21]), "=f"(acc_0[22]), "=f"(acc_0[23]), "=f"(acc_0[24]), "=f"(acc_0[25]), "=f"(acc_0[26]), "=f"(acc_0[27]), "=f"(acc_0[28]), "=f"(acc_0[29]), "=f"(acc_0[30]), "=f"(acc_0[31]) + : "r"(taddr + (unsigned int)(0 + wg * 64) + lane_base_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_0[32]), "=f"(acc_0[33]), "=f"(acc_0[34]), "=f"(acc_0[35]), "=f"(acc_0[36]), "=f"(acc_0[37]), "=f"(acc_0[38]), "=f"(acc_0[39]), "=f"(acc_0[40]), "=f"(acc_0[41]), "=f"(acc_0[42]), "=f"(acc_0[43]), "=f"(acc_0[44]), "=f"(acc_0[45]), "=f"(acc_0[46]), "=f"(acc_0[47]), "=f"(acc_0[48]), "=f"(acc_0[49]), "=f"(acc_0[50]), "=f"(acc_0[51]), "=f"(acc_0[52]), "=f"(acc_0[53]), "=f"(acc_0[54]), "=f"(acc_0[55]), "=f"(acc_0[56]), "=f"(acc_0[57]), "=f"(acc_0[58]), "=f"(acc_0[59]), "=f"(acc_0[60]), "=f"(acc_0[61]), "=f"(acc_0[62]), "=f"(acc_0[63]) + : "r"(taddr + (unsigned int)(0 + wg * 64 + 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } else { + #pragma unroll + for (int i_0 = 0; i_0 < 64; i_0++) { + acc_0[i_0] = 0.0f; + } + } + long long base_0 = (out_row * (long long)h_kv + (long long)h_1) * 256 + (long long)hd0; + if (store_row != 0) { + #pragma unroll + for (int i_1 = 0; i_1 < 16; i_1++) { + { + float4 _v4 = make_float4(acc_0[i_1 * 4 + 0], acc_0[i_1 * 4 + 1], acc_0[i_1 * 4 + 2], acc_0[i_1 * 4 + 3]); + *reinterpret_cast((dv_out + (base_0 + (long long)(i_1 * 4))) + 0) = _v4; + } + } + } + float acc_1[64]; + if (num_m_1 > 0) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_1[0]), "=f"(acc_1[1]), "=f"(acc_1[2]), "=f"(acc_1[3]), "=f"(acc_1[4]), "=f"(acc_1[5]), "=f"(acc_1[6]), "=f"(acc_1[7]), "=f"(acc_1[8]), "=f"(acc_1[9]), "=f"(acc_1[10]), "=f"(acc_1[11]), "=f"(acc_1[12]), "=f"(acc_1[13]), "=f"(acc_1[14]), "=f"(acc_1[15]), "=f"(acc_1[16]), "=f"(acc_1[17]), "=f"(acc_1[18]), "=f"(acc_1[19]), "=f"(acc_1[20]), "=f"(acc_1[21]), "=f"(acc_1[22]), "=f"(acc_1[23]), "=f"(acc_1[24]), "=f"(acc_1[25]), "=f"(acc_1[26]), "=f"(acc_1[27]), "=f"(acc_1[28]), "=f"(acc_1[29]), "=f"(acc_1[30]), "=f"(acc_1[31]) + : "r"(taddr + (unsigned int)(128 + wg * 64) + lane_base_1)); + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x32.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, %30, %31}, [%32];" + : "=f"(acc_1[32]), "=f"(acc_1[33]), "=f"(acc_1[34]), "=f"(acc_1[35]), "=f"(acc_1[36]), "=f"(acc_1[37]), "=f"(acc_1[38]), "=f"(acc_1[39]), "=f"(acc_1[40]), "=f"(acc_1[41]), "=f"(acc_1[42]), "=f"(acc_1[43]), "=f"(acc_1[44]), "=f"(acc_1[45]), "=f"(acc_1[46]), "=f"(acc_1[47]), "=f"(acc_1[48]), "=f"(acc_1[49]), "=f"(acc_1[50]), "=f"(acc_1[51]), "=f"(acc_1[52]), "=f"(acc_1[53]), "=f"(acc_1[54]), "=f"(acc_1[55]), "=f"(acc_1[56]), "=f"(acc_1[57]), "=f"(acc_1[58]), "=f"(acc_1[59]), "=f"(acc_1[60]), "=f"(acc_1[61]), "=f"(acc_1[62]), "=f"(acc_1[63]) + : "r"(taddr + (unsigned int)(128 + wg * 64 + 32) + lane_base_1)); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + } else { + #pragma unroll + for (int i_2 = 0; i_2 < 64; i_2++) { + acc_1[i_2] = 0.0f; + } + } + const float2 _scale2_6 = {softmax_scale, softmax_scale}; + #pragma unroll + for (int _ls = 0; _ls < 32; _ls++) + mul_f32x2_noftz_inplace(&reinterpret_cast(acc_1)[_ls], _scale2_6); + long long base_1 = (out_row * (long long)h_kv + (long long)h_1) * 256 + (long long)hd0; + if (store_row != 0) { + #pragma unroll + for (int i_3 = 0; i_3 < 16; i_3++) { + { + float4 _v4 = make_float4(acc_1[i_3 * 4 + 0], acc_1[i_3 * 4 + 1], acc_1[i_3 * 4 + 2], acc_1[i_3 * 4 + 3]); + *reinterpret_cast((dk_out + (base_1 + (long long)(i_3 * 4))) + 0) = _v4; + } + } + } + } + } + // ---- Role: mma ---- + if (warp == 12) { + { // mma_main + int seg_2 = blockIdx.z; + int n_blk_2 = blockIdx.x >> 1; + int q0_2 = cu_seqlens_q[seg_2]; + int q_len_2 = cu_seqlens_q[seg_2 + 1] - q0_2; + int k0_2 = cu_seqlens_k[seg_2]; + int k_len_2 = cu_seqlens_k[seg_2 + 1] - k0_2; + int found_2 = ((n_blk_2 * 128 < k_len_2) ? 1 : 0); + int off_2 = k_len_2 - q_len_2; + int first_key_2 = n_blk_2 * 128 - off_2; + int m_min_2 = ((first_key_2 > 0) ? first_key_2 : 0) >> 6; + int m_max_2 = q_len_2 + 63 >> 6; + int num_m_2 = ((found_2 != 0) ? m_max_2 - m_min_2 : 0); + int row_off_2 = q0_2 * 2 + seg_2 * 128 >> 7 << 7; + int n_groups_2 = group >> 1; + int t_2 = 0; + unsigned int _phase_q_full = 0; + unsigned int _phase_do_full = 0; + if ((cta_rank & 1) == 0) { + if (elect_sync()) { + mbarrier_wait_cluster_hint(kv_full_addr, 0, 10000000); + #pragma unroll 1 + for (int g_2 = 0; g_2 < n_groups_2; g_2++) { + #pragma unroll 1 + for (int mi_2 = 0; mi_2 < num_m_2; mi_2++) { + int first = ((t_2 == 0) ? 1 : 0); + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_0 = ((sk_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_1 = ((sq_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_0 = ((sk_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_1 = ((sq_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_0), "r"(_smem_desc_lo_1), "r"(tmem_s), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(s_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + mbarrier_wait_cluster_hint(do_full_addr, _phase_do_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_2 = ((sv_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_3 = ((sdo_addr) >> 4) & 0x7FFF; +#else + int _smem_desc_lo_2 = ((sv_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_3 = ((sdo_addr) >> 4) & 0x3FFF; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 136316048;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 506;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_2), "r"(_smem_desc_lo_3), "r"(tmem_dp), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(dp_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(do_empty_addr, (uint16_t)(3)); + _phase_do_full ^= 1; + mbarrier_wait_cluster_hint(q_full_addr, _phase_q_full, 10000000); + int _parity_ds_full = (int)((unsigned int)t_2 % 2); + mbarrier_wait_cluster_hint(ds_full_addr, _parity_ds_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_4 = ((sdst_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_5 = (((sqt_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_4 = ((sdst_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_5 = (((sqt_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_4), "r"(_smem_desc_lo_5), "r"(tmem_dk), "r"(((first) ? 0 : 1)) : "memory"); + tcgen05_commit_cg2_multicast(q_empty_addr, (uint16_t)(3)); + _phase_q_full ^= 1; + mbarrier_wait_cluster_hint(do_full_addr, _phase_do_full, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_6 = ((sp_addr) >> 4) & 0x7FFF; + int _smem_desc_lo_7 = (((sdot_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_6 = ((sp_addr) >> 4) & 0x3FFF; + int _smem_desc_lo_7 = (((sdot_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138478736;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 506;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 2;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_6), "r"(_smem_desc_lo_7), "r"(tmem_dv), "r"(((first) ? 0 : 1)) : "memory"); + tcgen05_commit_cg2_multicast(do_empty_addr, (uint16_t)(3)); + _phase_do_full ^= 1; + int _parity_xchg_done = (int)((unsigned int)t_2 % 2); + mbarrier_wait_cluster_hint(xchg_done_addr, _parity_xchg_done, 10000000); + int _parity_ds_empty_2 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(ds_empty_addr, _parity_ds_empty_2, 10000000); + int _parity_dq_empty = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dq_empty_addr, _parity_dq_empty, 10000000); + int _parity_dqlo_empty = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dqlo_empty_addr, _parity_dqlo_empty, 10000000); + asm volatile("tcgen05.fence::after_thread_sync;"); +#if __CUDA_ARCH__ == 1070 + int _smem_desc_lo_8 = (((sds_addr) >> 4) & 0x7FFF) | 0x4000000; + int _smem_desc_lo_9 = (((skt_addr) >> 4) & 0x7FFF) | 0x4000000; +#else + int _smem_desc_lo_8 = (((sds_addr) >> 4) & 0x3FFF) | 0x4000000; + int _smem_desc_lo_9 = (((skt_addr) >> 4) & 0x3FFF) | 0x4000000; +#endif + asm volatile( + "{\n\t" + ".reg .pred p0, p1;\n\t" + ".reg .b32 adhi, bdhi, alo, blo, id, m0, m1, m2, m3, m4, m5, m6, m7;\n\t" + ".reg .b64 da, db;\n\t" + "" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "mov.b32 m0, 0; mov.b32 m1, 0; mov.b32 m2, 0; mov.b32 m3, 0;\n\tmov.b32 m4, 0; mov.b32 m5, 0; mov.b32 m6, 0; mov.b32 m7, 0;\n\t" + "mov.b32 adhi, 0x40004040;\n\t" + "mov.b32 bdhi, 0x40004040;\n\t" + "mov.b32 id, 138511504;\n\t" + "mov.b32 alo, %0;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p0;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "add.u32 alo, alo, 128;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 da, {alo, adhi};\n\t" + "mov.b64 db, {blo, bdhi};\n\t" + "tcgen05.mma.cta_group::2.kind::f16 [%2], da, db, id, {m0, m1, m2, m3, m4, m5, m6, m7}, p1;\n\t" + "}\n" + :: "r"(_smem_desc_lo_8), "r"(_smem_desc_lo_9), "r"(tmem_dq), "r"(0) : "memory"); + tcgen05_commit_cg2_multicast(dq_full_addr, (uint16_t)(3)); + tcgen05_commit_cg2_multicast(ds_empty_addr, (uint16_t)(3)); + t_2 += 1; + } + } + tcgen05_commit_cg2_multicast(dkv_full_addr, (uint16_t)(3)); + int _parity_dq_empty_1 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dq_empty_addr, _parity_dq_empty_1, 10000000); + int _parity_dqlo_empty_1 = (int)((unsigned int)t_2 % 2) ^ 1; + mbarrier_wait_cluster_hint(dqlo_empty_addr, _parity_dqlo_empty_1, 10000000); + } + } + } + } + // ---- Role: load ---- + if (warp == 13) { + { // load_main + int seg_3 = blockIdx.z; + int n_blk_3 = blockIdx.x >> 1; + int q0_3 = cu_seqlens_q[seg_3]; + int q_len_3 = cu_seqlens_q[seg_3 + 1] - q0_3; + int k0_3 = cu_seqlens_k[seg_3]; + int k_len_3 = cu_seqlens_k[seg_3 + 1] - k0_3; + int found_3 = ((n_blk_3 * 128 < k_len_3) ? 1 : 0); + int off_3 = k_len_3 - q_len_3; + int first_key_3 = n_blk_3 * 128 - off_3; + int m_min_3 = ((first_key_3 > 0) ? first_key_3 : 0) >> 6; + int m_max_3 = q_len_3 + 63 >> 6; + int num_m_3 = ((found_3 != 0) ? m_max_3 - m_min_3 : 0); + int row_off_3 = q0_3 * 2 + seg_3 * 128 >> 7 << 7; + int rank_2 = cta_rank & 1; + int h_2 = blockIdx.y; + int kv_row0 = k0_3 + n_blk_3 * 128; + int n_groups_3 = group >> 1; + int t_3 = 0; + unsigned int _phase_q_empty = 1; + unsigned int _phase_do_empty = 1; + if (elect_sync()) { + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((kv_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(98304)) : "memory"); + tma_3d_gmem2smem_cta2(sk_addr + 0, (&K), 0, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 0, (&V), 0, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 8192, (&K), 64, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 8192, (&V), 64, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 16384, (&K), 128, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 16384, (&V), 128, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sk_addr + 24576, (&K), 192, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sv_addr + 24576, (&V), 192, h_2, kv_row0 + rank_2 * 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 0, (&K), rank_2 * 128 + 0, h_2, kv_row0 + 0, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 8192, (&K), rank_2 * 128 + 0, h_2, kv_row0 + 64, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 16384, (&K), rank_2 * 128 + 64, h_2, kv_row0 + 0, ((kv_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(skt_addr + 24576, (&K), rank_2 * 128 + 64, h_2, kv_row0 + 64, ((kv_full_addr) & 0xFEFFFFFF)); + #pragma unroll 1 + for (int g_3 = 0; g_3 < n_groups_3; g_3++) { + int head0 = h_2 * group + g_3 * 2; + int srow = (h_2 * n_groups_3 + g_3) * rows_pad + row_off_3; + #pragma unroll 1 + for (int mi_3 = 0; mi_3 < num_m_3; mi_3++) { + int tok = q0_3 + (m_min_3 + mi_3) * 64; + int stat_row = srow + (m_min_3 + mi_3) * 128; + mbarrier_wait(q_empty_addr, _phase_q_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sq_addr + 0, (&Q), 0, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 8192, (&Q), 64, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 16384, (&Q), 128, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sq_addr + 24576, (&Q), 192, head0, tok + rank_2 * 32, ((q_full_addr) & 0xFEFFFFFF)); + _phase_q_empty ^= 1; + int _parity_lse_empty = (int)((unsigned int)t_3 % 2) ^ 1; + mbarrier_wait(lse_empty_addr, _parity_lse_empty); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive_expect_tx(lse_full_addr, 512); + cp_async_bulk_gmem2smem(slse_addr, reinterpret_cast(reinterpret_cast(lse_log2) + ((unsigned long long)stat_row * 4)), 512, lse_full_addr); + mbarrier_wait(do_empty_addr, _phase_do_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((do_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sdo_addr + 0, (&dO), 0, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 8192, (&dO), 64, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 16384, (&dO), 128, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdo_addr + 24576, (&dO), 192, head0, tok + rank_2 * 32, ((do_full_addr) & 0xFEFFFFFF)); + _phase_do_empty ^= 1; + int _parity_dpsum_empty = (int)((unsigned int)t_3 % 2) ^ 1; + mbarrier_wait(dpsum_empty_addr, _parity_dpsum_empty); + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive_expect_tx(dpsum_full_addr, 512); + cp_async_bulk_gmem2smem(sdps_addr, reinterpret_cast(reinterpret_cast(dpsum) + ((unsigned long long)stat_row * 4)), 512, dpsum_full_addr); + mbarrier_wait(q_empty_addr, _phase_q_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((q_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sqt_addr + 0, (&Q), rank_2 * 128 + 0, head0, tok + 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 8192, (&Q), rank_2 * 128 + 0, head0, tok + 32, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 16384, (&Q), rank_2 * 128 + 64, head0, tok + 0, ((q_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sqt_addr + 24576, (&Q), rank_2 * 128 + 64, head0, tok + 32, ((q_full_addr) & 0xFEFFFFFF)); + _phase_q_empty ^= 1; + mbarrier_wait(do_empty_addr, _phase_do_empty); + asm volatile( + "mbarrier.arrive.expect_tx.release.cta.shared::cluster.b64 _, [%0], %1;" + :: "r"((do_full_addr) & 0xFEFFFFFF), "r"((uint32_t)(32768)) : "memory"); + tma_3d_gmem2smem_cta2(sdot_addr + 0, (&dO), rank_2 * 128 + 0, head0, tok + 0, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 8192, (&dO), rank_2 * 128 + 0, head0, tok + 32, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 16384, (&dO), rank_2 * 128 + 64, head0, tok + 0, ((do_full_addr) & 0xFEFFFFFF)); + tma_3d_gmem2smem_cta2(sdot_addr + 24576, (&dO), rank_2 * 128 + 64, head0, tok + 32, ((do_full_addr) & 0xFEFFFFFF)); + _phase_do_empty ^= 1; + t_3 += 1; + } + } + mbarrier_wait(q_empty_addr, _phase_q_empty); + mbarrier_wait(do_empty_addr, _phase_do_empty); + } + } + } + // ---- Role: relay ---- + if (warp == 14) { + { // relay_main + int seg_4 = blockIdx.z; + int n_blk_4 = blockIdx.x >> 1; + int q0_4 = cu_seqlens_q[seg_4]; + int q_len_4 = cu_seqlens_q[seg_4 + 1] - q0_4; + int k0_4 = cu_seqlens_k[seg_4]; + int k_len_4 = cu_seqlens_k[seg_4 + 1] - k0_4; + int found_4 = ((n_blk_4 * 128 < k_len_4) ? 1 : 0); + int off_4 = k_len_4 - q_len_4; + int first_key_4 = n_blk_4 * 128 - off_4; + int m_min_4 = ((first_key_4 > 0) ? first_key_4 : 0) >> 6; + int m_max_4 = q_len_4 + 63 >> 6; + int num_m_4 = ((found_4 != 0) ? m_max_4 - m_min_4 : 0); + int row_off_4 = q0_4 * 2 + seg_4 * 128 >> 7 << 7; + int rank_3 = cta_rank & 1; + int n_groups_4 = group >> 1; + int t_4 = 0; + #pragma unroll 1 + for (int g_4 = 0; g_4 < n_groups_4; g_4++) { + #pragma unroll 1 + for (int mi_4 = 0; mi_4 < num_m_4; mi_4++) { + int _parity_xchg_full = (int)((unsigned int)t_4 % 2); + mbarrier_wait_cluster_hint(xchg_full_addr, _parity_xchg_full, 10000000); + if (elect_sync()) { + asm volatile("fence.release.sync_restrict::shared::cta.cluster;" ::: "memory"); + asm volatile( + "mbarrier.arrive.relaxed.cluster.shared::cluster.b64 _, [%0];" + :: "r"((xchg_done_addr) & 0xFEFFFFFF) : "memory"); + asm volatile( + "{\n\t" + ".reg .b32 remAddr32;\n\t" + "mapa.shared::cluster.u32 remAddr32, %0, %1;\n\t" + "mbarrier.arrive.release.cluster.shared::cluster.b64 _, [remAddr32];\n\t" + "}" + :: "r"(ds_empty_addr), "r"(1 - rank_3) : "memory"); + } + t_4 += 1; + } + } + } + } + + // Cleanup + asm volatile("tcgen05.fence::before_thread_sync;" ::: "memory"); + asm volatile("barrier.cluster.arrive.release.aligned;"); + asm volatile("barrier.cluster.wait.acquire.aligned;"); + + if (warp == 12) { + asm volatile("tcgen05.dealloc.cta_group::2.sync.aligned.b32 %0, %1;" :: "r"(tmem_addr_storage[0]), "r"(512)); + } +} + +} // extern "C" diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu new file mode 100644 index 0000000..8cb0156 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu @@ -0,0 +1,100 @@ +/* + * 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 Cake: the tvm-ffi host binding of kernel_cake_gqa_train_76decc9a5d22dc91ddd1, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_gqa_train_76decc9a5d22dc91ddd1(float* __restrict__ dQ_accum, int* __restrict__ cu_seqlens_q, __nv_bfloat16* __restrict__ dQ, int T_q, int H_q, int num_segments, int rows_pad, float softmax_scale); + +namespace cake_binding_kernel_cake_gqa_train_76decc9a5d22dc91ddd1 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gqa_train_76decc9a5d22dc91ddd1"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void Run(TensorView a_dQ_accum, TensorView a_cu_seqlens_q, TensorView a_dQ, int64_t a_T_q, int64_t a_H_q, int64_t a_num_segments, int64_t a_rows_pad, double a_softmax_scale, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_dQ_accum.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_dQ_accum, "dQ_accum"); + check_dtype(a_dQ_accum, DLDataType{kDLFloat, 32, 1}, "dQ_accum"); + check_cuda_tensor(a_cu_seqlens_q, "cu_seqlens_q"); + check_dtype(a_cu_seqlens_q, DLDataType{kDLInt, 32, 1}, "cu_seqlens_q"); + check_cuda_tensor(a_dQ, "dQ"); + check_dtype(a_dQ, DLDataType{kDLBfloat, 16, 1}, "dQ"); + TVM_FFI_CHECK(a_T_q >= -2147483648LL && a_T_q <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T_q' = " << a_T_q + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_H_q >= -2147483648LL && a_H_q <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'H_q' = " << a_H_q + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_num_segments >= -2147483648LL && a_num_segments <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_segments' = " << a_num_segments + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_pad >= -2147483648LL && a_rows_pad <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_pad' = " << a_rows_pad + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_dQ_accum, "dQ_accum"); + check_contiguous(a_cu_seqlens_q, "cu_seqlens_q"); + check_same_device(a_cu_seqlens_q, a_dQ_accum, "cu_seqlens_q", "dQ_accum"); + check_contiguous(a_dQ, "dQ"); + check_same_device(a_dQ, a_dQ_accum, "dQ", "dQ_accum"); + TVM_FFI_CHECK((a_T_q >= 1LL), ValueError) + << KERNEL_NAME << ": requires (T_q >= 1)"; + TVM_FFI_CHECK((a_H_q >= 2LL), ValueError) + << KERNEL_NAME << ": requires (H_q >= 2)"; + TVM_FFI_CHECK((a_num_segments >= 1LL), ValueError) + << KERNEL_NAME << ": requires (num_segments >= 1)"; + TVM_FFI_CHECK((a_rows_pad >= 128LL), ValueError) + << KERNEL_NAME << ": requires (rows_pad >= 128)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_dQ_accum = DataPtr(a_dQ_accum); + void* v_cu_seqlens_q = DataPtr(a_cu_seqlens_q); + void* v_dQ = DataPtr(a_dQ); + int32_t v_T_q = static_cast(a_T_q); + int32_t v_H_q = static_cast(a_H_q); + int32_t v_num_segments = static_cast(a_num_segments); + int32_t v_rows_pad = static_cast(a_rows_pad); + float v_softmax_scale = static_cast(a_softmax_scale); + void* kernel_args[] = {&v_dQ_accum, &v_cu_seqlens_q, &v_dQ, &v_T_q, &v_H_q, &v_num_segments, &v_rows_pad, &v_softmax_scale}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(256, 1, 1); + config.dynamicSmemBytes = 0; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gqa_train_76decc9a5d22dc91ddd1), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gqa_train_76decc9a5d22dc91ddd1 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gqa_train_76decc9a5d22dc91ddd1::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu new file mode 100644 index 0000000..00f1909 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu @@ -0,0 +1,176 @@ +/* + * 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 256 + +#include + + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + +__device__ __forceinline__ unsigned int __as_u32(float v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "f"(v)); + return u; +} +__device__ __forceinline__ unsigned int __as_u32(__nv_bfloat162 v) { + return *reinterpret_cast(&v); +} +__device__ __forceinline__ unsigned int __as_u32(unsigned int v) { return v; } +__device__ __forceinline__ unsigned int __as_u32(int v) { + unsigned int u; + asm("mov.b32 %0, %1;" : "=r"(u) : "r"(v)); + return u; +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_cake_gqa_train_76decc9a5d22dc91ddd1(float* __restrict__ dQ_accum, int* __restrict__ cu_seqlens_q, __nv_bfloat16* __restrict__ dQ, int T_q, int H_q, int num_segments, int rows_pad, float softmax_scale) +{ + 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 (declaration order) === + int pair = blockIdx.y; + int row0 = blockIdx.x * 64; + int warp_row = row0 % 128 + warp * 8; + unsigned long long tile_base = ((unsigned long long)pair * (unsigned long long)rows_pad + (unsigned long long)(row0 - row0 % 128)) * 256; + asm volatile("griddepcontrol.wait;" ::: "memory"); + int limit = (row0 / 128 + 1) * 128; + int started_count = 0; + #pragma unroll 1 + for (int scan_base = 0; scan_base < num_segments; scan_base += 32) { + int b = scan_base + lane; + int _min_0 = ((b) < (num_segments - 1) ? (b) : (num_segments - 1)); + int cu_b = cu_seqlens_q[_min_0]; + unsigned int started = __ballot_sync(0xFFFFFFFF, b < num_segments && 2 * cu_b + 128 * b < limit); + started_count = started_count + (int)(unsigned int)__popc(started); + } + int seg = started_count - 1; + int seg_start = cu_seqlens_q[seg]; + int seg_len = cu_seqlens_q[seg + 1] - seg_start; + int local0 = row0 - (2 * seg_start + 128 * seg) / 128 * 128; + int in_window = ((local0 < (2 * seg_len + 128 - 1) / 128 * 128) ? 1 : 0); + int token_local0 = local0 + warp * 8 >> 1; + int tokens_left = ((in_window != 0) ? seg_len - token_local0 : 0); + float acc[64]; + int span_row = lane % 4; + int span_group = lane / 4; + #pragma unroll + for (int tb = 0; tb < 1; tb++) { + int row = warp_row + 2 * (tb * 4 + span_row); + #pragma unroll + for (int gb = 0; gb < 4; gb++) { + #pragma unroll + for (int hc = 0; hc < 2; hc++) { + { + unsigned _ldv8_1_0; + unsigned _ldv8_1_1; + unsigned _ldv8_1_2; + unsigned _ldv8_1_3; + unsigned _ldv8_1_4; + unsigned _ldv8_1_5; + unsigned _ldv8_1_6; + unsigned _ldv8_1_7; + asm volatile( + "ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(_ldv8_1_0), "=r"(_ldv8_1_1), "=r"(_ldv8_1_2), "=r"(_ldv8_1_3), "=r"(_ldv8_1_4), "=r"(_ldv8_1_5), "=r"(_ldv8_1_6), "=r"(_ldv8_1_7) : "l"((const void*)(dQ_accum + (tile_base + (unsigned long long)(row / 64 * 16384 + row % 64 * 4 + hc * 256 + (gb * 8 + span_group) * 512)) + (0))) : "memory"); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 0] = __uint_as_float(_ldv8_1_0); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 1] = __uint_as_float(_ldv8_1_1); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 2] = __uint_as_float(_ldv8_1_2); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 3] = __uint_as_float(_ldv8_1_3); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 4] = __uint_as_float(_ldv8_1_4); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 5] = __uint_as_float(_ldv8_1_5); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 6] = __uint_as_float(_ldv8_1_6); + acc[((tb * 4 + gb) * 2 + hc) * 8 + 7] = __uint_as_float(_ldv8_1_7); + } + } + } + } + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + float2 scale2 = make_float2(softmax_scale, softmax_scale); + #pragma unroll + for (int tb_1 = 0; tb_1 < 1; tb_1++) { + int tk = tb_1 * 4 + span_row; + int token = ((seg_start + token_local0 + tk) < (T_q - 1) ? (seg_start + token_local0 + tk) : (T_q - 1)); + #pragma unroll + for (int gb_1 = 0; gb_1 < 4; gb_1++) { + #pragma unroll + for (int hc_1 = 0; hc_1 < 2; hc_1++) { + unsigned int out[4]; + #pragma unroll + for (int j = 0; j < 4; j++) { + int a = ((tb_1 * 4 + gb_1) * 2 + hc_1) * 8 + 2 * j; + float2 _f32x2_mul_a_2 = make_float2(acc[a], acc[a + 1]); + float2 scaled; + asm("mul.rn.f32x2 %0, %1, %2;" : "=l"(*(unsigned long long*)&scaled) : "l"(*(const unsigned long long*)&_f32x2_mul_a_2), "l"(*(const unsigned long long*)&scale2)); + out[j] = __as_u32(__float22bfloat162_rn(scaled)); + } + if (tk < tokens_left) { + unsigned long long dst = ((unsigned long long)token * (unsigned long long)H_q + (unsigned long long)(2 * pair)) * 256 + (unsigned long long)(hc_1 * 128 + (gb_1 * 8 + span_group) * 4); + reinterpret_cast(reinterpret_cast(dQ + dst))[0] = reinterpret_cast(out)[0]; + reinterpret_cast(reinterpret_cast(dQ + (dst + 256)))[0] = reinterpret_cast(out + (2))[0]; + } + } + } + } +} + +} // extern "C" diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu new file mode 100644 index 0000000..7c8ea54 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu @@ -0,0 +1,119 @@ +/* + * 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 Cake: the tvm-ffi host binding of kernel_cake_gqa_train_a91752e8a2732bf2e246, for the device source that defines it. Do not edit. +#include +#include + +#include "cake_ffi_checks.h" + +#include + +extern "C" __global__ void kernel_cake_gqa_train_a91752e8a2732bf2e246(__nv_bfloat16* __restrict__ O, __nv_bfloat16* __restrict__ dO, float* __restrict__ LSE, int* __restrict__ cu_seqlens_q, float* __restrict__ LSE_log2, float* __restrict__ dPsum, float* __restrict__ dQ_accum, int T_q, int H_q, int num_segments, int rows_pad); + +namespace cake_binding_kernel_cake_gqa_train_a91752e8a2732bf2e246 { + +using tvm::ffi::TensorView; + +constexpr const char* KERNEL_NAME = "kernel_cake_gqa_train_a91752e8a2732bf2e246"; + +void CheckCuda(cudaError_t status, const char* what) { + TVM_FFI_CHECK(status == cudaSuccess, RuntimeError) + << KERNEL_NAME << ": " << what << " failed: " << cudaGetErrorString(status); +} + +void* DataPtr(const TensorView& t) { return static_cast(t.data_ptr()) + t.byte_offset(); } + +void Run(TensorView a_O, TensorView a_dO, TensorView a_LSE, TensorView a_cu_seqlens_q, TensorView a_LSE_log2, TensorView a_dPsum, TensorView a_dQ_accum, int64_t a_T_q, int64_t a_H_q, int64_t a_num_segments, int64_t a_rows_pad, int64_t grid_x, int64_t grid_y, int64_t grid_z) { + DLDevice dev = a_O.device(); + tvm::ffi::CUDADeviceGuard device_guard(dev.device_id); + check_cuda_tensor(a_O, "O"); + check_dtype(a_O, DLDataType{kDLBfloat, 16, 1}, "O"); + check_cuda_tensor(a_dO, "dO"); + check_dtype(a_dO, DLDataType{kDLBfloat, 16, 1}, "dO"); + check_cuda_tensor(a_LSE, "LSE"); + check_dtype(a_LSE, DLDataType{kDLFloat, 32, 1}, "LSE"); + check_cuda_tensor(a_cu_seqlens_q, "cu_seqlens_q"); + check_dtype(a_cu_seqlens_q, DLDataType{kDLInt, 32, 1}, "cu_seqlens_q"); + check_cuda_tensor(a_LSE_log2, "LSE_log2"); + check_dtype(a_LSE_log2, DLDataType{kDLFloat, 32, 1}, "LSE_log2"); + check_cuda_tensor(a_dPsum, "dPsum"); + check_dtype(a_dPsum, DLDataType{kDLFloat, 32, 1}, "dPsum"); + check_cuda_tensor(a_dQ_accum, "dQ_accum"); + check_dtype(a_dQ_accum, DLDataType{kDLFloat, 32, 1}, "dQ_accum"); + TVM_FFI_CHECK(a_T_q >= -2147483648LL && a_T_q <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'T_q' = " << a_T_q + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_H_q >= -2147483648LL && a_H_q <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'H_q' = " << a_H_q + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_num_segments >= -2147483648LL && a_num_segments <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'num_segments' = " << a_num_segments + << " is outside the i32 range [-2147483648, 2147483647]"; + TVM_FFI_CHECK(a_rows_pad >= -2147483648LL && a_rows_pad <= 2147483647LL, ValueError) + << KERNEL_NAME << ": 'rows_pad' = " << a_rows_pad + << " is outside the i32 range [-2147483648, 2147483647]"; + check_contiguous(a_O, "O"); + check_contiguous(a_dO, "dO"); + check_same_device(a_dO, a_O, "dO", "O"); + check_contiguous(a_LSE, "LSE"); + check_same_device(a_LSE, a_O, "LSE", "O"); + check_contiguous(a_cu_seqlens_q, "cu_seqlens_q"); + check_same_device(a_cu_seqlens_q, a_O, "cu_seqlens_q", "O"); + check_contiguous(a_LSE_log2, "LSE_log2"); + check_same_device(a_LSE_log2, a_O, "LSE_log2", "O"); + check_contiguous(a_dPsum, "dPsum"); + check_same_device(a_dPsum, a_O, "dPsum", "O"); + check_contiguous(a_dQ_accum, "dQ_accum"); + check_same_device(a_dQ_accum, a_O, "dQ_accum", "O"); + TVM_FFI_CHECK((a_T_q >= 1LL), ValueError) + << KERNEL_NAME << ": requires (T_q >= 1)"; + TVM_FFI_CHECK((a_H_q >= 2LL), ValueError) + << KERNEL_NAME << ": requires (H_q >= 2)"; + TVM_FFI_CHECK((a_num_segments >= 1LL), ValueError) + << KERNEL_NAME << ": requires (num_segments >= 1)"; + TVM_FFI_CHECK((a_rows_pad >= 128LL), ValueError) + << KERNEL_NAME << ": requires (rows_pad >= 128)"; + TVM_FFI_CHECK(grid_x > 0 && grid_y > 0 && grid_z > 0, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") must be positive"; + TVM_FFI_CHECK(grid_x <= 2147483647 && grid_y <= 65535 && grid_z <= 65535, ValueError) + << KERNEL_NAME << ": the grid (" << grid_x << ", " << grid_y << ", " << grid_z << ") exceeds CUDA's limits"; + cudaStream_t stream = static_cast(TVMFFIEnvGetStream(dev.device_type, dev.device_id)); + void* v_O = DataPtr(a_O); + void* v_dO = DataPtr(a_dO); + void* v_LSE = DataPtr(a_LSE); + void* v_cu_seqlens_q = DataPtr(a_cu_seqlens_q); + void* v_LSE_log2 = DataPtr(a_LSE_log2); + void* v_dPsum = DataPtr(a_dPsum); + void* v_dQ_accum = DataPtr(a_dQ_accum); + int32_t v_T_q = static_cast(a_T_q); + int32_t v_H_q = static_cast(a_H_q); + int32_t v_num_segments = static_cast(a_num_segments); + int32_t v_rows_pad = static_cast(a_rows_pad); + void* kernel_args[] = {&v_O, &v_dO, &v_LSE, &v_cu_seqlens_q, &v_LSE_log2, &v_dPsum, &v_dQ_accum, &v_T_q, &v_H_q, &v_num_segments, &v_rows_pad}; + cudaLaunchConfig_t config{}; + config.gridDim = dim3(static_cast(grid_x), static_cast(grid_y), static_cast(grid_z)); + config.blockDim = dim3(256, 1, 1); + config.dynamicSmemBytes = 0; + config.stream = stream; + CheckCuda(cudaLaunchKernelExC(&config, reinterpret_cast(kernel_cake_gqa_train_a91752e8a2732bf2e246), kernel_args), + "cudaLaunchKernelExC"); +} + +} // namespace cake_binding_kernel_cake_gqa_train_a91752e8a2732bf2e246 + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, cake_binding_kernel_cake_gqa_train_a91752e8a2732bf2e246::Run); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(build_cuda_version, cake_ffi_build_cuda_version); diff --git a/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu new file mode 100644 index 0000000..574e593 --- /dev/null +++ b/miles_ext_kernels/cake/_vendor/cake_gqa_train/csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu @@ -0,0 +1,174 @@ +/* + * 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 256 + +#include + + +__device__ __forceinline__ float mul_f32_rn_noftz(float a, float b) { + float r; + asm("mul.rn.f32 %0, %1, %2;" : "=f"(r) : "f"(a), "f"(b)); + return r; +} + +extern "C" { + +__global__ __launch_bounds__(256) void +kernel_cake_gqa_train_a91752e8a2732bf2e246(__nv_bfloat16* __restrict__ O, __nv_bfloat16* __restrict__ dO, float* __restrict__ LSE, int* __restrict__ cu_seqlens_q, float* __restrict__ LSE_log2, float* __restrict__ dPsum, float* __restrict__ dQ_accum, int T_q, int H_q, int num_segments, int rows_pad) +{ + 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 (declaration order) === + int pair = blockIdx.y; + int row0 = blockIdx.x * 32; + int half = tid / 16; + int col = lane % 16 * 16; + unsigned long long stats_base = (unsigned long long)pair * (unsigned long long)rows_pad + (unsigned long long)row0; + asm volatile("griddepcontrol.wait;" ::: "memory"); + int limit = (row0 / 128 + 1) * 128; + int started_count = 0; + #pragma unroll 1 + for (int scan_base = 0; scan_base < num_segments; scan_base += 32) { + int b = scan_base + lane; + int _min_0 = ((b) < (num_segments - 1) ? (b) : (num_segments - 1)); + int cu_b = cu_seqlens_q[_min_0]; + unsigned int started = __ballot_sync(0xFFFFFFFF, b < num_segments && 2 * cu_b + 128 * b < limit); + started_count = started_count + (int)(unsigned int)__popc(started); + } + int seg = started_count - 1; + int seg_start = cu_seqlens_q[seg]; + int seg_len = cu_seqlens_q[seg + 1] - seg_start; + int local0 = row0 - (2 * seg_start + 128 * seg) / 128 * 128; + int in_window = ((local0 < (2 * seg_len + 128 - 1) / 128 * 128) ? 1 : 0); + unsigned int o_words[16]; + unsigned int do_words[16]; + #pragma unroll + for (int i = 0; i < 2; i++) { + int local = local0 + i * 16 + half; + int load_token = ((seg_start + (local >> 1)) < (T_q - 1) ? (seg_start + (local >> 1)) : (T_q - 1)); + unsigned long long row_offset = ((unsigned long long)load_token * (unsigned long long)H_q + (unsigned long long)(2 * pair + (local & 1))) * 256 + (unsigned long long)col; + { + asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(o_words[i * 8 + 0]), "=r"(o_words[i * 8 + 1]), "=r"(o_words[i * 8 + 2]), "=r"(o_words[i * 8 + 3]), "=r"(o_words[i * 8 + 4]), "=r"(o_words[i * 8 + 5]), "=r"(o_words[i * 8 + 6]), "=r"(o_words[i * 8 + 7]) : "l"((const void*)((const char*)(O + row_offset + 0) + 0)) : "memory"); + } + { + asm volatile("ld.global.v8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];" + : "=r"(do_words[i * 8 + 0]), "=r"(do_words[i * 8 + 1]), "=r"(do_words[i * 8 + 2]), "=r"(do_words[i * 8 + 3]), "=r"(do_words[i * 8 + 4]), "=r"(do_words[i * 8 + 5]), "=r"(do_words[i * 8 + 6]), "=r"(do_words[i * 8 + 7]) : "l"((const void*)((const char*)(dO + row_offset + 0) + 0)) : "memory"); + } + } + int lse_local = local0 + tid % 32; + int lse_token = ((seg_start + (lse_local >> 1)) < (T_q - 1) ? (seg_start + (lse_local >> 1)) : (T_q - 1)); + float lse = LSE[(unsigned long long)(2 * pair + (lse_local & 1)) * (unsigned long long)T_q + (unsigned long long)lse_token]; + asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); + unsigned int zeros[8]; + #pragma unroll + for (int j = 0; j < 8; j++) { + zeros[j] = 0; + } + unsigned long long chunk_base = stats_base * 256 + (unsigned long long)(tid * 8); + #pragma unroll + for (int i_1 = 0; i_1 < 2; i_1++) { + unsigned long long elem = chunk_base + (unsigned long long)(i_1 * 4096); + { + const unsigned* _raw_stv8_3 = reinterpret_cast(zeros); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(reinterpret_cast(dQ_accum + elem))), "r"(_raw_stv8_3[0]), "r"(_raw_stv8_3[1]), "r"(_raw_stv8_3[2]), "r"(_raw_stv8_3[3]), "r"(_raw_stv8_3[4]), "r"(_raw_stv8_3[5]), "r"(_raw_stv8_3[6]), "r"(_raw_stv8_3[7]) : "memory"); + } + { + const unsigned* _raw_stv8_4 = reinterpret_cast(zeros); + asm volatile( + "st.global.v8.b32 [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "l"((void*)(reinterpret_cast(dQ_accum + (elem + 2048)))), "r"(_raw_stv8_4[0]), "r"(_raw_stv8_4[1]), "r"(_raw_stv8_4[2]), "r"(_raw_stv8_4[3]), "r"(_raw_stv8_4[4]), "r"(_raw_stv8_4[5]), "r"(_raw_stv8_4[6]), "r"(_raw_stv8_4[7]) : "memory"); + } + } + #pragma unroll + for (int i_2 = 0; i_2 < 2; i_2++) { + int local_1 = local0 + i_2 * 16 + half; + float even = 0.0f; + float odd = 0.0f; + #pragma unroll + for (int j_1 = 0; j_1 < 8; j_1++) { + float sums[2] = {even, odd}; + asm("fma.rn.f32.bf16 %0, %1, %2, %0;" : "+f"(sums[0]) : "h"(static_cast((o_words[i_2 * 8 + j_1]))), "h"(static_cast((do_words[i_2 * 8 + j_1])))); + asm("fma.rn.f32.bf16 %0, %1, %2, %0;" : "+f"(sums[1]) : "h"(static_cast((o_words[i_2 * 8 + j_1]) >> 16)), "h"(static_cast((do_words[i_2 * 8 + j_1]) >> 16))); + even = sums[0]; + odd = sums[1]; + } + float row_sum = even + odd; + float _shfl_xor_sync_5 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 8); + row_sum = row_sum + _shfl_xor_sync_5; + float _shfl_xor_sync_6 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 4); + row_sum = row_sum + _shfl_xor_sync_6; + float _shfl_xor_sync_7 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 2); + row_sum = row_sum + _shfl_xor_sync_7; + float _shfl_xor_sync_8 = __shfl_xor_sync(0xFFFFFFFF, row_sum, 1); + row_sum = row_sum + _shfl_xor_sync_8; + if (lane % 16 == 0 && in_window != 0) { + *(reinterpret_cast(dPsum + (stats_base + (unsigned long long)(i_2 * 16 + half))) + (0)) = ((local_1 >> 1 < seg_len) ? row_sum : 0.0f); + } + } + if (tid < 32 && in_window != 0) { + float lse_log2 = ((lse == -CUDART_INF_F) ? 0.0f : mul_f32_rn_noftz(lse, 1.442695f)); + *(reinterpret_cast(LSE_log2 + (stats_base + (unsigned long long)tid)) + (0)) = ((lse_local >> 1 < seg_len) ? lse_log2 : CUDART_INF_F); + } +} + +} // extern "C" diff --git a/miles_ext_kernels/cake/gqa.py b/miles_ext_kernels/cake/gqa.py new file mode 100644 index 0000000..da591e8 --- /dev/null +++ b/miles_ext_kernels/cake/gqa.py @@ -0,0 +1,106 @@ +""" +Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +Licensed under the Apache License, Version 2.0. +https://www.apache.org/licenses/LICENSE-2.0 + +Packed causal GQA attention training kernels, head dimension 256, forward and backward (SM100a / SM103a / SM107a). + +``varlen_gqa_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, *, max_seqlen_k, softmax_scale=None)`` takes BF16 +``q [T_q, H_q, 256]`` and ``k`` / ``v [T_k, H_kv, 256]`` with an even number of query heads per key/value head, +segments packed by int32 device ``cu_seqlens_q`` / ``cu_seqlens_k``, causal with FlashAttention's bottom-right +alignment per segment. ``max_seqlen_k`` is a hard contract: a host int at least every segment's key count. It sizes +the backward grid, so a smaller value leaves key blocks out of the backward and gives wrong dQ / dK / dV without an +error (the output is unaffected). ``--cake-gqa-check-seqlens`` checks it, with the segment boundaries, at every call. + +One variant, ``fast``: the output, the LSE and dK / dV repeat bit for bit run to run, but dQ is summed with FP32 +atomic adds in a run-dependent order. There is no ``det`` variant, so ``--deterministic-mode`` rejects the op. +:func:`check_supported` and :func:`is_supported` tell, without torch, whether the kernels run a head layout; miles' +argument validation applies them to the heads one tensor-parallel rank holds. + +miles selects the kernels with ``--full-attention-backend cake``. Importing this module imports neither torch nor the +implementation. +""" + +from argparse import Namespace + +from miles_ext_kernels.cake._runtime.spec import Flag, OpSpec, lazy_exports + +HEAD_DIM = 256 + + +def check_supported(num_heads_q: int, num_heads_kv: int, head_dim: int = HEAD_DIM) -> None: + """Raise ``ValueError`` naming what the kernels do not run, for the heads one call takes. + + The kernels take head dimension 256 and an even number of query heads per key/value head. Torch-free. + """ + errors = [] + if head_dim != HEAD_DIM: + errors.append(f"head dim {head_dim} (they take {HEAD_DIM})") + if num_heads_q < 1 or num_heads_kv < 1 or num_heads_q % num_heads_kv or (num_heads_q // num_heads_kv) % 2: + errors.append( + f"{num_heads_q} query heads over {num_heads_kv} key/value heads (they need an even number of query heads " + "per key/value head)" + ) + if errors: + raise ValueError("the cake gqa kernels do not run " + "; ".join(errors)) + + +def is_supported(num_heads_q: int, num_heads_kv: int, head_dim: int = HEAD_DIM) -> bool: + """Do the kernels run these heads (:func:`check_supported` without the reason)?""" + try: + check_supported(num_heads_q, num_heads_kv, head_dim) + except ValueError: + return False + return True + + +def _validate(args: Namespace) -> None: + """The heads of one tensor-parallel rank, from Megatron's arguments, must be heads the kernels run. + + A rank holds ``num_attention_heads / TP`` query heads and ``max(num_query_groups, TP) / TP`` key/value heads (the + key/value heads are replicated when there are fewer than TP). Without ``--kv-channels`` the head dimension is + Megatron's default, ``hidden_size / num_attention_heads``; with neither it is not checked here. + """ + heads = getattr(args, "num_attention_heads", None) + if heads is None: + return + tp = getattr(args, "tensor_model_parallel_size", None) or 1 + groups = getattr(args, "num_query_groups", None) if getattr(args, "group_query_attention", False) else None + groups = groups or heads + head_dim = getattr(args, "kv_channels", None) + if head_dim is None and getattr(args, "hidden_size", None): + head_dim = args.hidden_size // heads + try: + check_supported(heads // tp, max(groups, tp) // tp, HEAD_DIM if head_dim is None else head_dim) + except ValueError as exc: + raise ValueError( + f"--full-attention-backend cake with --num-attention-heads {heads}, {groups} query groups and " + f"--tensor-model-parallel-size {tp}: {exc}" + ) from None + + +OP = OpSpec( + name="gqa", + summary="packed causal GQA attention, head dim 256, forward + backward", + variants=("fast",), + selected_by=(("full_attention_backend", "cake"),), + packages=("cake_gqa_train",), + exports=("check_supported", "is_supported", "varlen_gqa_attention"), + flags=( + Flag( + "--cake-gqa-check-seqlens", + dict( + action="store_true", + default=False, + help=( + "Check at every call of the cake gqa kernels that cu_seqlens_q / cu_seqlens_k are well formed and " + "that max_seqlen_k is at least every segment's key count (a smaller one gives wrong gradients " + "without an error). It reads the boundaries on the host: one device synchronization per call." + ), + ), + ), + ), + validate=_validate, +) + +__all__, __getattr__, __dir__ = lazy_exports(__name__, OP) diff --git a/tests/cake/gqa/__init__.py b/tests/cake/gqa/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/cake/gqa/_cases.py b/tests/cake/gqa/_cases.py new file mode 100644 index 0000000..841b655 --- /dev/null +++ b/tests/cake/gqa/_cases.py @@ -0,0 +1,135 @@ +"""Inputs, the FP64 reference and runs of the gqa op, shared by its GPU tests and their child processes. + +Inputs are drawn on the CPU from a seeded generator and copied to the device, so two processes get the same bits. +Imports torch. +""" + +from __future__ import annotations + +import hashlib +import math + +import torch + +from miles_ext_kernels.cake import gqa +from miles_ext_kernels.cake._impl.gqa import ops +from miles_ext_kernels.cake._vendor.cake_gqa_train import cake_backend + +D = 256 +SCALE = D**-0.5 # the default, as miles runs it +# the tensors whose bits do not depend on the run (dQ is summed with FP32 atomic adds) +REPEATABLE = ("out", "lse", "dk", "dv") +# dQ of two runs of one input: one BF16 rounding step (DQ_RTOL) plus an absolute floor for the FP32 reordering error +# where partials cancel, DQ_ATOL_FRACTION of the tensor's largest magnitude +DQ_RTOL = 2**-7 +DQ_ATOL_FRACTION = 2**-16 + +# A miles micro-batch: three sequences, then the pad that aligns the packed tokens to a multiple of 128 as its own +# segment; self-attention, so the query and key boundaries are one. +MILES_SEGMENTS = (1000, 517, 2049, 18) + + +def cu(lens) -> list[int]: + bounds = [0] + for n in lens: + bounds.append(bounds[-1] + int(n)) + return bounds + + +def inputs(q_lens, k_lens, h_q, h_kv, device, seed=0) -> dict: + """``q``, ``k``, ``v``, ``dout`` (BF16), the device boundaries ``cu_q`` / ``cu_k`` and their host copies.""" + g = torch.Generator().manual_seed(seed) + t_q, t_k = sum(q_lens), sum(k_lens) + + def randn(*shape): + return torch.randn(*shape, generator=g, dtype=torch.float32).to(torch.bfloat16).to(device) + + return { + "q": randn(t_q, h_q, D), + "k": randn(t_k, h_kv, D), + "v": randn(t_k, h_kv, D), + "dout": randn(t_q, h_q, D), + "cu_q": torch.tensor(cu(q_lens), dtype=torch.int32, device=device), + "cu_k": torch.tensor(cu(k_lens), dtype=torch.int32, device=device), + "cq": cu(q_lens), + "ck": cu(k_lens), + } + + +def longest_keys(x) -> int: + ck = x["ck"] + return max((ck[i + 1] - ck[i] for i in range(len(ck) - 1)), default=0) + + +def run(x, *, max_seqlen_k: int | None = None, softmax_scale: float | None = SCALE) -> dict: + """Forward and backward through autograd, plus the forward's LSE: ``{"out", "lse", "dq", "dk", "dv"}``. + + ``max_seqlen_k`` defaults to the longest key segment, as miles passes it. + """ + bound = longest_keys(x) if max_seqlen_k is None else max_seqlen_k + q, k, v = (x[t].clone().requires_grad_() for t in ("q", "k", "v")) + out = gqa.varlen_gqa_attention(q, k, v, x["cu_q"], x["cu_k"], max_seqlen_k=bound, softmax_scale=softmax_scale) + out.backward(x["dout"]) + _, lse = cake_backend.gqa_attention_fwd( + x["q"], + x["k"], + x["v"], + x["cu_q"], + x["cu_k"], + softmax_scale=SCALE if softmax_scale is None else softmax_scale, + workspace=ops.PoolWorkspace(x["q"].device), + ) + return {"out": out.detach(), "lse": lse, "dq": q.grad, "dk": k.grad, "dv": v.grad} + + +def repeatable(result: dict) -> dict: + return {name: result[name] for name in REPEATABLE} + + +def digests(result: dict) -> dict[str, str]: + """sha256 of the bytes of every repeatable tensor.""" + return { + name: hashlib.sha256(result[name].detach().cpu().contiguous().view(torch.uint8).numpy().tobytes()).hexdigest() + for name in REPEATABLE + } + + +def assert_dq_close(got: torch.Tensor, ref: torch.Tensor) -> None: + """Two dQ of one input (see DQ_RTOL / DQ_ATOL_FRACTION).""" + floor = DQ_ATOL_FRACTION * float(ref.float().abs().max()) if ref.numel() else 0.0 + torch.testing.assert_close(got, ref, rtol=DQ_RTOL, atol=floor) + + +def reference(x) -> dict: + """FP64 out, lse and gradients of the same BF16 inputs (bottom-right causal per segment).""" + q = x["q"].double().requires_grad_() + k = x["k"].double().requires_grad_() + v = x["v"].double().requires_grad_() + group = q.shape[1] // k.shape[1] + outs, lses = [], [] + for b in range(len(x["cq"]) - 1): + q0, q1, k0, k1 = x["cq"][b], x["cq"][b + 1], x["ck"][b], x["ck"][b + 1] + ql, kl = q1 - q0, k1 - k0 + ks = k[k0:k1].repeat_interleave(group, dim=1) + vs = v[k0:k1].repeat_interleave(group, dim=1) + s = torch.einsum("qhd,khd->hqk", q[q0:q1], ks) * SCALE + seen = torch.arange(kl, device=q.device)[None, :] <= torch.arange(ql, device=q.device)[:, None] + (kl - ql) + has_key = seen.any(dim=-1) + # rows without a key take finite dummy scores (no NaN reaches the gradients) and a zero P + s = s.masked_fill(~seen, -math.inf).masked_fill(~has_key[None, :, None], 0.0) + lse = torch.logsumexp(s, dim=-1) + p = torch.exp(s - lse[..., None]) * seen + outs.append(torch.einsum("hqk,khd->qhd", p, vs)) + lses.append(lse.detach().masked_fill(~has_key[None, :], -math.inf)) + out = torch.cat(outs) + out.backward(x["dout"].double()) + return {"out": out.detach(), "lse": torch.cat(lses, dim=1), "dq": q.grad, "dk": k.grad, "dv": v.grad} + + +def no_key_rows(x) -> torch.Tensor: + """Mask ``[T_q]`` of the query rows that see no key (``i + k_len - q_len < 0``).""" + rows = [] + for b in range(len(x["cq"]) - 1): + ql, kl = x["cq"][b + 1] - x["cq"][b], x["ck"][b + 1] - x["ck"][b] + rows += [i + kl - ql < 0 for i in range(ql)] + return torch.tensor(rows, dtype=torch.bool, device=x["q"].device) diff --git a/tests/cake/gqa/test_arguments.py b/tests/cake/gqa/test_arguments.py new file mode 100644 index 0000000..e02311a --- /dev/null +++ b/tests/cake/gqa/test_arguments.py @@ -0,0 +1,176 @@ +"""The gqa op's arguments, variant and head-layout predicate, through the hooks miles calls (torch-free).""" + +import argparse +from types import SimpleNamespace + +import pytest + +import miles_ext_kernels +from miles_ext_kernels import cake +from miles_ext_kernels.cake import gqa + +NO_DET = "the cake gqa kernels have no det variant" + + +def _parse(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser() + miles_ext_kernels.add_arguments(parser) + return parser.parse_args(argv) + + +def _miles(argv: tuple[str, ...] = (), **values) -> argparse.Namespace: + """miles' parsed flat namespace selecting the op: a Qwen3.5-style full-attention layer (16 query heads, 2 + key/value groups, head dim 256) on one tensor-parallel rank, unless ``values`` say otherwise.""" + args = _parse(list(argv)) + vars(args).update( + full_attention_backend="cake", + deterministic_mode=False, + num_attention_heads=16, + num_query_groups=2, + group_query_attention=True, + kv_channels=256, + hidden_size=2048, + tensor_model_parallel_size=1, + ) + vars(args).update(values) + return args + + +def test_the_op_is_registered_with_its_one_variant_and_no_variant_flag(): + spec = cake.op_spec("gqa") + + assert "gqa" in cake.OPS and spec is gqa.OP + assert spec.variants == ("fast",) and spec.packages == ("cake_gqa_train",) + assert spec.selected_by == (("full_attention_backend", "cake"),) + assert not spec.has_variant_flag + assert "cake_gqa_variant" not in vars(_parse([])) + with pytest.raises(SystemExit): + _parse(["--cake-gqa-variant", "fast"]) + + +def test_the_check_option_is_off_by_default_and_needs_the_backend(): + assert _parse([]).cake_gqa_check_seqlens is False + args = _parse(["--cake-gqa-check-seqlens"]) + assert args.cake_gqa_check_seqlens is True + + with pytest.raises(ValueError, match="--cake-gqa-check-seqlens requires --full-attention-backend cake"): + miles_ext_kernels.validate_args(args) + miles_ext_kernels.validate_args(_miles(["--cake-gqa-check-seqlens"])) + + +@pytest.mark.parametrize("backend", [None, "megatron", "cake"]) +def test_every_backend_validates_outside_deterministic_mode(backend): + args = _miles(full_attention_backend=backend) + + miles_ext_kernels.validate_args(args) + + assert "cake_gqa_variant" not in vars(args) + + +def test_deterministic_mode_rejects_the_op(): + with pytest.raises(ValueError, match=f"{NO_DET}: --full-attention-backend cake is not deterministic"): + miles_ext_kernels.validate_args(_miles(deterministic_mode=True)) + + +def test_the_typed_configuration_s_deterministic_mode_rejects_it_too(): + # the configuration miles builds after parsing carries the flag as args.backend.deterministic_mode + args = _miles(backend=SimpleNamespace(deterministic_mode=True)) + del args.deterministic_mode + + with pytest.raises(ValueError, match=NO_DET): + miles_ext_kernels.validate_args(args) + + +def test_another_attention_backend_validates_in_deterministic_mode(): + miles_ext_kernels.validate_args(_miles(full_attention_backend="megatron", deterministic_mode=True)) + + +@pytest.mark.parametrize( + "heads, kv_heads, head_dim", + [(2, 1, 256), (4, 2, 256), (8, 2, 256), (6, 1, 256), (16, 2, 256), (24, 4, 256), (32, 2, 256), (8, 1, 256)], +) +def test_the_predicate_accepts_an_even_number_of_query_heads_per_key_head(heads, kv_heads, head_dim): + gqa.check_supported(heads, kv_heads, head_dim) + + assert gqa.is_supported(heads, kv_heads, head_dim) + + +@pytest.mark.parametrize( + "heads, kv_heads, head_dim, reason", + [ + (3, 1, 256, "3 query heads over 1 key/value heads"), + (1, 1, 256, "1 query heads over 1 key/value heads"), + (24, 8, 256, "24 query heads over 8 key/value heads"), + (8, 3, 256, "8 query heads over 3 key/value heads"), + (0, 1, 256, "0 query heads over 1 key/value heads"), + (8, 0, 256, "8 query heads over 0 key/value heads"), + (8, 2, 128, "head dim 128"), + ], +) +def test_the_predicate_names_what_the_kernels_do_not_run(heads, kv_heads, head_dim, reason): + with pytest.raises(ValueError, match=f"the cake gqa kernels do not run .*{reason}"): + gqa.check_supported(heads, kv_heads, head_dim) + + assert not gqa.is_supported(heads, kv_heads, head_dim) + + +@pytest.mark.parametrize( + "values, local", + [ + # 24 query heads over 4 groups at TP 8: 3 query heads over 1 replicated key/value head per rank + (dict(num_attention_heads=24, num_query_groups=4, tensor_model_parallel_size=8), "3 query heads over 1"), + (dict(num_attention_heads=8, num_query_groups=8, tensor_model_parallel_size=8), "1 query heads over 1"), + (dict(num_attention_heads=16, num_query_groups=4, tensor_model_parallel_size=8), None), + # without --group-query-attention every query head has its own key/value head + (dict(group_query_attention=False), "16 query heads over 16"), + (dict(kv_channels=128), "head dim 128"), + # no --kv-channels: Megatron's default, hidden_size / num_attention_heads + (dict(kv_channels=None, hidden_size=4096, num_attention_heads=16), None), + (dict(kv_channels=None, hidden_size=2048, num_attention_heads=16), "head dim 128"), + (dict(num_attention_heads=32, num_query_groups=4, tensor_model_parallel_size=2), None), + (dict(num_attention_heads=16, num_query_groups=2, tensor_model_parallel_size=4), None), + (dict(num_attention_heads=32, num_query_groups=2, tensor_model_parallel_size=8), None), + ], +) +def test_validate_args_checks_the_heads_of_one_tensor_parallel_rank(values, local): + args = _miles(**values) + if local is None: + miles_ext_kernels.validate_args(args) + return + with pytest.raises(ValueError, match=f"--full-attention-backend cake with --num-attention-heads .*{local}"): + miles_ext_kernels.validate_args(args) + + +def test_validate_args_leaves_heads_it_cannot_see_to_the_layer(): + miles_ext_kernels.validate_args(_miles(num_attention_heads=None, kv_channels=None, hidden_size=None)) + + +def test_an_unconfigured_process_runs_fast(): + assert cake.variant("gqa") == "fast" + assert cake.option("gqa", "cake_gqa_check_seqlens") is False + + +def test_configure_holds_the_deterministic_rules_and_the_option(): + cake.configure(deterministic_mode=True) + with pytest.raises(ValueError, match=NO_DET): + cake.variant("gqa") + with pytest.raises(ValueError, match="fast is not bit-reproducible under --deterministic-mode"): + cake.configure(cake_gqa_variant="fast", deterministic_mode=True) + with pytest.raises(ValueError, match="must be one of"): + cake.configure(cake_gqa_variant="det") + + cake.configure(cake_gqa_variant="fast", cake_gqa_check_seqlens=True) + assert cake.variant("gqa") == "fast" + assert cake.option("gqa", "cake_gqa_check_seqlens") is True + cake.configure() + assert cake.option("gqa", "cake_gqa_check_seqlens") is False + + +def test_the_values_miles_keeps_configure_a_worker(): + args = _miles(["--cake-gqa-check-seqlens"]) + miles_ext_kernels.validate_args(args) + + miles_ext_kernels.configure({key: value for key, value in vars(args).items() if key.startswith("cake_")}) + + assert cake.variant("gqa") == "fast" + assert cake.option("gqa", "cake_gqa_check_seqlens") is True diff --git a/tests/cake/gqa/test_gqa_cpu_plan.py b/tests/cake/gqa/test_gqa_cpu_plan.py new file mode 100644 index 0000000..8279b33 --- /dev/null +++ b/tests/cake/gqa/test_gqa_cpu_plan.py @@ -0,0 +1,129 @@ +"""The launch plan of the vendored GQA kernels, without a device (torch-free). + +``cake_jit`` (the registry, the grids and the scratch layout; it imports only ``typing``) is loaded by path, so neither +torch nor the vendored package's loader is imported; ``MIN_CUDA`` comes from the torch-free runtime. Checked: every +stage ``ops.STAGES`` launches has exactly one library per target architecture, which the oldest toolkit the runtime +accepts for that architecture builds (``MIN_CUDA``: CUDA 12.8 for SM100a, 12.9 for SM103a, 13.4 for SM107a); the +forward grid covers every tile; the backward grid covers the key blocks of ``max_seqlen_k`` and no more (the +``max_seqlen_k`` contract), or of ``T_k`` without it; the backward scratch is aligned and sized. +""" + +import ast +import importlib.util +import json +import random +from pathlib import Path + +import pytest + +from miles_ext_kernels.cake import MIN_CUDA + +REPO_ROOT = Path(__file__).resolve().parents[3] +PACKAGE_DIR = REPO_ROOT / "miles_ext_kernels" / "cake" / "_vendor" / "cake_gqa_train" +OPS = REPO_ROOT / "miles_ext_kernels" / "cake" / "_impl" / "gqa" / "ops.py" +# ops.STAGES, the forward and the BF16 dK / dV backward, which is the path the entry point runs +STAGES = ("fwd", "bwd_pre", "bwd_bf16", "bwd_post") +# target architecture -> the oldest toolkit the runtime builds it with (README, Install); no library may ask for a +# newer one, which that toolkit would refuse +TOOLKIT = dict(MIN_CUDA) +ARCHES = ("sm_100a", "sm_103a", "sm_107a") +D = 256 + + +@pytest.fixture(scope="module") +def cake_jit(): + path = PACKAGE_DIR / "cake_jit.py" + assert path.is_file(), f"the vendored package is missing: {PACKAGE_DIR}" + spec = importlib.util.spec_from_file_location("vendored_gqa_cake_jit", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(scope="module") +def manifest(): + return json.loads((PACKAGE_DIR / "build.json").read_text()) + + +def test_ops_launches_the_stages_checked_here(): + # ops imports torch: read its STAGES from the source + tree = ast.parse(OPS.read_text()) + values = [ + ast.literal_eval(node.value) + for node in tree.body + if isinstance(node, ast.Assign) and [getattr(t, "id", None) for t in node.targets] == ["STAGES"] + ] + assert values == [STAGES] + + +def test_every_stage_has_one_library_per_target_arch_the_toolkit_builds(cake_jit, manifest): + entries = {entry["name"]: entry for entry in manifest["libraries"]} + assert manifest["schema"] == "cake-build/1" + assert sorted(manifest["arch_flags"]) == sorted(TOOLKIT) == sorted(ARCHES) + for arch in ARCHES: + for stage in STAGES: + serving = [ + name for name, rec in cake_jit.MODULES.items() if rec["stage"] == stage and arch in rec["arches"] + ] + assert serving == [cake_jit.stage_record(stage, arch)], (stage, arch, serving) + entry = entries[f"{cake_jit.MODULES[serving[0]]['library']}_{arch}"] + assert entry["arch"] == arch + assert tuple(int(part) for part in entry["min_cuda"].split(".")) <= TOOLKIT[arch], entry["name"] + launch = cake_jit.MODULES[serving[0]]["launch"] + assert launch["cooperative"] is False and launch["use_pdl"] is False, stage + + +def test_grids_cover_every_tile(cake_jit): + rng = random.Random(7073) + for _ in range(500): + n = rng.randint(1, 40) + q_lens = [rng.choice((0, 1, rng.randint(1, 600), rng.randint(1, 9000))) for _ in range(n)] + k_lens = [rng.choice((0, q, q + rng.randint(0, 20000))) for q in q_lens] + h_q, h_kv = rng.choice(((24, 4), (8, 2), (6, 1), (16, 2), (8, 1))) + dims = dict( + total_q=sum(q_lens), + total_k=sum(k_lens), + num_heads_q=h_q, + num_heads_kv=h_kv, + num_segments=n, + sm_count=148, + ) + fwd_tiles = sum(-(-q // cake_jit.FWD_TILE_ROWS) for q in q_lens) + bound = dims["total_q"] // cake_jit.FWD_TILE_ROWS + n + assert bound >= fwd_tiles + # the forward is persistent: any pair count >= 1 covers every tile, capped at one pair per two SMs + assert cake_jit.grid("fwd", **dims) == (2 * max(1, min(148 // 2, h_q * bound)), 1, 1) + # the backward: a CTA pair per (128-key block of the longest segment, key head, segment) + longest = max(-(-k // cake_jit.BWD_TILE_KEYS) for k in k_lens) + gx, gy, gz = cake_jit.grid("bwd_bf16", **dims) + assert gx % 2 == 0 and gx // 2 >= longest and (gy, gz) == (h_kv, n) + assert cake_jit.grid("bwd_bf16", **dims, max_seqlen_k=max(k_lens)) == (2 * max(1, longest), h_kv, n) + rows = cake_jit.rows_pad(dims["total_q"], n) + assert rows % 128 == 0 and rows >= sum(-(-2 * q // 128) * 128 for q in q_lens) + assert cake_jit.grid("bwd_pre", **dims) == (rows // 32, h_q // 2, 1) + assert cake_jit.grid("bwd_post", **dims) == (rows // 64, h_q // 2, 1) + + +@pytest.mark.parametrize("longest", [128, 129, 1024, 1920, 2048]) +def test_the_backward_grid_covers_the_key_blocks_of_max_seqlen_k_only(cake_jit, longest): + """Why ``max_seqlen_k`` is a contract: the grid has the key blocks of the value given, so a segment with more keys + than that has key blocks no CTA pair visits (their dK / dV are never written).""" + dims = dict(total_q=8 * 2048, total_k=8 * 2048, num_heads_q=16, num_heads_kv=2, num_segments=8, sm_count=148) + blocks = cake_jit.grid("bwd_bf16", **dims, max_seqlen_k=longest)[0] // 2 + + assert blocks == -(-longest // cake_jit.BWD_TILE_KEYS) + assert (blocks * cake_jit.BWD_TILE_KEYS >= 2048) == (longest > 1920) + # without a bound the grid covers T_k keys: correct for any segment, with pairs that only retire + assert cake_jit.grid("bwd_bf16", **dims)[0] // 2 == 8 * 2048 // cake_jit.BWD_TILE_KEYS + + +def test_scratch_layout_is_aligned_and_sized(cake_jit): + for total_q, heads, segs in ((1, 24, 1), (30193, 24, 12), (77, 8, 3), (7744, 6, 4), (3584, 16, 4)): + rows = cake_jit.rows_pad(total_q, segs) + layout = cake_jit.scratch_layout(total_q, heads, segs) + assert layout["dq_accum"][1] == (heads // 2, rows, D) + assert layout["lse_log2"][1] == layout["dpsum"][1] == (heads // 2, rows) + for offset, _shape in layout.values(): + assert offset % cake_jit.SCRATCH_ALIGN == 0 + floats = (heads // 2) * rows * (D + 2) + assert 4 * floats <= cake_jit.scratch_bytes(total_q, heads, segs) <= 4 * floats + 3 * cake_jit.SCRATCH_ALIGN diff --git a/tests/cake/gqa/test_gqa_gpu.py b/tests/cake/gqa/test_gqa_gpu.py new file mode 100644 index 0000000..530f212 --- /dev/null +++ b/tests/cake/gqa/test_gqa_gpu.py @@ -0,0 +1,346 @@ +"""Packed causal GQA attention (``miles_ext_kernels.cake.gqa``) on SM100a / SM103a / SM107a. + +Against an FP64 reference of the same BF16 inputs, on FlashInfer's suite layouts (unequal lengths, tails that are no +multiple of a tile, keys fewer than queries, empty segments, long keys) and on the tensor-parallel-local head pairs +miles runs (8/2, 6/1, 16/2, 8/1) over a packed micro-batch whose last segment is its alignment pad: the output, the LSE +and every gradient within FlashInfer's declared budget and twice the measured error, rows that see no key exactly zero. +The output, the LSE and dK / dV repeat bit for bit run to run (in one process and in a new one), on the default and a +side stream, for any valid ``max_seqlen_k``; dQ is within its FP32 reordering. Also: the opt-in check of +``max_seqlen_k`` and the boundaries, document isolation, checkpointing, no device synchronization in steady state, the +pooled backward scratch and ``release``, the rejected operands, and two GPUs giving the same bits. Needs a B200 / +GB200, B300 / GB300 or R200 and skips elsewhere. +""" + +import json +import os +import subprocess +import sys +from pathlib import Path + +import pytest + +torch = pytest.importorskip("torch") +if not torch.cuda.is_available(): + pytest.skip("CUDA device required", allow_module_level=True) + +from torch.utils.checkpoint import checkpoint # noqa: E402 + +import miles_ext_kernels # noqa: E402 +from miles_ext_kernels import cake # noqa: E402 +from miles_ext_kernels.cake import gqa # noqa: E402 +from miles_ext_kernels.cake._runtime import workspace # noqa: E402 +from miles_ext_kernels.cake._vendor.cake_gqa_train import cake_backend, cake_jit # noqa: E402 + +from ..helpers import POOLED, assert_bitwise, assert_rel_rms # noqa: E402 +from . import _cases as c # noqa: E402 + +if tuple(torch.cuda.get_device_capability()) not in cake.ARCH_BY_CAPABILITY: + pytest.skip( + f"the kernels run on {sorted(cake.ARCH_BY_CAPABILITY)}, not {torch.cuda.get_device_capability()}", + allow_module_level=True, + ) + +REPO_ROOT = Path(__file__).resolve().parents[3] +MIB = 2**20 + +# FlashInfer's accuracy budget for this package: per layout and tensor, the relative L2 error against the FP64 +# reference of the same BF16 inputs is at most 1.10 times FlashAttention CuTe's (its dedicated head-dim-256 backward) +# on B300 against the same reference. +FA4_REL_L2 = {"out": 2.3e-3, "lse": 1.5e-7, "dq": 2.6e-3, "dk": 2.6e-3, "dv": 2.4e-3} +# The largest error per tensor over LAYOUTS, measured with this package on a B200 (sm_100a, CUDA 13.0) from +# standard-normal BF16 inputs drawn on the device (other bits than the CPU draws here): out 2.34e-3 and dV 2.35e-3 on +# long_keys, the LSE 8.04e-8 on keys_fewer_than_queries, dQ 2.62e-3 and dK 2.58e-3 on empty_segments. The test prints +# the errors on the device it runs on (REL_RMS). +MEASURED_REL_L2 = {"out": 2.34e-3, "lse": 8.04e-8, "dq": 2.62e-3, "dk": 2.58e-3, "dv": 2.35e-3} +# FlashInfer's budget, capped at twice the measured error: out, dQ, dK and dV keep FlashInfer's (1.08 to 1.12 times the +# measured error); the LSE, whose errors (7.0e-8 to 8.0e-8 over the layouts) are at FP32 rounding (2^-24 = 6.0e-8), +# gets twice its measured error, 1.61e-7, below FlashInfer's 1.65e-7. +BUDGET = {name: min(1.10 * FA4_REL_L2[name], 2 * MEASURED_REL_L2[name]) for name in FA4_REL_L2} +# Checkpointing: memory kept after a non-reentrant checkpointed forward, less the output storage. +CKPT_RETAINED_LIMIT_MIB = 2.0 + +# (name, q_lens, k_lens, H_q, H_kv) +FI_LAYOUTS = ( + ("tails", (77, 300, 1, 129), (77, 700, 1000, 129), 24, 4), + ("keys_fewer_than_queries", (200, 150, 64), (77, 0, 63), 24, 4), + ("empty_segments", (0, 100, 0, 37, 0), (0, 300, 50, 37, 0), 24, 4), + ("zigzag_cp2", (250, 250, 129, 130), (250, 1000, 129, 517), 24, 4), + ("long_keys", (256, 255), (4099, 2049), 24, 4), + ("group_2", (300, 77), (513, 77), 8, 4), +) +MILES_HEADS = ((8, 2), (6, 1), (16, 2), (8, 1)) +MILES_LAYOUTS = tuple( + (f"miles_{h_q}_{h_kv}", c.MILES_SEGMENTS, c.MILES_SEGMENTS, h_q, h_kv) for h_q, h_kv in MILES_HEADS +) +LAYOUTS = FI_LAYOUTS + MILES_LAYOUTS +# the reviewer's layout of the max_seqlen_k contract: 2048-key segments +EQUAL_SEGMENTS = (2048,) * 4 + + +def _ids(layouts): + return [layout[0] for layout in layouts] + + +@pytest.fixture(scope="module") +def device() -> torch.device: + dev = torch.device("cuda", torch.cuda.current_device()) + # build every library of the forward and backward before the first checked call + assert len(cake.prebuild("gqa", device=dev)) == 4 + return dev + + +@pytest.mark.parametrize("layout", LAYOUTS, ids=_ids(LAYOUTS)) +def test_matches_the_fp64_reference(layout, device): + name, q_lens, k_lens, h_q, h_kv = layout + x = c.inputs(q_lens, k_lens, h_q, h_kv, device) + ref = c.reference(x) + # miles hands over the longest segment; FlashInfer's layouts run the T_k bound, always correct + got = c.run(x, max_seqlen_k=None if name.startswith("miles") else sum(k_lens)) + torch.cuda.synchronize() + finite = torch.isfinite(ref["lse"]) + errors = assert_rel_rms( + {**got, "lse": got["lse"][finite]}, {**ref, "lse": ref["lse"][finite]}, BUDGET, detect_scale=1.02 + ) + print(f"REL_RMS {name} " + " ".join(f"{t}={e:.3e}" for t, e in errors.items())) + # rows that see no key: out 0, lse -inf, dq 0, exactly + no_key = c.no_key_rows(x) + assert torch.equal(got["out"][no_key], torch.zeros_like(got["out"][no_key])) + assert bool((got["lse"][:, no_key] == -float("inf")).all()) + assert bool((got["lse"][:, ~no_key] > -float("inf")).all()) + assert torch.equal(got["dq"][no_key], torch.zeros_like(got["dq"][no_key])) + # keys no query sees (a segment without queries) get exactly zero gradients + unseen = (ref["dk"].abs().amax(dim=(1, 2)) == 0) & (ref["dv"].abs().amax(dim=(1, 2)) == 0) + assert torch.equal(got["dk"][unseen], torch.zeros_like(got["dk"][unseen])) + assert torch.equal(got["dv"][unseen], torch.zeros_like(got["dv"][unseen])) + + +@pytest.mark.parametrize("layout", MILES_LAYOUTS, ids=_ids(MILES_LAYOUTS)) +def test_out_lse_dk_dv_repeat_bit_for_bit_in_one_process(layout, device): + _, q_lens, k_lens, h_q, h_kv = layout + x = c.inputs(q_lens, k_lens, h_q, h_kv, device, seed=5) + first, second = c.run(x), c.run(x) + torch.cuda.synchronize() + + assert_bitwise(c.repeatable(first), c.repeatable(second)) + c.assert_dq_close(second["dq"], first["dq"]) + + +_NEW_PROCESS = """ +import json, sys +import torch +from cake.gqa import _cases as c +h_q, h_kv, seed = (int(arg) for arg in sys.argv[1:]) +x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, h_q, h_kv, torch.device("cuda"), seed=seed) +result = c.run(x) +torch.cuda.synchronize() +print(json.dumps(c.digests(result))) +""" + + +@pytest.mark.parametrize("heads", [(16, 2), (8, 1)], ids=["16_2", "8_1"]) +def test_out_lse_dk_dv_are_the_bits_of_a_new_process(heads, device): + h_q, h_kv = heads + x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, h_q, h_kv, device, seed=11) + here = c.digests(c.run(x)) + paths = [str(REPO_ROOT), str(REPO_ROOT / "tests"), os.environ.get("PYTHONPATH")] + result = subprocess.run( + [sys.executable, "-c", _NEW_PROCESS, str(h_q), str(h_kv), "11"], + cwd=REPO_ROOT, + env={**os.environ, "PYTHONPATH": os.pathsep.join(filter(None, paths))}, + capture_output=True, + text=True, + timeout=1800, + ) + + assert result.returncode == 0, result.stderr[-5000:] + assert json.loads(result.stdout.splitlines()[-1]) == here + + +@pytest.mark.parametrize("segments", [c.MILES_SEGMENTS, EQUAL_SEGMENTS], ids=["miles", "equal_2048"]) +def test_any_valid_max_seqlen_k_gives_the_same_out_lse_dk_dv(segments, device): + """``max_seqlen_k`` sizes the backward grid only: the exact bound, a larger one and ``T_k`` agree bit for bit.""" + x = c.inputs(segments, segments, 16, 2, device, seed=3) + longest = c.longest_keys(x) + exact = c.run(x, max_seqlen_k=longest) + for bound in (longest + 1, 2 * longest, sum(segments)): + other = c.run(x, max_seqlen_k=bound) + torch.cuda.synchronize() + assert_bitwise(c.repeatable(exact), c.repeatable(other)) + c.assert_dq_close(other["dq"], exact["dq"]) + + +@pytest.mark.parametrize("bound", [1024, 1920, 2047]) +def test_the_check_rejects_a_max_seqlen_k_below_the_longest_segment(bound, device): + x = c.inputs(EQUAL_SEGMENTS, EQUAL_SEGMENTS, 16, 2, device) + cake.configure(cake_gqa_check_seqlens=True) + + with pytest.raises(ValueError, match=f"max_seqlen_k is {bound}, but segment 0 has 2048 keys"): + c.run(x, max_seqlen_k=bound) + + +def test_the_check_passes_a_valid_call_and_changes_no_bit(device): + x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, 8, 2, device, seed=7) + plain = c.run(x) + cake.configure(cake_gqa_check_seqlens=True) + checked = c.run(x) + torch.cuda.synchronize() + + assert_bitwise(c.repeatable(plain), c.repeatable(checked)) + + +def test_documents_are_isolated(device): + q_lens, k_lens = (130, 77, 300, 41), (500, 77, 301, 1000) + x = c.inputs(q_lens, k_lens, 8, 2, device) + base = c.run(x) + again = c.run(x) + y = dict(x) + b = 2 + q0, q1, k0, k1 = x["cq"][b], x["cq"][b + 1], x["ck"][b], x["ck"][b + 1] + for key, lo, hi in (("q", q0, q1), ("dout", q0, q1), ("k", k0, k1), ("v", k0, k1)): + y[key] = x[key].clone() + y[key][lo:hi] = -y[key][lo:hi] * 1.5 + pert = c.run(y) + torch.cuda.synchronize() + q_other = torch.ones(sum(q_lens), dtype=torch.bool, device=device) + q_other[q0:q1] = False + k_other = torch.ones(sum(k_lens), dtype=torch.bool, device=device) + k_other[k0:k1] = False + assert not torch.equal(pert["out"][~q_other], base["out"][~q_other]) + assert torch.equal(pert["out"][q_other], base["out"][q_other]) + assert torch.equal(pert["lse"][:, q_other], base["lse"][:, q_other]) + assert torch.equal(pert["dk"][k_other], base["dk"][k_other]) + assert torch.equal(pert["dv"][k_other], base["dv"][k_other]) + # dQ is reduced with atomic adds: the other segments agree to the reduction-order rounding of two plain runs + c.assert_dq_close(again["dq"], base["dq"]) + c.assert_dq_close(pert["dq"][q_other], base["dq"][q_other]) + + +@pytest.mark.parametrize("reentrant", [False, True], ids=["non_reentrant", "reentrant"]) +def test_checkpointing_matches_eager(reentrant, device): + x = c.inputs((1000, 3000, 128), (1000, 3000, 128), 8, 2, device) + + def fn(q, k, v): + return gqa.varlen_gqa_attention(q, k, v, x["cu_q"], x["cu_k"], max_seqlen_k=3000, softmax_scale=c.SCALE) + + def step(use_checkpoint): + q, k, v = (x[t].clone().requires_grad_() for t in ("q", "k", "v")) + out = checkpoint(fn, q, k, v, use_reentrant=reentrant) if use_checkpoint else fn(q, k, v) + out.backward(x["dout"]) + torch.cuda.synchronize() + return out.detach(), q.grad, k.grad, v.grad + + if not reentrant: + q, k, v = (x[t].clone().requires_grad_() for t in ("q", "k", "v")) + torch.cuda.synchronize() + before = torch.cuda.memory_allocated() + out = checkpoint(fn, q, k, v, use_reentrant=False) + torch.cuda.synchronize() + kept = torch.cuda.memory_allocated() - before - out.untyped_storage().nbytes() + assert kept / MIB <= CKPT_RETAINED_LIMIT_MIB, f"a checkpointed forward kept {kept / MIB:.2f} MiB" + del out + eager = step(False) + ckpt = step(True) + assert torch.equal(ckpt[0], eager[0]) + c.assert_dq_close(ckpt[1], eager[1]) + assert torch.equal(ckpt[2], eager[2]) and torch.equal(ckpt[3], eager[3]) + + +def test_steady_state_does_not_synchronize_and_the_check_does(device): + xs = [c.inputs(q, q, 8, 2, device, seed=i) for i, q in enumerate(((700, 77, 3), (300,)))] + + def step(x): + q, k, v = (x[t].clone().requires_grad_() for t in ("q", "k", "v")) + out = gqa.varlen_gqa_attention(q, k, v, x["cu_q"], x["cu_k"], max_seqlen_k=700) + out.backward(x["dout"]) + + for x in xs: + step(x) + torch.cuda.synchronize() + previous = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + for x in xs + xs: + step(x) + cake.configure(cake_gqa_check_seqlens=True) + with pytest.raises(RuntimeError, match="synchroniz"): + step(xs[0]) + finally: + torch.cuda.set_sync_debug_mode(previous) + torch.cuda.synchronize() + + +def test_rejects_unsupported_operands(device): + x = c.inputs((10,), (10,), 24, 4, device) + q, k, v, cq, ck = x["q"], x["k"], x["v"], x["cu_q"], x["cu_k"] + with pytest.raises(ValueError, match="even multiple"): + gqa.varlen_gqa_attention(q[:, :12], k[:, :4], v[:, :4], cq, ck, max_seqlen_k=10) # group 3 + with pytest.raises(ValueError, match="bfloat16"): + gqa.varlen_gqa_attention(q[..., :128], k[..., :128], v[..., :128], cq, ck, max_seqlen_k=10) # head dim 128 + with pytest.raises(ValueError, match="bfloat16"): + gqa.varlen_gqa_attention(q.float(), k.float(), v.float(), cq, ck, max_seqlen_k=10) + with pytest.raises(ValueError, match="int32"): + gqa.varlen_gqa_attention(q, k, v, cq.long(), ck.long(), max_seqlen_k=10) + assert not gqa.is_supported(12, 4) and gqa.is_supported(24, 4) + + +def test_default_and_side_stream_calls_agree_bit_for_bit(device, pooled): + """The pool keeps one backward scratch per stream: alternating streams changes no bit of out / lse / dK / dV.""" + x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, 16, 2, device, seed=9) + default, side = torch.cuda.current_stream(device), torch.cuda.Stream(device) + results = [] + for stream in (default, side, default, side): + stream.wait_stream(default) + with torch.cuda.stream(stream): + results.append(c.run(x)) + torch.cuda.synchronize() + + for other in results[1:]: + assert_bitwise(c.repeatable(results[0]), c.repeatable(other)) + c.assert_dq_close(other["dq"], results[0]["dq"]) + assert workspace.stats()[device.index]["entries"] == 2 # one ("gqa", "bwd_scratch") buffer per stream + + +def test_the_backward_scratch_comes_from_the_pool_and_release_frees_it(device, pooled): + x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, 16, 2, device) + c.run(x) + torch.cuda.synchronize() + need = cake_jit.scratch_bytes(sum(c.MILES_SEGMENTS), 16, len(c.MILES_SEGMENTS)) + + assert workspace.stats()[device.index]["bytes"] >= need + # the vendored per-stream workspaces, kept for the life of the process, stay unused + assert cake_backend._DEFAULT_WORKSPACES == {} + assert miles_ext_kernels.release() >= need + assert workspace.stats()[device.index]["bytes"] == 0 + + +def test_a_pool_bound_of_zero_keeps_nothing_and_changes_no_bit(device): + x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, 8, 2, device, seed=4) + cake.configure(**POOLED) + pooled = c.run(x) + cake.release() + cake.configure(cake_workspace_pool_bytes=0) + unpooled = c.run(x) + torch.cuda.synchronize() + + assert workspace.stats()[device.index]["bytes"] == 0 + assert_bitwise(c.repeatable(pooled), c.repeatable(unpooled)) + + +@pytest.mark.multi_gpu +def test_two_gpus_give_the_same_bits(device, pooled): + if torch.cuda.device_count() < 2: + pytest.skip("needs two GPUs") + first, second = torch.device("cuda", 0), torch.device("cuda", 1) + models = {(p.name, p.multi_processor_count) for p in map(torch.cuda.get_device_properties, (first, second))} + if len(models) != 1: + pytest.skip(f"the two GPUs differ in model or SM count: {sorted(models)}") + cake.prebuild("gqa", device=second) + results = [] + # cuda:1 runs while cuda:0 is the current device: the launches follow the operands' device + for dev in (first, second): + x = c.inputs(c.MILES_SEGMENTS, c.MILES_SEGMENTS, 16, 2, dev, seed=13) + results.append({name: tensor.cpu() for name, tensor in c.repeatable(c.run(x)).items()}) + torch.cuda.synchronize(dev) + + assert_bitwise(results[0], results[1]) + assert workspace.stats()[1]["entries"] == 1 diff --git a/tests/cake/gqa/test_gqa_host.py b/tests/cake/gqa/test_gqa_host.py new file mode 100644 index 0000000..3eec48a --- /dev/null +++ b/tests/cake/gqa/test_gqa_host.py @@ -0,0 +1,152 @@ +"""The host side of the gqa entry point, on the CPU: what it checks before any launch, the opt-in check of the +boundaries and ``max_seqlen_k``, the pooled backward scratch and the prebuild errors. + +Needs torch and tvm_ffi (the implementation imports them); no GPU. CPU tensors reach the kernels' own operand check, +which refuses them, so a call that gets that far shows that the checks before it passed. +""" + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("tvm_ffi") + +from miles_ext_kernels import cake # noqa: E402 +from miles_ext_kernels.cake import gqa # noqa: E402 +from miles_ext_kernels.cake._impl.gqa import ops # noqa: E402 +from miles_ext_kernels.cake._runtime import config, loader, workspace # noqa: E402 +from miles_ext_kernels.cake._vendor.cake_gqa_train import cake_jit # noqa: E402 + +NO_DET = "the cake gqa kernels have no det variant" +REACHED_THE_KERNELS = "must be on q's CUDA device" +# a micro-batch of three segments: 1000, 517 and 2049 keys +LENS = (1000, 517, 2049) +LONGEST = max(LENS) + + +def _cu(lens) -> torch.Tensor: + return torch.tensor([0, *torch.tensor(lens).cumsum(0).tolist()], dtype=torch.int32) + + +def _operands(lens=LENS, h_q=8, h_kv=2): + total = sum(lens) + q = torch.zeros(total, h_q, 256, dtype=torch.bfloat16) + k = torch.zeros(total, h_kv, 256, dtype=torch.bfloat16) + return q, k, k.clone(), _cu(lens), _cu(lens) + + +def test_a_process_configured_for_deterministic_mode_rejects_the_call_when_it_runs(): + cake.configure(deterministic_mode=True) + with pytest.raises(ValueError, match=NO_DET): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=LONGEST) + + cake.configure() # read at the call: the next call goes on to the kernels' operand check + with pytest.raises(ValueError, match=REACHED_THE_KERNELS): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=LONGEST) + + +def test_max_seqlen_k_is_a_required_host_int(): + with pytest.raises(TypeError, match="max_seqlen_k"): + gqa.varlen_gqa_attention(*_operands()) + for value in (None, 2049.0, True, torch.tensor(LONGEST)): + with pytest.raises(TypeError, match="max_seqlen_k must be a host int"): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=value) + with pytest.raises(ValueError, match="max_seqlen_k must be >= 0"): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=-1) + + +def test_the_check_accepts_the_contract(): + q, k, _, cu_q, cu_k = _operands() + + for bound in (LONGEST, LONGEST + 1, sum(LENS)): + ops.check_seqlens(cu_q, cu_k, q.shape[0], k.shape[0], bound) + ops.check_seqlens(_cu(()), _cu(()), 0, 0, 0) # no segment + ops.check_seqlens(_cu((0, 5, 0)), _cu((7, 0, 3)), 5, 10, 7) # empty segments, keys without queries + + +@pytest.mark.parametrize("bound", [0, 1024, 1920, LONGEST - 1]) +def test_the_check_rejects_a_max_seqlen_k_below_the_longest_segment(bound): + q, k, _, cu_q, cu_k = _operands() + + with pytest.raises(ValueError, match=f"max_seqlen_k is {bound}, but segment 2 has 2049 keys: the backward"): + ops.check_seqlens(cu_q, cu_k, q.shape[0], k.shape[0], bound) + + +@pytest.mark.parametrize( + "which, bounds, total, message", + [ + ("q", [1, 1000, 1517, 3566], 3566, "cu_seqlens_q must start at 0"), + ("k", [0, 1000, 900, 3566], 3566, "cu_seqlens_k must start at 0, never decrease"), + ("k", [0, 1000, 1517, 3565], 3566, r"cu_seqlens_k .* end at 3566; got \[0, 1000, 1517, 3565\]"), + ("q", [], 0, r"cu_seqlens_q .*; got \[\]"), + ], +) +def test_the_check_rejects_malformed_boundaries(which, bounds, total, message): + good = _cu(LENS) + cu = torch.tensor(bounds, dtype=torch.int32) + cu_q, cu_k = (cu, good) if which == "q" else (good, cu) + + with pytest.raises(ValueError, match=message): + ops.check_seqlens(cu_q, cu_k, total if which == "q" else sum(LENS), total if which == "k" else sum(LENS), 4096) + + +def test_the_option_checks_before_any_launch_and_is_off_by_default(): + too_small = LONGEST - 1 + with pytest.raises(ValueError, match=REACHED_THE_KERNELS): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=too_small) + + cake.configure(cake_gqa_check_seqlens=True) + with pytest.raises(ValueError, match="but segment 2 has 2049 keys"): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=too_small) + with pytest.raises(ValueError, match=REACHED_THE_KERNELS): + gqa.varlen_gqa_attention(*_operands(), max_seqlen_k=LONGEST) + + +def test_the_backward_scratch_is_one_pooled_buffer_in_the_vendored_layout(pooled): + total_q, heads, segments = 3566, 8, 3 + hits = workspace.stats().get("cpu", {}).get("hits", 0) + views = ops.PoolWorkspace(torch.device("cpu")).scratch(total_q, heads, segments) + layout = cake_jit.scratch_layout(total_q, heads, segments) + nbytes = cake_jit.scratch_bytes(total_q, heads, segments) + + assert [tuple(view.shape) for view in views] == [shape for _, shape in layout.values()] + assert all(view.dtype == torch.float32 and view.is_contiguous() for view in views) + base = views[0].data_ptr() + assert [view.data_ptr() - base for view in views] == [offset for offset, _ in layout.values()] + assert workspace.stats()["cpu"]["bytes"] == nbytes and workspace.stats()["cpu"]["entries"] == 1 + # the same shapes again reuse the buffer; release() frees it + again = ops.PoolWorkspace(torch.device("cpu")).scratch(total_q, heads, segments) + assert again[0].data_ptr() == base and workspace.stats()["cpu"]["hits"] == hits + 1 + assert cake.release() == nbytes and workspace.stats()["cpu"]["bytes"] == 0 + + +def test_at_the_default_bound_of_zero_the_backward_scratch_is_not_kept_and_release_frees_nothing(): + """Never configured, the bound is 0: every backward takes a scratch buffer of its own, in the vendored layout, and + the pool keeps none of them, so ``release`` frees nothing.""" + assert config.option(None, "cake_workspace_pool_bytes") == 0 + total_q, heads, segments = 3566, 8, 3 + before = workspace.stats().get("cpu", {"hits": 0, "misses": 0}) + layout = cake_jit.scratch_layout(total_q, heads, segments) + for _ in range(2): + views = ops.PoolWorkspace(torch.device("cpu")).scratch(total_q, heads, segments) + assert [tuple(view.shape) for view in views] == [shape for _, shape in layout.values()] + after = workspace.stats()["cpu"] + assert (after["bytes"], after["entries"]) == (0, 0) + assert (after["hits"], after["misses"]) == (before["hits"], before["misses"] + 2) + assert cake.release() == 0 + + +def test_prebuild_takes_no_model_parameter_and_rejects_another_gpu(monkeypatch): + with pytest.raises(TypeError): + cake.prebuild("gqa", num_heads=16) + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (9, 0)) + with pytest.raises(NotImplementedError, match="SM100a, SM103a or SM107a, not compute capability 9.0"): + cake.prebuild("gqa") + + +def test_prebuild_refuses_a_toolkit_too_old_for_sm107a(monkeypatch): + monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (10, 7)) + cuda_13_0 = loader.CudaVersion(13, 0, "13.0.88") + monkeypatch.setattr(loader, "_toolkit", lambda nvcc=None: ("/usr/local/cuda/bin/nvcc", cuda_13_0)) + + with pytest.raises(RuntimeError, match=r"13\.4"): + cake.prebuild("gqa") diff --git a/tests/cake/gqa/test_prebuild_gpu.py b/tests/cake/gqa/test_prebuild_gpu.py new file mode 100644 index 0000000..07add6b --- /dev/null +++ b/tests/cake/gqa/test_prebuild_gpu.py @@ -0,0 +1,142 @@ +"""The prebuild contract of the gqa op: after ``prebuild("gqa")``, no call compiles (README, Prebuild). + +``ROUTES`` names every call-time path of the op, each run forward and backward: the exact ``max_seqlen_k`` and the +``T_k`` bound, the head groups miles runs, keys fewer than queries and empty segments, long keys, no query and no +key, a softmax scale, the option ``cake_gqa_check_seqlens``, a side stream and the workspace pool turned on (the +default bound, 0, keeps nothing). The op takes no model parameter. In a fresh process, the libraries a call loads are +the four ``prebuild`` built, and flashinfer is never imported. Needs a B200 / GB200, B300 / GB300 or R200 and skips +elsewhere. +""" + +import os +import subprocess +import sys +from functools import partial +from pathlib import Path + +from miles_ext_kernels import cake +from miles_ext_kernels.cake import gqa + +from ..helpers import POOLED, check_prebuild, requires_arch + +pytestmark = requires_arch("sm_100a", "sm_103a", "sm_107a") + +REPO_ROOT = Path(__file__).resolve().parents[3] +MILES = (1000, 517, 2049, 18) + + +def _fwd_bwd(q_lens, k_lens, h_q, h_kv, *, max_seqlen_k=None, softmax_scale=None, stream=None): + import torch + + from . import _cases as c + + device = torch.device("cuda", torch.cuda.current_device()) + x = c.inputs(q_lens, k_lens, h_q, h_kv, device) + bound = c.longest_keys(x) if max_seqlen_k is None else max_seqlen_k + stream = stream or torch.cuda.current_stream(device) + stream.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(stream): + q, k, v = (x[t].clone().requires_grad_() for t in ("q", "k", "v")) + out = gqa.varlen_gqa_attention(q, k, v, x["cu_q"], x["cu_k"], max_seqlen_k=bound, softmax_scale=softmax_scale) + out.backward(x["dout"]) + torch.cuda.synchronize(device) + + +def _route(variant, *args, options=None, side_stream=False, **kwargs): + if options: + cake.configure(cake_gqa_variant=variant, **options) + if side_stream: + import torch + + kwargs["stream"] = torch.cuda.Stream() + _fwd_bwd(*args, **kwargs) + + +ROUTES = { + "miles_exact_max_seqlen_k": partial(_route, q_lens=MILES, k_lens=MILES, h_q=8, h_kv=2), + "max_seqlen_k_total_keys": partial(_route, q_lens=MILES, k_lens=MILES, h_q=8, h_kv=2, max_seqlen_k=sum(MILES)), + "group_2": partial(_route, q_lens=(300, 77), k_lens=(513, 77), h_q=8, h_kv=4), + "group_6": partial(_route, q_lens=MILES, k_lens=MILES, h_q=6, h_kv=1), + "group_8": partial(_route, q_lens=MILES, k_lens=MILES, h_q=16, h_kv=2), + "keys_fewer_than_queries": partial(_route, q_lens=(200, 150, 64), k_lens=(77, 0, 63), h_q=24, h_kv=4), + "empty_segments": partial(_route, q_lens=(0, 100, 0, 37, 0), k_lens=(0, 300, 50, 37, 0), h_q=24, h_kv=4), + "long_keys": partial(_route, q_lens=(256, 255), k_lens=(4099, 2049), h_q=24, h_kv=4), + "single_segment": partial(_route, q_lens=(4096,), k_lens=(4096,), h_q=16, h_kv=2), + "no_queries": partial(_route, q_lens=(0, 0), k_lens=(64, 32), h_q=8, h_kv=2), + "no_keys": partial(_route, q_lens=(64, 32), k_lens=(0, 0), h_q=8, h_kv=2), + "softmax_scale": partial(_route, q_lens=MILES, k_lens=MILES, h_q=8, h_kv=2, softmax_scale=0.05), + "check_seqlens": partial( + _route, q_lens=MILES, k_lens=MILES, h_q=8, h_kv=2, options={"cake_gqa_check_seqlens": True} + ), + "side_stream": partial(_route, q_lens=MILES, k_lens=MILES, h_q=8, h_kv=2, side_stream=True), + "pooled": partial(_route, q_lens=MILES, k_lens=MILES, h_q=8, h_kv=2, options=dict(POOLED)), +} + + +def test_one_prebuild_serves_every_route(): + check_prebuild("gqa", ROUTES) + + +def test_prebuild_returns_the_four_libraries_once_built(): + from miles_ext_kernels.cake._impl.gqa import ops + + paths = cake.prebuild("gqa") + again = cake.prebuild("gqa") + + assert again == paths and len(paths) == len(ops.STAGES) and all(path.is_file() for path in paths) + assert [path.parent.name for path in paths] == [ops.PACKAGE] * len(ops.STAGES) + + +_FRESH_PROCESS = """ +import sys + +import torch + +from miles_ext_kernels import cake +from miles_ext_kernels.cake import gqa +from miles_ext_kernels.cake._impl.gqa import ops +from miles_ext_kernels.cake._runtime import loader +from miles_ext_kernels.cake._vendor.cake_gqa_train import cake_build + +loaded = [] +load_library = cake_build.load_library + + +def recording_load_library(name): + loaded.append(name) + return load_library(name) + + +cake_build.load_library = recording_load_library +paths = cake.prebuild("gqa") +g = torch.Generator().manual_seed(0) +cu = torch.tensor([0, 300, 811], dtype=torch.int32, device="cuda") +q, k, v = (torch.randn(811, h, 256, generator=g).bfloat16().cuda().requires_grad_() for h in (8, 2, 2)) +with loader.no_build(): + out = gqa.varlen_gqa_attention(q, k, v, cu, cu, max_seqlen_k=511) + out.backward(torch.randn_like(out)) +torch.cuda.synchronize() +expected = ops.libraries(cake.device_arch()) +assert sorted(loaded) == sorted(expected), (loaded, expected) +assert sorted(path.name.rsplit("_", 1)[0] for path in paths) == sorted( + f"{name}_{cake_build.libraries_by_name()[name]['closure_sha256'][:20]}" for name in expected +), paths +flashinfer = sorted(n for n in sys.modules if n == "flashinfer" or n.startswith("flashinfer.")) +assert flashinfer == [], flashinfer +print("loaded", len(loaded)) +""" + + +def test_a_fresh_process_loads_only_the_prebuilt_libraries_and_never_flashinfer(): + pythonpath = os.pathsep.join(filter(None, [str(REPO_ROOT), os.environ.get("PYTHONPATH")])) + result = subprocess.run( + [sys.executable, "-c", _FRESH_PROCESS], + cwd=REPO_ROOT, + env={**os.environ, "PYTHONPATH": pythonpath}, + capture_output=True, + text=True, + timeout=1800, + ) + + assert result.returncode == 0, result.stderr[-5000:] + assert result.stdout.split() == ["loaded", "4"]