diff --git a/benchmarks/bench_recurrent_kda_prefill.py b/benchmarks/bench_recurrent_kda_prefill.py new file mode 100644 index 00000000000..51425a775eb --- /dev/null +++ b/benchmarks/bench_recurrent_kda_prefill.py @@ -0,0 +1,635 @@ +# 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. + +"""CUPTI benchmark for the six frozen recurrent-KDA prefill contract shapes. + +The FlashInfer candidate is always invoked through the public +``recurrent_kda`` API. With ``--flash-kda-peer``, two commit-verified +MoonshotAI/FlashKDA measurements are reported: + +* the raw ``_fwd_raw`` kernel timing scope; +* a public-semantics adapter that follows ``_fwd_raw`` with the same-stream + state copy-back required to emulate ``recurrent_kda`` on FlashKDA. + +All paths use the same deterministic tensors and seeds. Preinitialized +rotating state buffers ensure every timed invocation sees the same initial +state. The FlashInfer path updates each state slot in place inside the kernel; +it has no state scratch or copy-back. Allocation, metadata, sequence ordering, +build/JIT, and state-pool reset are outside the measured region. +""" + +import argparse +import hashlib +import json +import subprocess +from dataclasses import dataclass +from importlib import import_module +from importlib.metadata import version +from pathlib import Path +from typing import Callable, Optional + +import numpy as np +import torch + +from flashinfer.kda import recurrent_kda +from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace +from flashinfer.testing import bench_gpu_time +from flashinfer.utils import get_compute_capability + +FLASH_KDA_PEER_COMMIT = "d2ff19a6a0c82f39f796f637ebd1c36090b1268f" +FLASH_KDA_CUTLASS_COMMIT = "5c149f52a436782210263fb2f19b354443a61c6a" +DEFAULT_STATE_ROTATIONS = 512 + + +@dataclass(frozen=True) +class Case: + name: str + num_heads: int + seq_lens: tuple[int, ...] + packed: bool + seed: int + + +@dataclass +class PreparedCase: + candidate_run: Callable[[], tuple[torch.Tensor, Optional[torch.Tensor]]] + peer_raw_run: Optional[Callable[[], None]] + peer_adapted_run: Optional[Callable[[], None]] + reset_state_pools: Callable[[], None] + candidate_output: torch.Tensor + candidate_state_pool: torch.Tensor + peer_raw_output: Optional[torch.Tensor] + peer_raw_final_state: Optional[torch.Tensor] + peer_adapted_output: Optional[torch.Tensor] + peer_adapted_state_pool: Optional[torch.Tensor] + state_cursors: dict[str, list[int]] + metadata: dict + + +CASES = ( + Case("h96_fixed8192", 96, (8192,), False, 10000), + Case("h96_mixed", 96, (1300, 547, 2048, 963, 271, 3063), True, 10001), + Case("h96_uniform", 96, (1024,) * 8, True, 10002), + Case("h64_fixed8192", 64, (8192,), False, 10003), + Case("h64_mixed", 64, (1300, 547, 2048, 963, 271, 3063), True, 10004), + Case("h64_uniform", 64, (1024,) * 8, True, 10005), +) + + +def _require_cupti() -> None: + try: + from cupti import cupti # noqa: F401 + except ImportError as error: + raise RuntimeError("cupti-python >= 13 is required") from error + cupti_version = version("cupti-python") + if int(cupti_version.split(".", 1)[0]) < 13: + raise RuntimeError(f"cupti-python >= 13 is required, found {cupti_version}") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _git_output(root: Path, *args: str) -> str: + try: + return subprocess.check_output( + ["git", "-C", str(root), *args], + text=True, + stderr=subprocess.STDOUT, + ).strip() + except subprocess.CalledProcessError as error: + raise RuntimeError( + f"failed to verify FlashKDA source provenance at {root}: " + f"{error.output.strip()}" + ) from error + + +def _verify_peer_provenance(flash_kda, source_dir: Path) -> dict: + source_dir = source_dir.resolve(strict=True) + package_path = Path(flash_kda.__file__).resolve(strict=True) + if not package_path.is_relative_to(source_dir): + raise RuntimeError( + "flash_kda must be imported from the verified source checkout: " + f"module={package_path}, checkout={source_dir}" + ) + + source_commit = _git_output(source_dir, "rev-parse", "HEAD") + if source_commit != FLASH_KDA_PEER_COMMIT: + raise RuntimeError( + "unexpected FlashKDA source revision: " + f"expected {FLASH_KDA_PEER_COMMIT}, got {source_commit}" + ) + cutlass_dir = source_dir / "cutlass" + cutlass_commit = _git_output(cutlass_dir, "rev-parse", "HEAD") + if cutlass_commit != FLASH_KDA_CUTLASS_COMMIT: + raise RuntimeError( + "unexpected FlashKDA CUTLASS revision: " + f"expected {FLASH_KDA_CUTLASS_COMMIT}, got {cutlass_commit}" + ) + submodule_record = _git_output(source_dir, "ls-tree", "HEAD", "cutlass").split() + if len(submodule_record) < 3 or submodule_record[2] != cutlass_commit: + raise RuntimeError( + "FlashKDA CUTLASS checkout does not match the pinned gitlink" + ) + tracked_changes = _git_output( + source_dir, + "status", + "--porcelain", + "--untracked-files=no", + ) + if tracked_changes: + raise RuntimeError( + f"verified FlashKDA checkout has tracked modifications:\n{tracked_changes}" + ) + + extension = import_module("flash_kda_C") + extension_path = Path(extension.__file__).resolve(strict=True) + if not extension_path.is_relative_to(source_dir): + raise RuntimeError( + "flash_kda_C must be loaded from the verified source checkout: " + f"extension={extension_path}, checkout={source_dir}" + ) + return { + "repository": "https://github.com/MoonshotAI/FlashKDA.git", + "source_dir": str(source_dir), + "source_commit": source_commit, + "cutlass_commit": cutlass_commit, + "package_path": str(package_path), + "package_sha256": _sha256(package_path), + "extension_path": str(extension_path), + "extension_sha256": _sha256(extension_path), + } + + +def _make_state_pool( + initial_state: torch.Tensor, + rotations: int, +) -> torch.Tensor: + return initial_state.unsqueeze(0).expand(rotations, *initial_state.shape).clone() + + +def _make_case(case: Case, *, state_rotations: int, flash_kda=None) -> PreparedCase: + total_tokens = sum(case.seq_lens) + shape = (1, total_tokens, case.num_heads, 128) + generator = torch.Generator(device="cuda").manual_seed(case.seed) + q = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + k = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + v = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + g = torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) + beta = torch.randn( + (1, total_tokens, case.num_heads), + generator=generator, + device="cuda", + ).to(torch.bfloat16) + A_log = torch.rand( + (case.num_heads,), + generator=generator, + device="cuda", + dtype=torch.float32, + ) + dt_bias = torch.rand( + (case.num_heads, 128), + generator=generator, + device="cuda", + dtype=torch.float32, + ) + initial_state = ( + torch.randn( + (len(case.seq_lens), case.num_heads, 128, 128), + generator=generator, + device="cuda", + ) + * 0.25 + ).to(torch.bfloat16) + candidate_state_pool = _make_state_pool(initial_state, state_rotations) + candidate_output = torch.empty_like(q) + candidate_workspace = RecurrentKDAPrefillWorkspace(q.device) + state_cursors = {"pr": [0], "adapted": [0]} + + offsets = [0] + for seq_len in case.seq_lens: + offsets.append(offsets[-1] + seq_len) + cu_seqlens = ( + torch.tensor(offsets, dtype=torch.int64, device="cuda") if case.packed else None + ) + seq_order = ( + torch.tensor( + sorted( + range(len(case.seq_lens)), + key=case.seq_lens.__getitem__, + reverse=True, + ), + dtype=torch.int32, + device="cuda", + ) + if case.packed + else None + ) + scale = float(1.0 / np.sqrt(128.0)) + + def candidate_run(): + state_index = state_cursors["pr"][0] + if state_index >= state_rotations: + raise RuntimeError( + f"PR state rotations exhausted: {state_index} >= {state_rotations}" + ) + state_cursors["pr"][0] += 1 + return recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=candidate_state_pool[state_index], + output=candidate_output, + output_final_state=False, + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=-5.0, + cu_seqlens=cu_seqlens, + beta_is_logit=True, + seq_order=seq_order, + prefill_workspace=candidate_workspace, + ) + + peer_raw_run = None + peer_adapted_run = None + peer_raw_output = None + peer_raw_final_state = None + peer_adapted_output = None + peer_adapted_state_pool = None + peer_raw_initial_state = None + if flash_kda is not None: + peer_raw_initial_state = initial_state.clone() + peer_raw_final_state = torch.empty_like(initial_state) + peer_raw_output = torch.empty_like(q) + peer_adapted_state_pool = _make_state_pool( + initial_state, + state_rotations, + ) + peer_adapted_final_state = torch.empty_like(initial_state) + peer_adapted_output = torch.empty_like(q) + workspace_size = flash_kda.get_workspace_size( + total_tokens, + case.num_heads, + len(case.seq_lens), + ) + peer_raw_workspace = torch.empty( + workspace_size, + dtype=torch.uint8, + device="cuda", + ) + peer_adapted_workspace = torch.empty( + workspace_size, + dtype=torch.uint8, + device="cuda", + ) + + def peer_raw_run() -> None: + flash_kda._fwd_raw( + q, + k, + v, + g, + beta, + scale, + peer_raw_output, + peer_raw_workspace, + A_log, + dt_bias, + -5.0, + initial_state=peer_raw_initial_state, + final_state=peer_raw_final_state, + cu_seqlens=cu_seqlens, + ) + + def peer_adapted_run() -> None: + state_index = state_cursors["adapted"][0] + if state_index >= state_rotations: + raise RuntimeError( + "adapted-peer state rotations exhausted: " + f"{state_index} >= {state_rotations}" + ) + state_cursors["adapted"][0] += 1 + adapted_state = peer_adapted_state_pool[state_index] + flash_kda._fwd_raw( + q, + k, + v, + g, + beta, + scale, + peer_adapted_output, + peer_adapted_workspace, + A_log, + dt_bias, + -5.0, + initial_state=adapted_state, + final_state=peer_adapted_final_state, + cu_seqlens=cu_seqlens, + ) + adapted_state.copy_(peer_adapted_final_state) + + def reset_state_pools() -> None: + candidate_state_pool.copy_(initial_state.unsqueeze(0)) + state_cursors["pr"][0] = 0 + if peer_raw_initial_state is not None: + peer_raw_initial_state.copy_(initial_state) + if peer_adapted_state_pool is not None: + peer_adapted_state_pool.copy_(initial_state.unsqueeze(0)) + state_cursors["adapted"][0] = 0 + + metadata = { + "name": case.name, + "num_heads": case.num_heads, + "seq_lens": list(case.seq_lens), + "total_tokens": total_tokens, + "layout": "packed" if case.packed else "fixed", + "variant": "m64" if case.name == "h64_fixed8192" else "m128", + "seed": case.seed, + "state_rotation_capacity": state_rotations, + } + return PreparedCase( + candidate_run=candidate_run, + peer_raw_run=peer_raw_run, + peer_adapted_run=peer_adapted_run, + reset_state_pools=reset_state_pools, + candidate_output=candidate_output, + candidate_state_pool=candidate_state_pool, + peer_raw_output=peer_raw_output, + peer_raw_final_state=peer_raw_final_state, + peer_adapted_output=peer_adapted_output, + peer_adapted_state_pool=peer_adapted_state_pool, + state_cursors=state_cursors, + metadata=metadata, + ) + + +def _check_peer(prepared: PreparedCase) -> dict[str, float]: + assert prepared.peer_raw_run is not None + assert prepared.peer_adapted_run is not None + assert prepared.peer_raw_output is not None + assert prepared.peer_raw_final_state is not None + assert prepared.peer_adapted_output is not None + assert prepared.peer_adapted_state_pool is not None + prepared.reset_state_pools() + prepared.candidate_run() + prepared.peer_raw_run() + prepared.peer_adapted_run() + torch.cuda.synchronize() + + candidate_state = prepared.candidate_state_pool[0] + adapted_state = prepared.peer_adapted_state_pool[0] + comparisons = ( + ( + "raw_output_max_abs", + prepared.candidate_output, + prepared.peer_raw_output, + ), + ( + "raw_state_max_abs", + candidate_state, + prepared.peer_raw_final_state, + ), + ( + "adapted_output_max_abs", + prepared.candidate_output, + prepared.peer_adapted_output, + ), + ("adapted_state_max_abs", candidate_state, adapted_state), + ) + diagnostics = {} + for name, actual, expected in comparisons: + diagnostics[name] = float((actual.float() - expected.float()).abs().max()) + torch.testing.assert_close( + actual, + expected, + atol=1e-2, + rtol=1e-2, + ) + return diagnostics + + +def _measure( + run: Callable[[], object], + *, + warmup_ms: int, + bench_ms: int, +) -> tuple[float, list[float]]: + measurements = bench_gpu_time( + run, + enable_cupti=True, + cold_l2_cache=True, + use_cuda_graph=False, + dry_run_time_ms=warmup_ms, + repeat_time_ms=bench_ms, + ) + samples_ms = [float(value) for value in measurements] + return float(np.median(samples_ms)), samples_ms + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--warmup-ms", type=int, default=20) + parser.add_argument("--bench-ms", type=int, default=100) + parser.add_argument( + "--state-rotations", + type=int, + default=DEFAULT_STATE_ROTATIONS, + help="Number of preinitialized same-input state slots per mutable path.", + ) + parser.add_argument( + "--flash-kda-peer", + action="store_true", + help=( + "Compare against FlashKDA commit " + f"{FLASH_KDA_PEER_COMMIT} using both raw and " + "public-semantics-adapted scopes." + ), + ) + parser.add_argument( + "--flash-kda-source-dir", + type=Path, + help=( + "Required with --flash-kda-peer. The imported editable FlashKDA " + "package and extension must resolve inside this exact checkout." + ), + ) + parser.add_argument( + "--json", + type=Path, + help="Optionally write the result list as JSON.", + ) + args = parser.parse_args() + + if args.warmup_ms <= 0 or args.bench_ms <= 0: + parser.error("--warmup-ms and --bench-ms must be positive") + if args.state_rotations <= 0: + parser.error("--state-rotations must be positive") + if args.flash_kda_peer != (args.flash_kda_source_dir is not None): + parser.error( + "--flash-kda-peer and --flash-kda-source-dir must be provided together" + ) + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is required") + if get_compute_capability(torch.device("cuda")) != (10, 0): + raise RuntimeError("frozen recurrent-KDA prefill requires B200 (cc 10.0)") + _require_cupti() + + flash_kda = None + peer_provenance = None + if args.flash_kda_peer: + try: + import flash_kda as imported_flash_kda + except ImportError as error: + raise RuntimeError( + "install MoonshotAI/FlashKDA at " + f"{FLASH_KDA_PEER_COMMIT} to run the peer comparison" + ) from error + flash_kda = imported_flash_kda + assert args.flash_kda_source_dir is not None + peer_provenance = _verify_peer_provenance( + flash_kda, + args.flash_kda_source_dir, + ) + + results = [] + for case in CASES: + prepared = _make_case( + case, + state_rotations=args.state_rotations, + flash_kda=flash_kda, + ) + result = dict(prepared.metadata) + if prepared.peer_raw_run is None: + prepared.reset_state_pools() + prepared.candidate_run() + torch.cuda.synchronize() + prepared.reset_state_pools() + candidate_ms, candidate_samples = _measure( + prepared.candidate_run, + warmup_ms=args.warmup_ms, + bench_ms=args.bench_ms, + ) + candidate_block_medians = [candidate_ms] + result["correctness_peer"] = "not_requested" + else: + assert prepared.peer_adapted_run is not None + correctness = _check_peer(prepared) + result.update( + { + "correctness_peer": "passed", + **correctness, + } + ) + samples = {"pr": [], "raw": [], "adapted": []} + block_medians = {"pr": [], "raw": [], "adapted": []} + state_slots_used = {"pr": [], "adapted": []} + # Symmetric ABCCBA order bounds temperature/clock drift while + # retaining two independent medians for every timing scope. + for backend, run in ( + ("pr", prepared.candidate_run), + ("raw", prepared.peer_raw_run), + ("adapted", prepared.peer_adapted_run), + ("adapted", prepared.peer_adapted_run), + ("raw", prepared.peer_raw_run), + ("pr", prepared.candidate_run), + ): + prepared.reset_state_pools() + torch.cuda.synchronize() + block_median, block_samples = _measure( + run, + warmup_ms=args.warmup_ms, + bench_ms=args.bench_ms, + ) + block_medians[backend].append(block_median) + samples[backend].extend(block_samples) + if backend in state_slots_used: + state_slots_used[backend].append(prepared.state_cursors[backend][0]) + del run + + candidate_block_medians = block_medians["pr"] + candidate_samples = samples["pr"] + candidate_ms = float(np.median(candidate_block_medians)) + raw_ms = float(np.median(block_medians["raw"])) + adapted_ms = float(np.median(block_medians["adapted"])) + result.update( + { + "flash_kda_peer_raw_ms": raw_ms, + "flash_kda_peer_raw_samples_ms": samples["raw"], + "flash_kda_peer_raw_block_medians_ms": (block_medians["raw"]), + "speedup_vs_flash_kda_peer_raw": raw_ms / candidate_ms, + "flash_kda_peer_adapted_ms": adapted_ms, + "flash_kda_peer_adapted_samples_ms": samples["adapted"], + "flash_kda_peer_adapted_block_medians_ms": ( + block_medians["adapted"] + ), + "speedup_vs_flash_kda_peer_adapted": (adapted_ms / candidate_ms), + "peer_raw_timing_scope": ("flash_kda_raw_fwd"), + "peer_adapted_timing_scope": ( + "raw_fwd_plus_public_state_copy_back" + ), + "pair_order": "PR/raw/adapted/adapted/raw/PR", + "same_initial_state_per_timed_call": True, + "state_slots_used_per_block": state_slots_used, + "flash_kda_peer_provenance": peer_provenance, + } + ) + + result.update( + { + "median_ms": candidate_ms, + "median_us": candidate_ms * 1000.0, + "samples_ms": candidate_samples, + "block_medians_ms": candidate_block_medians, + "timing_backend": "cupti", + "cold_l2": True, + "cuda_graph": False, + "timing_scope": ("public_recurrent_kda_with_inplace_state_update"), + "warmup_ms": args.warmup_ms, + "bench_ms": args.bench_ms, + } + ) + results.append(result) + if prepared.peer_raw_run is None: + print( + f"{result['name']:<18} {result['variant']:<4} " + f"{result['median_us']:10.3f} us" + ) + else: + print( + f"{result['name']:<18} {result['variant']:<4} " + f"PR {result['median_us']:10.3f} us " + f"raw {result['flash_kda_peer_raw_ms'] * 1000.0:10.3f} us " + f"{result['speedup_vs_flash_kda_peer_raw']:.4f}x " + f"adapted " + f"{result['flash_kda_peer_adapted_ms'] * 1000.0:10.3f} us " + f"{result['speedup_vs_flash_kda_peer_adapted']:.4f}x" + ) + del prepared + torch.cuda.empty_cache() + + if args.json is not None: + args.json.write_text(json.dumps(results, indent=2) + "\n") + + +if __name__ == "__main__": + main() diff --git a/csrc/kda/flashkda_bf16_fused_m128.cu b/csrc/kda/flashkda_bf16_fused_m128.cu new file mode 100644 index 00000000000..ffaada1f9f6 --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m128.cu @@ -0,0 +1,2555 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// clang-format off +// Generated by tools/export-generated-programs (device kernel). +// Provenance: generated Loom schedule 'flashkda_bf16_fused_m128'; embedded in the host TU as flashkda_bf16_fused_m128_f0217be48b. +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_TMEM_STATE_OFFSET 64 +#define TMEM_TMEM_STATE_INP_OFFSET 0 +#define TMEM_TMEM_U_ACC_OFFSET 224 +#define TMEM_TMEM_U2_INP_OFFSET 224 +#define TMEM_TMEM_U2_ACC_OFFSET 0 +#define TMEM_TMEM_OUT_OFFSET 192 +#define TMEM_TMEM_STATE_OUT_OFFSET 64 +#define NUM_CHUNK_PIPE_STAGES 5 +#define SMEM_SMEM_QD_OFF 1024 +#define SMEM_SMEM_QD_STAGE_BYTES 8192 +#define SMEM_SMEM_QD_STRIDE 41984 +#define SMEM_SMEM_G_RAW_OFF 1024 +#define SMEM_SMEM_G_RAW_STAGE_BYTES 8192 +#define SMEM_SMEM_G_RAW_STRIDE 41984 +#define SMEM_SMEM_G_RAW_ALL_OFF 1024 +#define SMEM_SMEM_G_RAW_ALL_STAGE_BYTES 176128 +#define SMEM_SMEM_G_RAW_ALL_STRIDE 176128 +#define SMEM_SMEM_KD_OFF 9216 +#define SMEM_SMEM_KD_STAGE_BYTES 8192 +#define SMEM_SMEM_KD_STRIDE 41984 +#define SMEM_SMEM_Q_RAW_PREFETCH_OFF 17408 +#define SMEM_SMEM_Q_RAW_PREFETCH_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_RAW_PREFETCH_STRIDE 41984 +#define SMEM_SMEM_FINAL_TRANS_OFF 17408 +#define SMEM_SMEM_FINAL_TRANS_STAGE_BYTES 12288 +#define SMEM_SMEM_FINAL_TRANS_STRIDE 41984 +#define SMEM_SMEM_KR_TRANS_OFF 17408 +#define SMEM_SMEM_KR_TRANS_STAGE_BYTES 8192 +#define SMEM_SMEM_KR_TRANS_STRIDE 41984 +#define SMEM_SMEM_MQK_TRANS_OFF 25600 +#define SMEM_SMEM_MQK_TRANS_STAGE_BYTES 2048 +#define SMEM_SMEM_MQK_TRANS_STRIDE 41984 +#define SMEM_SMEM_INV_OFF 29696 +#define SMEM_SMEM_INV_STAGE_BYTES 2048 +#define SMEM_SMEM_INV_STRIDE 41984 +#define SMEM_SMEM_V_OFF 32384 +#define SMEM_SMEM_V_STAGE_BYTES 8192 +#define SMEM_SMEM_V_STRIDE 41984 +#define SMEM_SMEM_KI_OFF 17408 +#define SMEM_SMEM_KI_STAGE_BYTES 8192 +#define SMEM_SMEM_KI_STRIDE 41984 +#define SMEM_SMEM_GATE_OFF 25600 +#define SMEM_SMEM_GATE_STAGE_BYTES 16384 +#define SMEM_SMEM_GATE_STRIDE 41984 +#define SMEM_SMEM_BETA_RAW_OFF 41984 +#define SMEM_SMEM_BETA_RAW_STAGE_BYTES 512 +#define SMEM_SMEM_BETA_RAW_STRIDE 41984 +#define SMEM_SMEM_INV_WORK_OFF 32384 +#define SMEM_SMEM_INV_WORK_STAGE_BYTES 4096 +#define SMEM_SMEM_INV_WORK_STRIDE 41984 +#define SMEM_SMEM_OUT_OFF 210944 +#define SMEM_SMEM_OUT_STAGE_BYTES 8192 +#define SMEM_SMEM_OUT_STRIDE 8192 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_OFF 41984 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STAGE_BYTES 168452 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STRIDE 168452 +#define SMEM_SMEM_GT_PREFIX_ALL_OFF 41472 +#define SMEM_SMEM_GT_PREFIX_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_PREFIX_ALL_STRIDE 168448 +#define SMEM_SMEM_GT_ALL_OFF 31744 +#define SMEM_SMEM_GT_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_ALL_STRIDE 168448 +#define SMEM_SMEM_PREP_BETA_ALL_OFF 42500 +#define SMEM_SMEM_PREP_BETA_ALL_STAGE_BYTES 168064 +#define SMEM_SMEM_PREP_BETA_ALL_STRIDE 168064 +#define SMEM_SMEM_GATE_RATE_ALL_OFF 42628 +#define SMEM_SMEM_GATE_RATE_ALL_STAGE_BYTES 167940 +#define SMEM_SMEM_GATE_RATE_ALL_STRIDE 167940 +#define SMEM_SMEM_V_ALL_OFF 32384 +#define SMEM_SMEM_V_ALL_STAGE_BYTES 176128 +#define SMEM_SMEM_V_ALL_STRIDE 176128 +#define SMEM_SMEM_GATE_ALL_OFF 25600 +#define SMEM_SMEM_GATE_ALL_STAGE_BYTES 184320 +#define SMEM_SMEM_GATE_ALL_STRIDE 184320 +#define SMEM_TOTAL 227328 +#define THREADS 1024 + +#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)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + 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, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ts_step( + int taddr_out, int taddr_a, int b_lo, uint32_t b_dhi, + uint32_t i_desc, int enable_d) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 dhi;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %5, 0;\n\t" + "mov.b32 dhi, %3;\n\t" + "mov.b64 db, {%2, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], db, %4, p;\n\t" + "}\n" + :: "r"(taddr_out), "r"(taddr_a), "r"(b_lo), "r"(b_dhi), + "r"(i_desc), "r"(enable_d)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__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__ void tmem_ld_x32(float* dst, int tmem_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"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__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__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__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_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.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)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + +__device__ __forceinline__ void elect_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + "}\n" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + 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::cta.global" + ".mbarrier::complete_tx::bytes" + " [%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 tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(1024) void +// FLASHINFER INTEGRATION BEGIN: allow exact state alias +kernel_flashkda_bf16_fused_m128(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +// FLASHINFER INTEGRATION END: allow exact state alias +{ + // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps + // CUDA kernel-start ordering does not acquire the tensor-map proxy. + // One thread acquires each 128-byte map; the CTA barrier publishes those + // acquires to every thread before any TMA instruction can use a map. + if (threadIdx.x == 0) { + asm volatile( + "fence.proxy.tensormap::generic.acquire.gpu [%0], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%1], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%2], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%3], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%4], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%5], 128;\n" + :: "l"(q_tma), "l"(k_tma), "l"(v_tma), "l"(g_tma), + "l"(beta_tma), "l"(out_tma) + : "memory"); + } + __syncthreads(); + // FLASHINFER INTEGRATION END: acquire global tensor maps + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_qd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_qd_addr = smem + 1024; + __nv_bfloat16* smem_g_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_addr = smem + 1024; + __nv_bfloat16* smem_g_raw_all = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_all_addr = smem + 1024; + __nv_bfloat16* smem_kd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_kd_addr = smem + 9216; + __nv_bfloat16* smem_q_raw_prefetch = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_q_raw_prefetch_addr = smem + 17408; + __nv_bfloat16* smem_final_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_final_trans_addr = smem + 17408; + __nv_bfloat16* smem_kr_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_kr_trans_addr = smem + 17408; + __nv_bfloat16* smem_mqk_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 25600); + const int smem_mqk_trans_addr = smem + 25600; + __nv_bfloat16* smem_inv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 29696); + const int smem_inv_addr = smem + 29696; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_v_addr = smem + 32384; + __nv_bfloat16* smem_ki = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_ki_addr = smem + 17408; + float* smem_gate = reinterpret_cast(smem_raw + 25600); + const int smem_gate_addr = smem + 25600; + __nv_bfloat16* smem_beta_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 41984); + const int smem_beta_raw_addr = smem + 41984; + __nv_bfloat16* smem_inv_work = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_inv_work_addr = smem + 32384; + __nv_bfloat16* smem_out = reinterpret_cast<__nv_bfloat16*>(smem_raw + 210944); + const int smem_out_addr = smem + 210944; + float* smem_restore_factor_all = reinterpret_cast(smem_raw + 41984); + const int smem_restore_factor_all_addr = smem + 41984; + float* smem_gt_prefix_all = reinterpret_cast(smem_raw + 41472); + const int smem_gt_prefix_all_addr = smem + 41472; + float* smem_gt_all = reinterpret_cast(smem_raw + 31744); + const int smem_gt_all_addr = smem + 31744; + float* smem_prep_beta_all = reinterpret_cast(smem_raw + 42500); + const int smem_prep_beta_all_addr = smem + 42500; + float* smem_gate_rate_all = reinterpret_cast(smem_raw + 42628); + const int smem_gate_rate_all_addr = smem + 42628; + __nv_bfloat16* smem_v_all = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_v_all_addr = smem + 32384; + float* smem_gate_all = reinterpret_cast(smem_raw + 25600); + const int smem_gate_all_addr = smem + 25600; + + // Mbarrier init (17 groups, 77 barriers) + // Mbarriers at smem_raw[0..616) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'chunk_pipe' --- + // qk_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + // gate_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + mbarrier_init_pred(smem + 48, 1, leader); + mbarrier_init_pred(smem + 56, 1, leader); + mbarrier_init_pred(smem + 64, 1, leader); + mbarrier_init_pred(smem + 72, 1, leader); + // qk_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + mbarrier_init_pred(smem + 88, 1, leader); + mbarrier_init_pred(smem + 96, 1, leader); + mbarrier_init_pred(smem + 104, 1, leader); + mbarrier_init_pred(smem + 112, 1, leader); + // v_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 120, 1, leader); + mbarrier_init_pred(smem + 128, 1, leader); + mbarrier_init_pred(smem + 136, 1, leader); + mbarrier_init_pred(smem + 144, 1, leader); + mbarrier_init_pred(smem + 152, 1, leader); + // v_free: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 160, 4, leader); + mbarrier_init_pred(smem + 168, 4, leader); + mbarrier_init_pred(smem + 176, 4, leader); + mbarrier_init_pred(smem + 184, 4, leader); + mbarrier_init_pred(smem + 192, 4, leader); + // smem_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 200, 1, leader); + mbarrier_init_pred(smem + 208, 1, leader); + mbarrier_init_pred(smem + 216, 1, leader); + mbarrier_init_pred(smem + 224, 1, leader); + mbarrier_init_pred(smem + 232, 1, leader); + // raw_inputs_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 240, 1, leader); + mbarrier_init_pred(smem + 248, 1, leader); + mbarrier_init_pred(smem + 256, 1, leader); + mbarrier_init_pred(smem + 264, 1, leader); + mbarrier_init_pred(smem + 272, 1, leader); + // state_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 280, 4, leader); + mbarrier_init_pred(smem + 288, 4, leader); + mbarrier_init_pred(smem + 296, 4, leader); + mbarrier_init_pred(smem + 304, 4, leader); + mbarrier_init_pred(smem + 312, 4, leader); + // old_out_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 320, 1, leader); + mbarrier_init_pred(smem + 328, 1, leader); + mbarrier_init_pred(smem + 336, 1, leader); + mbarrier_init_pred(smem + 344, 1, leader); + mbarrier_init_pred(smem + 352, 1, leader); + // u_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 360, 4, leader); + mbarrier_init_pred(smem + 368, 4, leader); + mbarrier_init_pred(smem + 376, 4, leader); + mbarrier_init_pred(smem + 384, 4, leader); + mbarrier_init_pred(smem + 392, 4, leader); + // u2_acc_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 400, 1, leader); + mbarrier_init_pred(smem + 408, 1, leader); + mbarrier_init_pred(smem + 416, 1, leader); + mbarrier_init_pred(smem + 424, 1, leader); + mbarrier_init_pred(smem + 432, 1, leader); + // u2_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 440, 4, leader); + mbarrier_init_pred(smem + 448, 4, leader); + mbarrier_init_pred(smem + 456, 4, leader); + mbarrier_init_pred(smem + 464, 4, leader); + mbarrier_init_pred(smem + 472, 4, leader); + // final_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 480, 1, leader); + mbarrier_init_pred(smem + 488, 1, leader); + mbarrier_init_pred(smem + 496, 1, leader); + mbarrier_init_pred(smem + 504, 1, leader); + mbarrier_init_pred(smem + 512, 1, leader); + // out_empty: 1 barriers, init_count=1 + mbarrier_init_pred(smem + 520, 1, leader); + // tmem_dealloc_ready: 1 barriers, init_count=2 + mbarrier_init_pred(smem + 528, 2, leader); + // prep_diag_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 536, 2, leader); + mbarrier_init_pred(smem + 544, 2, leader); + mbarrier_init_pred(smem + 552, 2, leader); + mbarrier_init_pred(smem + 560, 2, leader); + mbarrier_init_pred(smem + 568, 2, leader); + // prep_inv16_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 576, 2, leader); + mbarrier_init_pred(smem + 584, 2, leader); + mbarrier_init_pred(smem + 592, 2, leader); + mbarrier_init_pred(smem + 600, 2, leader); + mbarrier_init_pred(smem + 608, 2, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 616); + if (warp == 0) { + int _tmem_hold = smem + 616; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define qk_full_addr (mbar_base + 0) + #define gate_raw_full_addr (mbar_base + 40) + #define qk_raw_full_addr (mbar_base + 80) + #define v_full_addr (mbar_base + 120) + #define v_free_addr (mbar_base + 160) + #define smem_free_addr (mbar_base + 200) + #define raw_inputs_free_addr (mbar_base + 240) + #define state_inp_ready_addr (mbar_base + 280) + #define old_out_ready_addr (mbar_base + 320) + #define u_inp_ready_addr (mbar_base + 360) + #define u2_acc_ready_addr (mbar_base + 400) + #define u2_inp_ready_addr (mbar_base + 440) + #define final_ready_addr (mbar_base + 480) + #define out_empty_addr (mbar_base + 520) + #define tmem_dealloc_ready_addr (mbar_base + 528) + #define prep_diag_ready_addr (mbar_base + 536) + #define prep_inv16_ready_addr (mbar_base + 576) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_state = taddr + 64; + const int tmem_tmem_state_inp = taddr; + const int tmem_tmem_u_acc = taddr + 224; + const int tmem_tmem_u2_inp = taddr + 224; + const int tmem_tmem_u2_acc = taddr; + const int tmem_tmem_out = taddr + 192; + const int tmem_tmem_state_out = taddr + 64; + + // ---- Register redistribution for WGs split across roles ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + } + + // ---- Role: compute ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 168;"); + { // compute_main + int task_idx = blockIdx.x; + int seq_idx = seq_order[task_idx / num_heads]; + int head_idx = task_idx % num_heads; + long long bos = cu_seqlens[seq_idx]; + long long eos = cu_seqlens[seq_idx + 1]; + int seq_len = (int)(eos - bos); + int num_chunks = (seq_len + 32 - 1) / 32; + int warp_in_wg = warp % 4; + const int tmem_row_base = warp_in_wg * 32 << 16; + int state_row = warp_in_wg * 32 + lane; + int warp_id_in_role = (warp - 0); + int compute_local_warp = warp_id_in_role; + long long state_base = (((long long)seq_idx * (long long)num_heads + (long long)head_idx) * 128 + (long long)state_row) * 128; + #pragma unroll + for (int state_col_block = 0; state_col_block < 4; state_col_block++) { + float state_frag[32]; + state_frag[0] = 0.0f; + state_frag[1] = 0.0f; + state_frag[2] = 0.0f; + state_frag[3] = 0.0f; + state_frag[4] = 0.0f; + state_frag[5] = 0.0f; + state_frag[6] = 0.0f; + state_frag[7] = 0.0f; + state_frag[8] = 0.0f; + state_frag[9] = 0.0f; + state_frag[10] = 0.0f; + state_frag[11] = 0.0f; + state_frag[12] = 0.0f; + state_frag[13] = 0.0f; + state_frag[14] = 0.0f; + state_frag[15] = 0.0f; + state_frag[16] = 0.0f; + state_frag[17] = 0.0f; + state_frag[18] = 0.0f; + state_frag[19] = 0.0f; + state_frag[20] = 0.0f; + state_frag[21] = 0.0f; + state_frag[22] = 0.0f; + state_frag[23] = 0.0f; + state_frag[24] = 0.0f; + state_frag[25] = 0.0f; + state_frag[26] = 0.0f; + state_frag[27] = 0.0f; + state_frag[28] = 0.0f; + state_frag[29] = 0.0f; + state_frag[30] = 0.0f; + state_frag[31] = 0.0f; + if (use_initial_state != 0) { + { + const uint4* _vptr_0 = reinterpret_cast(initial_state + state_base + (long long)(state_col_block * 32)); + uint4 _vld_0[2]; + #pragma unroll + for (int _blk = 0; _blk < 2; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_frag[0 + _blk * 8 + _pair * 2])[0]), "=f"((&state_frag[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(initial_state + state_base + (long long)(state_col_block * 32) + 16); + uint4 _vld_1[2]; + #pragma unroll + for (int _blk = 0; _blk < 2; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&state_frag[16 + _blk * 8 + _pair * 2])[0]), "=f"((&state_frag[16 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + tmem_st_x32_f32(taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_block * 32), state_frag); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + unsigned int compute_stage = 0; + unsigned int _phase_qk_full = 0; + unsigned int _phase_v_full = 0; + unsigned int _phase_old_out_ready = 0; + unsigned int _phase_u2_acc_ready = 0; + unsigned int _phase_final_ready = 0; + #pragma unroll 1 + for (int chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + mbarrier_wait(qk_full_addr + (compute_stage) * 8, _phase_qk_full); + #pragma unroll 1 + for (int state_col_block_1 = 0; state_col_block_1 < 4; state_col_block_1++) { + int state_addr = taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_block_1 * 32); + float _tmem_load_0[32]; + tmem_ld_x32(&_tmem_load_0[0], state_addr); + uint32_t _tmem_load_0_bf16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + 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"(taddr + (unsigned int)tmem_row_base + (unsigned int)(state_col_block_1 * 16)), "r"(*reinterpret_cast(&_tmem_load_0_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[7])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[8])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[9])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[10])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[11])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[12])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[13])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[14])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[15])) + : "memory"); + float state_scale[16]; + #pragma unroll + for (int state_half = 0; state_half < 2; state_half++) { + #pragma unroll + for (int state_col = 0; state_col < 16; state_col++) { + state_scale[state_col] = smem_gt_all[compute_stage * 10496 + (unsigned int)(state_col_block_1 * 32) + (unsigned int)(state_half * 16) + (unsigned int)state_col]; + } + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast((_tmem_load_0 + state_half * 16))[_ls], reinterpret_cast(state_scale)[_ls]); + } + tmem_st_x32_f32(state_addr, _tmem_load_0); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(state_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(v_full_addr + (compute_stage) * 8, _phase_v_full); + mbarrier_wait(old_out_ready_addr + (compute_stage) * 8, _phase_old_out_ready); + float _tmem_load_1[32]; + tmem_ld_x32(&_tmem_load_1[0], taddr + 224 + (unsigned int)tmem_row_base); + #pragma unroll + for (int residual_half = 0; residual_half < 2; residual_half++) { + float residual_v[16]; + float residual_beta[16]; + #pragma unroll + for (int residual_col = 0; residual_col < 16; residual_col++) { + int token_col = residual_half * 16 + residual_col; + __nv_bfloat16 v_value = smem_v_all[compute_stage * 20992 + (unsigned int)(token_col * 128) + (unsigned int)state_row]; + float _cvt_f32_2 = __bfloat162float(v_value); + residual_v[residual_col] = _cvt_f32_2; + residual_beta[residual_col] = smem_prep_beta_all[compute_stage * 10496 + (unsigned int)token_col]; + } + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + sub_f32x2_inplace(&reinterpret_cast(residual_v)[_ls], reinterpret_cast((_tmem_load_1 + residual_half * 16))[_ls]); + #pragma unroll + for (int _ls = 0; _ls < 8; _ls++) + mul_f32x2_inplace(&reinterpret_cast(residual_v)[_ls], reinterpret_cast(residual_beta)[_ls]); + uint32_t residual_v_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(residual_v[_lp*2 + 0], residual_v[_lp*2+1 + 0])); + residual_v_bf16[_lp] = *(uint32_t*)&_bf2; + } + tmem_st_x8_u32(taddr + 224 + (unsigned int)tmem_row_base + (unsigned int)(residual_half * 8), (const uint32_t*)residual_v_bf16); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(v_free_addr + (compute_stage) * 8); + mbarrier_arrive(u_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(u2_acc_ready_addr + (compute_stage) * 8, _phase_u2_acc_ready); + float _tmem_load_2[32]; + tmem_ld_x32(&_tmem_load_2[0], taddr + (unsigned int)tmem_row_base); + uint32_t _tmem_load_2_bf16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_2[_lp*2 + 0], _tmem_load_2[_lp*2+1 + 0])); + _tmem_load_2_bf16[_lp] = *(uint32_t*)&_bf2; + } + 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"(taddr + 224 + (unsigned int)tmem_row_base), "r"(*reinterpret_cast(&_tmem_load_2_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[7])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[8])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[9])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[10])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[11])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[12])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[13])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[14])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[15])) + : "memory"); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(u2_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(final_ready_addr + (compute_stage) * 8, _phase_final_ready); + compute_stage += 1; + if (compute_stage == 5) { compute_stage = 0; _phase_qk_full ^= 1; _phase_v_full ^= 1; _phase_old_out_ready ^= 1; _phase_u2_acc_ready ^= 1; _phase_final_ready ^= 1; } + } + if (store_final_state != 0) { + #pragma unroll + for (int state_col_block_2 = 0; state_col_block_2 < 4; state_col_block_2++) { + float _tmem_load_3[32]; + tmem_ld_x32(&_tmem_load_3[0], taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_block_2 * 32)); + { + __nv_bfloat162 _pk[8]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_3[0 + 0], _tmem_load_3[0 + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_3[0 + 2], _tmem_load_3[0 + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_3[0 + 4], _tmem_load_3[0 + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_3[0 + 6], _tmem_load_3[0 + 7]); + _pk[4] = __floats2bfloat162_rn(_tmem_load_3[0 + 8], _tmem_load_3[0 + 9]); + _pk[5] = __floats2bfloat162_rn(_tmem_load_3[0 + 10], _tmem_load_3[0 + 11]); + _pk[6] = __floats2bfloat162_rn(_tmem_load_3[0 + 12], _tmem_load_3[0 + 13]); + _pk[7] = __floats2bfloat162_rn(_tmem_load_3[0 + 14], _tmem_load_3[0 + 15]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32))))[0]) = *reinterpret_cast(&_pk[0]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32))))[8]) = *reinterpret_cast(&_pk[4]); + } + { + __nv_bfloat162 _pk[8]; + _pk[0] = __floats2bfloat162_rn(_tmem_load_3[16 + 0], _tmem_load_3[16 + 1]); + _pk[1] = __floats2bfloat162_rn(_tmem_load_3[16 + 2], _tmem_load_3[16 + 3]); + _pk[2] = __floats2bfloat162_rn(_tmem_load_3[16 + 4], _tmem_load_3[16 + 5]); + _pk[3] = __floats2bfloat162_rn(_tmem_load_3[16 + 6], _tmem_load_3[16 + 7]); + _pk[4] = __floats2bfloat162_rn(_tmem_load_3[16 + 8], _tmem_load_3[16 + 9]); + _pk[5] = __floats2bfloat162_rn(_tmem_load_3[16 + 10], _tmem_load_3[16 + 11]); + _pk[6] = __floats2bfloat162_rn(_tmem_load_3[16 + 12], _tmem_load_3[16 + 13]); + _pk[7] = __floats2bfloat162_rn(_tmem_load_3[16 + 14], _tmem_load_3[16 + 15]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32) + 16)))[0]) = *reinterpret_cast(&_pk[0]); + *reinterpret_cast(&((__nv_bfloat16*)(final_state + (state_base + (long long)(state_col_block_2 * 32) + 16)))[8]) = *reinterpret_cast(&_pk[4]); + } + } + } + asm volatile("barrier.sync 10, 128;" ::: "memory"); + if (compute_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: epilogue ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // epilogue_main + int task_idx_1 = blockIdx.x; + int seq_idx_1 = seq_order[task_idx_1 / num_heads]; + int head_idx_1 = task_idx_1 % num_heads; + long long bos_1 = cu_seqlens[seq_idx_1]; + long long eos_1 = cu_seqlens[seq_idx_1 + 1]; + int seq_len_1 = (int)(eos_1 - bos_1); + int num_chunks_1 = (seq_len_1 + 32 - 1) / 32; + int warp_id_in_role_1 = (warp - 4); + int epilogue_local_warp = warp_id_in_role_1; + int warp_in_wg_1 = warp % 4; + const int tmem_row_base_1 = warp_in_wg_1 * 32 << 16; + int state_row_1 = warp_in_wg_1 * 32 + lane; + unsigned int epilogue_stage = 0; + unsigned int output_stage = 0; + unsigned int _phase_final_ready_1 = 0; + #pragma unroll 1 + for (int chunk_idx_1 = 0; chunk_idx_1 < num_chunks_1; chunk_idx_1++) { + mbarrier_wait(final_ready_addr + (epilogue_stage) * 8, _phase_final_ready_1); + int chunk_is_full = ((seq_len_1 >= (chunk_idx_1 + 1) * 32) ? 1 : 0); + if (chunk_is_full != 0) { + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1) + : "memory"); + float _tmem_load_5[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_5[0])), "=r"(*reinterpret_cast(&_tmem_load_5[1])), "=r"(*reinterpret_cast(&_tmem_load_5[2])), "=r"(*reinterpret_cast(&_tmem_load_5[3])), "=r"(*reinterpret_cast(&_tmem_load_5[4])), "=r"(*reinterpret_cast(&_tmem_load_5[5])), "=r"(*reinterpret_cast(&_tmem_load_5[6])), "=r"(*reinterpret_cast(&_tmem_load_5[7])), "=r"(*reinterpret_cast(&_tmem_load_5[8])), "=r"(*reinterpret_cast(&_tmem_load_5[9])), "=r"(*reinterpret_cast(&_tmem_load_5[10])), "=r"(*reinterpret_cast(&_tmem_load_5[11])), "=r"(*reinterpret_cast(&_tmem_load_5[12])), "=r"(*reinterpret_cast(&_tmem_load_5[13])), "=r"(*reinterpret_cast(&_tmem_load_5[14])), "=r"(*reinterpret_cast(&_tmem_load_5[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1 + 1048576) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + } + if (epilogue_local_warp == 0) { + if (chunk_idx_1 >= 2) { + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + int out_stage_addr = smem_out_addr + output_stage * 8192; + #pragma unroll + for (int dim_half = 0; dim_half < 2; dim_half++) { + unsigned int out_packed[8]; + if (dim_half == 0) { + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + out_packed[_lp] = *(uint32_t*)&_bf2; + } + } else { + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_5[_lp*2 + 0], _tmem_load_5[_lp*2+1 + 0])); + out_packed[_lp] = *(uint32_t*)&_bf2; + } + } + #pragma unroll + for (int token_group = 0; token_group < 2; token_group++) { + int mtx_idx = lane / 8; + int row_addr = lane & 7; + int dim_base = epilogue_local_warp * 32 + dim_half * 16 + (mtx_idx & 1) * 8; + int token_base = token_group * 16 + mtx_idx / 2 * 8; + int token_addr = token_base + row_addr; + int token_pair = token_addr / 2; + int token_parity = token_addr & 1; + int raw_row = token_pair + dim_base / 64 * 16; + int raw_col = (dim_base & 63 ^ (token_pair & 3) << 4 ^ token_parity << 3) + token_parity * 64; + int stsm_offset = (raw_row * 128 + raw_col) * 2; + const int pack_base = token_group * 4; + uint32_t _stmatrix_addr_0 = static_cast((unsigned long long)(out_stage_addr + stsm_offset)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_0), "r"(*reinterpret_cast(&out_packed[pack_base])), "r"(*reinterpret_cast(&out_packed[pack_base + 1])), "r"(*reinterpret_cast(&out_packed[pack_base + 2])), "r"(*reinterpret_cast(&out_packed[pack_base + 3])) + : "memory"); + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + tma_store_4d(out_tma, 0, (int)(bos_1 + (long long)(chunk_idx_1 * 32)), head_idx_1, 0, smem_out_addr + output_stage * 8192); + } + asm volatile("cp.async.bulk.commit_group;"); + } + output_stage = output_stage ^ 1; + } else { + float _tmem_load_6[32]; + tmem_ld_x32(&_tmem_load_6[0], taddr + 192 + (unsigned int)tmem_row_base_1); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + } + #pragma unroll + for (int token_col_1 = 0; token_col_1 < 32; token_col_1++) { + long long out_token = bos_1 + (long long)(chunk_idx_1 * 32 + token_col_1); + if (out_token < eos_1) { + long long out_idx = (out_token * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_1; + out[out_idx] = _tmem_load_6[token_col_1]; + } + } + } + epilogue_stage += 1; + if (epilogue_stage == 5) { epilogue_stage = 0; _phase_final_ready_1 ^= 1; } + } + if (epilogue_local_warp == 0) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: mma ---- + } else if (warp == 9) { + { // mma_main + int task_idx_2 = blockIdx.x; + int seq_idx_2 = seq_order[task_idx_2 / num_heads]; + long long bos_2 = cu_seqlens[seq_idx_2]; + long long eos_2 = cu_seqlens[seq_idx_2 + 1]; + int seq_len_2 = (int)(eos_2 - bos_2); + int num_chunks_2 = (seq_len_2 + 32 - 1) / 32; + unsigned int mma_stage = 0; + unsigned int _phase_qk_full_1 = 0; + unsigned int _phase_state_inp_ready = 0; + unsigned int _phase_out_empty_0 = 1; + unsigned int _phase_u_inp_ready = 0; + unsigned int _phase_u2_inp_ready = 0; + #pragma unroll 1 + for (int _chunk_idx = 0; _chunk_idx < num_chunks_2; _chunk_idx++) { + mbarrier_wait(qk_full_addr + (mma_stage) * 8, _phase_qk_full_1); + mbarrier_wait(state_inp_ready_addr + (mma_stage) * 8, _phase_state_inp_ready); + mbarrier_wait(out_empty_addr, _phase_out_empty_0); + _phase_out_empty_0 ^= 1; + int _mma_b_addr_0 = smem_qd_addr + mma_stage * 41984; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 134743184;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_out), "r"(_mma_b_lo_0), "r"(tmem_tmem_state_inp), "r"(0)); + int _mma_b_addr_1 = smem_kd_addr + mma_stage * 41984; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 134743184;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u_acc), "r"(_mma_b_lo_1), "r"(tmem_tmem_state_inp), "r"(0)); + elect_commit2(old_out_ready_addr + (mma_stage) * 8, raw_inputs_free_addr + (mma_stage) * 8); + mbarrier_wait(u_inp_ready_addr + (mma_stage) * 8, _phase_u_inp_ready); + int _mma_b_addr_2 = smem_inv_addr + mma_stage * 41984; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0xC0004010;\n\t" + "mov.b32 id, 134743184;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 64;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u2_acc), "r"(_mma_b_lo_2), "r"(tmem_tmem_u2_inp), "r"(0)); + elect_commit(u2_acc_ready_addr + (mma_stage) * 8); + mbarrier_wait(u2_inp_ready_addr + (mma_stage) * 8, _phase_u2_inp_ready); + int _mma_b_addr_3 = smem_final_trans_addr + mma_stage * 41984; + int _mma_b_lo_3 = make_warp_uniform(((_mma_b_addr_3 >> 4) & 0x3FFF) | 0x1000000); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 136905872;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_state_out), "r"(_mma_b_lo_3), "r"(tmem_tmem_u2_inp), "r"(1)); + elect_commit2(final_ready_addr + (mma_stage) * 8, smem_free_addr + (mma_stage) * 8); + mma_stage += 1; + if (mma_stage == 5) { mma_stage = 0; _phase_qk_full_1 ^= 1; _phase_state_inp_ready ^= 1; _phase_u_inp_ready ^= 1; _phase_u2_inp_ready ^= 1; } + } + unsigned int _phase_tmem_dealloc_ready_0 = 0; + mbarrier_wait(tmem_dealloc_ready_addr, _phase_tmem_dealloc_ready_0); + _phase_tmem_dealloc_ready_0 ^= 1; + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + // ---- Role: load ---- + } else if (warp == 10) { + { // load_main + int task_idx_3 = blockIdx.x; + int seq_idx_3 = seq_order[task_idx_3 / num_heads]; + int head_idx_2 = task_idx_3 % num_heads; + long long bos_3 = cu_seqlens[seq_idx_3]; + long long eos_3 = cu_seqlens[seq_idx_3 + 1]; + int seq_len_3 = (int)(eos_3 - bos_3); + int num_chunks_3 = (seq_len_3 + 32 - 1) / 32; + unsigned int load_stage = 0; + unsigned int _phase_v_free = 1; + unsigned int _phase_qk_full_2 = 0; + #pragma unroll 1 + for (int chunk_idx_2 = 0; chunk_idx_2 < num_chunks_3; chunk_idx_2++) { + mbarrier_wait(v_free_addr + (load_stage) * 8, _phase_v_free); + mbarrier_wait(qk_full_addr + (load_stage) * 8, _phase_qk_full_2); + int chunk_is_full_1 = ((seq_len_3 >= (chunk_idx_2 + 1) * 32) ? 1 : 0); + if (elect_sync()) { + if (chunk_is_full_1 != 0) { + mbarrier_arrive_expect_tx(v_full_addr + (load_stage) * 8, 8192); + tma_3d_gmem2smem(smem_v_addr + load_stage * 41984, v_tma, 0, head_idx_2, (int)(bos_3 + (long long)(chunk_idx_2 * 32)), v_full_addr + (load_stage) * 8); + } + } + if (chunk_is_full_1 == 0) { + #pragma unroll + for (int v_load_iter = 0; v_load_iter < 16; v_load_iter++) { + int v_item = v_load_iter * 32 + lane; + int row = v_item / 16; + int segment = v_item % 16; + long long token = bos_3 + (long long)(chunk_idx_2 * 32 + row); + int token_valid = ((token < eos_3) ? 1 : 0); + long long v_src = (token * (long long)num_heads + (long long)head_idx_2) * 128 + (long long)(segment * 8); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(smem_v_addr + load_stage * 41984 + (unsigned int)((row * 128 + segment * 8) * 2)), "l"(v + v_src), "r"((token_valid != 0) ? 16 : 0)); + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } + asm volatile("barrier.sync 8, 32;" ::: "memory"); + if (elect_sync()) { + if (chunk_is_full_1 == 0) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(v_full_addr + (load_stage) * 8); + } + } + load_stage += 1; + if (load_stage == 5) { load_stage = 0; _phase_v_free ^= 1; _phase_qk_full_2 ^= 1; } + } + } + // ---- Role: prep ---- + } else if (warp >= 12 && warp <= 31) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // prep_main + int task_idx_4 = blockIdx.x; + int seq_idx_4 = seq_order[task_idx_4 / num_heads]; + int head_idx_3 = task_idx_4 % num_heads; + long long bos_4 = cu_seqlens[seq_idx_4]; + long long eos_4 = cu_seqlens[seq_idx_4 + 1]; + int seq_len_4 = (int)(eos_4 - bos_4); + int num_chunks_4 = (seq_len_4 + 32 - 1) / 32; + int instance_id = (warp - 12) / 4; + int prep_instance = instance_id; + int warp_id_in_role_2 = (warp - 12); + int prep_local_warp = warp_id_in_role_2 - prep_instance * 4; + int prep_tid = prep_local_warp * 32 + lane; + int num_prep_iters = (num_chunks_4 + 4 - prep_instance) / 5; + unsigned int prep_stage = (unsigned int)prep_instance; + int gate_rate_stage_f32 = prep_instance * 10496; + if (prep_tid == 0) { + float _expf_0 = __expf(A_log[head_idx_3]); + smem_gate_rate_all[gate_rate_stage_f32] = _expf_0; + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + unsigned int _phase_raw_inputs_free = 1; + unsigned int _phase_gate_raw_full = 0; + unsigned int _phase_smem_free = 1; + unsigned int _phase_qk_raw_full = 0; + unsigned int _phase_prep_diag_ready = 0; + unsigned int _phase_prep_inv16_ready = 0; + #pragma unroll 1 + for (int prep_iter = 0; prep_iter < num_prep_iters; prep_iter++) { + int chunk_idx_3 = prep_iter * 5 + prep_instance; + int stage_f32 = prep_stage * 10496; + int stage_bf16 = prep_stage * 20992; + int chunk_is_full_2 = ((seq_len_4 >= (chunk_idx_3 + 1) * 32) ? 1 : 0); + float early_beta_value = 0.0f; + float early_gate0 = 0.0f; + if (chunk_is_full_2 != 0) { + mbarrier_wait(raw_inputs_free_addr + (prep_stage) * 8, _phase_raw_inputs_free); + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(gate_raw_full_addr + (prep_stage) * 8, 8704); + tma_3d_gmem2smem(smem_g_raw_addr + prep_stage * 41984, g_tma, 0, head_idx_3, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + tma_2d_gmem2smem(smem_beta_raw_addr + prep_stage * 41984, beta_tma, head_idx_3 / 8 * 8, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + mbarrier_arrive_expect_tx(qk_raw_full_addr + (prep_stage) * 8, 16384); + tma_4d_gmem2smem(smem_kd_addr + prep_stage * 41984, k_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + mbarrier_wait(gate_raw_full_addr + (prep_stage) * 8, _phase_gate_raw_full); + if (prep_local_warp == 2 && lane < 32) { + unsigned int beta_raw_pair[1]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&beta_raw_pair[0])) : "r"(smem_beta_raw_addr + prep_stage * 41984 + (unsigned int)(lane * 16) + (unsigned int)(head_idx_3 % 8 / 2 * 4))); + float beta_raw_pair_fp32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&beta_raw_pair_fp32[_pair * 2])[0]), "=f"((&beta_raw_pair_fp32[_pair * 2])[1]) + : "r"(beta_raw_pair[_pair + 0])); + } + float beta_logit = beta_raw_pair_fp32[0]; + if (head_idx_3 % 2 != 0) { + beta_logit = beta_raw_pair_fp32[1]; + } + float _tanh_approx_0; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_0) : "f"(beta_logit * 0.5f)); + early_beta_value = _tanh_approx_0 * 0.5f + 0.5f; + } + if (prep_tid < 128) { + float early_gate_rate = smem_gate_rate_all[stage_f32]; + float early_gate_bias = dt_bias[head_idx_3 * 128 + prep_tid]; + __nv_bfloat16 early_gate_raw = smem_g_raw_all[stage_bf16 + prep_tid]; + float _cvt_f32_0 = __bfloat162float(early_gate_raw); + float early_gate_arg = early_gate_rate * (_cvt_f32_0 + early_gate_bias); + float _tanh_approx_1; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_1) : "f"(early_gate_arg * 0.5f)); + float early_gate_sigmoid = _tanh_approx_1 * 0.5f + 0.5f; + early_gate0 = lower_bound * 1.4426950408889634f * early_gate_sigmoid; + } + } + mbarrier_wait(smem_free_addr + (prep_stage) * 8, _phase_smem_free); + if (chunk_is_full_2 != 0) { + if (prep_local_warp == 0) { + if (elect_sync()) { + tma_4d_gmem2smem(smem_q_raw_prefetch_addr + prep_stage * 41984, q_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + } + if (chunk_is_full_2 == 0) { + #pragma unroll + for (int gate_load_pass = 0; gate_load_pass < 4; gate_load_pass++) { + int gate_load_item = gate_load_pass * 128 + prep_tid; + int gate_load_row = gate_load_item / 16; + int gate_load_segment = gate_load_item % 16; + long long gate_load_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_load_row); + long long gate_load_base = (gate_load_token * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(gate_load_segment * 8); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(smem_g_raw_addr + prep_stage * 41984 + (unsigned int)(gate_load_item * 16)), "l"(g + gate_load_base), "r"((gate_load_token < eos_4) ? 16 : 0)); + } + } + if (chunk_is_full_2 == 0) { + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + } + if (prep_local_warp == 2 && lane < 32) { + float beta_value = early_beta_value; + if (chunk_is_full_2 == 0) { + long long beta_token = bos_4 + (long long)(chunk_idx_3 * 32 + lane); + if (beta_token < eos_4) { + float beta_logit_1 = (float)beta[beta_token * (long long)num_heads + (long long)head_idx_3]; + float _tanh_approx_2; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_2) : "f"(beta_logit_1 * 0.5f)); + beta_value = _tanh_approx_2 * 0.5f + 0.5f; + } + } + smem_prep_beta_all[stage_f32 + lane] = beta_value; + } + if (prep_tid < 128) { + int gate_col = prep_tid; + float gate_rate = smem_gate_rate_all[stage_f32]; + float gate_bias = dt_bias[head_idx_3 * 128 + gate_col]; + float prefix_log2 = 0.0f; + for (int gate_row = 0; gate_row < 32; gate_row++) { + long long gate_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_row); + float gate_log2 = 0.0f; + int gate_needs_compute = 1; + if (gate_row == 0) { + if (chunk_is_full_2 != 0) { + gate_log2 = early_gate0; + gate_needs_compute = 0; + } + } + if (gate_needs_compute != 0) { + if (gate_token < eos_4) { + __nv_bfloat16 gate_raw = smem_g_raw_all[stage_bf16 + gate_row * 128 + gate_col]; + float _cvt_f32_1 = __bfloat162float(gate_raw); + float gate_arg = gate_rate * (_cvt_f32_1 + gate_bias); + float _tanh_approx_3; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_3) : "f"(gate_arg * 0.5f)); + float gate_sigmoid = _tanh_approx_3 * 0.5f + 0.5f; + gate_log2 = lower_bound * 1.4426950408889634f * gate_sigmoid; + } + } + prefix_log2 += gate_log2; + smem_gate_all[stage_f32 + gate_row * 128 + gate_col] = prefix_log2; + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (chunk_is_full_2 != 0) { + mbarrier_wait(qk_raw_full_addr + (prep_stage) * 8, _phase_qk_raw_full); + } + if (prep_tid < 128) { + float total_log2 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_0 = approx_exp2(total_log2 - lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + prep_tid] = _exp2_0; + } + if (prep_tid == 0) { + float _exp2_1 = approx_exp2(lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + 128] = _exp2_1; + } + #pragma unroll 1 + for (int work_pass = 0; work_pass < 4; work_pass++) { + int work_item = work_pass * 128 + prep_tid; + int row_1 = work_item / 16; + int segment_1 = work_item % 16; + long long token_1 = bos_4 + (long long)(chunk_idx_3 * 32 + row_1); + int token_valid_1 = ((token_1 < eos_4) ? 1 : 0); + long long gmem_base = (token_1 * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(segment_1 * 8); + float q_raw_vec[8]; + float k_raw_vec[8]; + q_raw_vec[0] = 0.0f; + q_raw_vec[1] = 0.0f; + q_raw_vec[2] = 0.0f; + q_raw_vec[3] = 0.0f; + q_raw_vec[4] = 0.0f; + q_raw_vec[5] = 0.0f; + q_raw_vec[6] = 0.0f; + q_raw_vec[7] = 0.0f; + k_raw_vec[0] = 0.0f; + k_raw_vec[1] = 0.0f; + k_raw_vec[2] = 0.0f; + k_raw_vec[3] = 0.0f; + k_raw_vec[4] = 0.0f; + k_raw_vec[5] = 0.0f; + k_raw_vec[6] = 0.0f; + k_raw_vec[7] = 0.0f; + if (chunk_is_full_2 != 0) { + unsigned int packed[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed[0])), "=r"(*reinterpret_cast(&packed[(0) + 1])), "=r"(*reinterpret_cast(&packed[(0) + 2])), "=r"(*reinterpret_cast(&packed[(0) + 3])) + : "r"((smem_q_raw_prefetch_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32[_pair * 2])[0]), "=f"((&packed_fp32[_pair * 2])[1]) + : "r"(packed[_pair + 0])); + } + #pragma unroll + for (int value_idx = 0; value_idx < 8; value_idx++) { + q_raw_vec[value_idx] = packed_fp32[value_idx]; + } + unsigned int packed_0[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0[0])), "=r"(*reinterpret_cast(&packed_0[(0) + 1])), "=r"(*reinterpret_cast(&packed_0[(0) + 2])), "=r"(*reinterpret_cast(&packed_0[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32[_pair * 2])[0]), "=f"((&packed_0_fp32[_pair * 2])[1]) + : "r"(packed_0[_pair + 0])); + } + #pragma unroll + for (int value_idx_1 = 0; value_idx_1 < 8; value_idx_1++) { + k_raw_vec[value_idx_1] = packed_0_fp32[value_idx_1]; + } + } else if (token_valid_1 != 0) { + { + const uint4* _vptr_0 = reinterpret_cast(q + gmem_base); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(k + gmem_base); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + float q_sum = 0.0f; + float k_sum = 0.0f; + for (int elem_in_segment = 0; elem_in_segment < 8; elem_in_segment++) { + float q_raw = q_raw_vec[elem_in_segment]; + float k_raw = k_raw_vec[elem_in_segment]; + float _fma_0 = __fmaf_rn(q_raw, q_raw, q_sum); + q_sum = _fma_0; + float _fma_1 = __fmaf_rn(k_raw, k_raw, k_sum); + k_sum = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 8); + q_sum += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 8); + k_sum += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 4); + q_sum += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 4); + k_sum += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 2); + q_sum += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 2); + k_sum += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 1); + q_sum += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 1); + k_sum += _shfl_xor_7; + float _rsqrt_0 = rsqrtf(q_sum + 1e-06f); + float q_inv = _rsqrt_0; + float _rsqrt_1 = rsqrtf(k_sum + 1e-06f); + float k_inv = _rsqrt_1; + const float2 _scale2_2 = {q_inv, q_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(q_raw_vec)[_ls], _scale2_2); + const float2 _scale2_3 = {k_inv, k_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(k_raw_vec)[_ls], _scale2_3); + float qd_vec[8]; + float kd_vec[8]; + float ki_vec[8]; + for (int elem_in_segment_1 = 0; elem_in_segment_1 < 8; elem_in_segment_1++) { + int col = segment_1 * 8 + elem_in_segment_1; + float prefix = smem_gate_all[stage_f32 + row_1 * 128 + col]; + float common_log2 = lower_bound * 1.4426950408889634f * 16.0f; + float _exp2_2 = approx_exp2(prefix - common_log2); + float decay = _exp2_2; + qd_vec[elem_in_segment_1] = decay; + kd_vec[elem_in_segment_1] = decay; + ki_vec[elem_in_segment_1] = k_raw_vec[elem_in_segment_1] / decay; + } + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], reinterpret_cast(q_raw_vec)[_ls]); + const float2 _scale2_4 = {scale, scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], _scale2_4); + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(kd_vec)[_ls], reinterpret_cast(k_raw_vec)[_ls]); + unsigned int packed_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(qd_vec[_lp*2 + 0], qd_vec[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word = 0; word < 4; word++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word * 4)), "r"(packed_1[word])); + } + unsigned int packed_0_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(kd_vec[_lp*2 + 0], kd_vec[_lp*2+1 + 0])); + packed_0_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_1 = 0; word_1 < 4; word_1++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_1 * 4)), "r"(packed_0_1[word_1])); + } + unsigned int packed_1_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ki_vec[_lp*2 + 0], ki_vec[_lp*2+1 + 0])); + packed_1_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_2 = 0; word_2 < 4; word_2++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_2 * 4)), "r"(packed_1_1[word_2])); + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + int pair_row_base = prep_local_warp / 2 * 16; + int pair_col_base = prep_local_warp % 2 * 16; + unsigned int a_frag[4]; + unsigned int b_frag[4]; + float acc[8]; + if (pair_row_base >= pair_col_base) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + int row0 = pair_row_base + lane / 4; + int row1 = row0 + 8; + int col0 = pair_col_base + lane % 4 * 2; + float beta0 = smem_prep_beta_all[stage_f32 + row0]; + float beta1 = smem_prep_beta_all[stage_f32 + row1]; + float seed[8]; + seed[0] = 0.0f; + seed[1] = 0.0f; + seed[2] = 0.0f; + seed[3] = 0.0f; + seed[4] = 0.0f; + seed[5] = 0.0f; + seed[6] = 0.0f; + seed[7] = 0.0f; + if (row0 > col0) { + seed[0] = acc[0] * beta0; + } + if (row0 > col0 + 1) { + seed[1] = acc[1] * beta0; + } + if (row1 > col0) { + seed[2] = acc[2] * beta1; + } + if (row1 > col0 + 1) { + seed[3] = acc[3] * beta1; + } + if (row0 > col0 + 8) { + seed[4] = acc[4] * beta0; + } + if (row0 > col0 + 9) { + seed[5] = acc[5] * beta0; + } + if (row1 > col0 + 8) { + seed[6] = acc[6] * beta1; + } + if (row1 > col0 + 9) { + seed[7] = acc[7] * beta1; + } + unsigned int seed_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(seed[_lp*2 + 0], seed[_lp*2+1 + 0])); + seed_packed[_lp] = *(uint32_t*)&_bf2; + } + int seed_lane_row = lane % 16; + int seed_lane_col = lane / 16 * 8; + int byte_off = (pair_row_base + seed_lane_row) * 128 + (pair_col_base + seed_lane_col) * 2; + int swizzled_off = byte_off ^ (byte_off >> 7 & 7) << 4; + int seed_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off; + uint32_t _stmatrix_addr_5 = static_cast((unsigned long long)seed_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_5), "r"(*reinterpret_cast(&seed_packed[0])), "r"(*reinterpret_cast(&seed_packed[1])), "r"(*reinterpret_cast(&seed_packed[2])), "r"(*reinterpret_cast(&seed_packed[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + } else { + acc[0] = 0.0f; + acc[1] = 0.0f; + acc[2] = 0.0f; + acc[3] = 0.0f; + acc[4] = 0.0f; + acc[5] = 0.0f; + acc[6] = 0.0f; + acc[7] = 0.0f; + } + int row0_1 = pair_row_base + lane / 4; + int row1_1 = row0_1 + 8; + int col0_1 = pair_col_base + lane % 4 * 2; + float mqk[8]; + mqk[0] = 0.0f; + mqk[1] = 0.0f; + mqk[2] = 0.0f; + mqk[3] = 0.0f; + mqk[4] = 0.0f; + mqk[5] = 0.0f; + mqk[6] = 0.0f; + mqk[7] = 0.0f; + if (row0_1 >= col0_1) { + mqk[0] = acc[0]; + } + if (row0_1 >= col0_1 + 1) { + mqk[1] = acc[1]; + } + if (row1_1 >= col0_1) { + mqk[2] = acc[2]; + } + if (row1_1 >= col0_1 + 1) { + mqk[3] = acc[3]; + } + if (row0_1 >= col0_1 + 8) { + mqk[4] = acc[4]; + } + if (row0_1 >= col0_1 + 9) { + mqk[5] = acc[5]; + } + if (row1_1 >= col0_1 + 8) { + mqk[6] = acc[6]; + } + if (row1_1 >= col0_1 + 9) { + mqk[7] = acc[7]; + } + unsigned int mqk_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(mqk[_lp*2 + 0], mqk[_lp*2+1 + 0])); + mqk_packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int publish_pair = 0; publish_pair < 2; publish_pair++) { + int publish_row = pair_col_base + publish_pair * 8 + (lane & 7); + int publish_col = 128 + pair_row_base + lane / 8 * 8; + uint32_t _stmatrix_addr_6 = static_cast((unsigned long long)(smem_final_trans_addr + prep_stage * 41984 + (unsigned int)(publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 ^ (publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 >> 7 & 7) << 4))); + asm volatile("stmatrix.sync.aligned.m8n8.x2.trans.shared.b16 [%0], {%1, %2};\n" + :: "r"(_stmatrix_addr_6), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2])), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2 + 1])) + : "memory"); + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_tid < 128) { + float total_log2_1 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_3 = approx_exp2(total_log2_1); + smem_gt_all[stage_f32 + prep_tid] = _exp2_3; + } + if (prep_local_warp >= 2) { + int stage_f32_0 = prep_stage * 10496; + float restore_scale = smem_restore_factor_all[stage_f32_0 + 128]; + float restore_factor[8]; + int restore_segment = lane & 15; + #pragma unroll + for (int restore_elem = 0; restore_elem < 8; restore_elem++) { + int restore_col = restore_segment * 8 + restore_elem; + restore_factor[restore_elem] = smem_restore_factor_all[stage_f32_0 + restore_col]; + } + #pragma unroll 1 + for (int restore_pass = 0; restore_pass < 6; restore_pass++) { + int restore_row = 8 + (prep_local_warp - 2) * 12 + restore_pass * 2 + (lane >> 4); + float restore_qd_values[8]; + float restore_kd_values[8]; + float restore_ki_values[8]; + unsigned int packed_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_2[0])), "=r"(*reinterpret_cast(&packed_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_2[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_1[_pair * 2])[0]), "=f"((&packed_fp32_1[_pair * 2])[1]) + : "r"(packed_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_2 = 0; value_idx_2 < 8; value_idx_2++) { + restore_qd_values[value_idx_2] = packed_fp32_1[value_idx_2]; + } + unsigned int packed_0_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_2[0])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_1[_pair * 2])[0]), "=f"((&packed_0_fp32_1[_pair * 2])[1]) + : "r"(packed_0_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_3 = 0; value_idx_3 < 8; value_idx_3++) { + restore_kd_values[value_idx_3] = packed_0_fp32_1[value_idx_3]; + } + unsigned int packed_1_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_2[0])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32[_pair * 2])[0]), "=f"((&packed_1_fp32[_pair * 2])[1]) + : "r"(packed_1_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_4 = 0; value_idx_4 < 8; value_idx_4++) { + restore_ki_values[value_idx_4] = packed_1_fp32[value_idx_4]; + } + float restore_kr_values[8]; + #pragma unroll + for (int restore_elem_1 = 0; restore_elem_1 < 8; restore_elem_1++) { + restore_kr_values[restore_elem_1] = restore_ki_values[restore_elem_1] * restore_factor[restore_elem_1]; + } + const float2 _scale2_7 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values)[_ls], _scale2_7); + const float2 _scale2_8 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values)[_ls], _scale2_8); + unsigned int packed_2_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values[_lp*2 + 0], restore_qd_values[_lp*2+1 + 0])); + packed_2_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_3 = 0; word_3 < 4; word_3++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_3 * 4)), "r"(packed_2_1[word_3])); + } + unsigned int packed_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values[_lp*2 + 0], restore_kd_values[_lp*2+1 + 0])); + packed_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_4 = 0; word_4 < 4; word_4++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_4 * 4)), "r"(packed_3[word_4])); + } + unsigned int packed_4[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values[_lp*2 + 0], restore_kr_values[_lp*2+1 + 0])); + packed_4[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_5 = 0; word_5 < 4; word_5++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_5 * 4)), "r"(packed_4[word_5])); + } + } + } + if (prep_local_warp == 0) { + int inverse_row = lane; + int diag_block = inverse_row / 8; + int lane_in_diag = lane & 7; + float inv_row[8]; + unsigned int packed_5[4]; + int byte_off_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_1 = byte_off_1 ^ (byte_off_1 >> 7 & 7) << 4; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_5[0])), "=r"(*reinterpret_cast(&packed_5[(0) + 1])), "=r"(*reinterpret_cast(&packed_5[(0) + 2])), "=r"(*reinterpret_cast(&packed_5[(0) + 3])) + : "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1)); + float packed_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_2[_pair * 2])[0]), "=f"((&packed_fp32_2[_pair * 2])[1]) + : "r"(packed_5[_pair + 0])); + } + #pragma unroll + for (int value_idx_5 = 0; value_idx_5 < 8; value_idx_5++) { + inv_row[value_idx_5] = packed_fp32_2[value_idx_5]; + } + #pragma unroll + for (int diag_elem = 0; diag_elem < 8; diag_elem++) { + if (lane_in_diag == diag_elem) { + inv_row[diag_elem] = 1.0f; + } + } + int diag_group_base = lane - lane_in_diag; + #pragma unroll + for (int src_row = 0; src_row < 7; src_row++) { + float row_scale = -inv_row[src_row]; + #pragma unroll + for (int prev_col = 0; prev_col < src_row; prev_col++) { + int pivot_lane = diag_group_base + src_row; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, inv_row[prev_col], pivot_lane); + float pivot = _shfl_0; + if (lane_in_diag > src_row) { + float _fma_2 = __fmaf_rn(row_scale, pivot, inv_row[prev_col]); + inv_row[prev_col] = _fma_2; + } + } + if (lane_in_diag > src_row) { + inv_row[src_row] = row_scale; + } + } + unsigned int packed_0_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(inv_row[_lp*2 + 0], inv_row[_lp*2+1 + 0])); + packed_0_3[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_1_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_2 = byte_off_1_1 ^ (byte_off_1_1 >> 7 & 7) << 4; + #pragma unroll + for (int word_6 = 0; word_6 < 4; word_6++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_2 + (unsigned int)(word_6 * 4)), "r"(packed_0_3[word_6])); + } + } + if (prep_local_warp < 2) { + if (elect_sync()) { + mbarrier_arrive(prep_diag_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_diag_ready_addr + (prep_stage) * 8, _phase_prep_diag_ready); + } + if (prep_local_warp < 2) { + int lane_row = lane & 7; + int byte_off_2 = (prep_local_warp * 16 + 8 + lane_row) * 128 + (prep_local_warp * 16 + 8) * 2; + int swizzled_off_3 = byte_off_2 ^ (byte_off_2 >> 7 & 7) << 4; + int d_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3; + int byte_off_0 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_1_1 = byte_off_0 ^ (byte_off_0 >> 7 & 7) << 4; + int c_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_1; + int byte_off_2_1 = (prep_local_warp * 16 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_3_1 = byte_off_2_1 ^ (byte_off_2_1 >> 7 & 7) << 4; + int a_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_1; + unsigned int d_frag[2]; + unsigned int c_frag[1]; + float dc_acc[4]; + unsigned int dc_bf16[2]; + unsigned int inv_a_frag[1]; + float o_acc[4]; + unsigned int o_bf16[2]; + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[0]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[1]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(c_frag[0]) + : "r"(c_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc_acc[0]), "=f"(dc_acc[1]), "=f"(dc_acc[2]), "=f"(dc_acc[3]) + : "r"(d_frag[0]), "r"(d_frag[1]), "r"(c_frag[0])); + const float2 _scale2_9 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 2; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc_acc)[_ls], _scale2_9); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc_acc[_lp*2 + 0], dc_acc[_lp*2+1 + 0])); + dc_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(inv_a_frag[0]) + : "r"(a_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o_acc[0]), "=f"(o_acc[1]), "=f"(o_acc[2]), "=f"(o_acc[3]) + : "r"(dc_bf16[0]), "r"(dc_bf16[1]), "r"(inv_a_frag[0])); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o_acc[_lp*2 + 0], o_acc[_lp*2+1 + 0])); + o_bf16[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_4 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_5 = byte_off_4 ^ (byte_off_4 >> 7 & 7) << 4; + int o_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_5; + uint32_t _stmatrix_addr_10 = static_cast((unsigned long long)o_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x1.shared.b16 [%0], {%1};\n" + :: "r"(_stmatrix_addr_10), "r"(*reinterpret_cast(&o_bf16[0])) + : "memory"); + if (elect_sync()) { + mbarrier_arrive(prep_inv16_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_inv16_ready_addr + (prep_stage) * 8, _phase_prep_inv16_ready); + } + if (prep_local_warp == 0) { + int lane_row_1 = lane % 16; + int lane_col = lane / 16 * 8; + int byte_off_3 = (16 + lane_row_1) * 128 + (16 + lane_col) * 2; + int swizzled_off_4 = byte_off_3 ^ (byte_off_3 >> 7 & 7) << 4; + int d_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_4; + int byte_off_0_1 = (16 + lane_row_1) * 128 + lane_col * 2; + int swizzled_off_1_2 = byte_off_0_1 ^ (byte_off_0_1 >> 7 & 7) << 4; + int c_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_2; + int byte_off_2_2 = lane_row_1 * 128 + lane_col * 2; + int swizzled_off_3_2 = byte_off_2_2 ^ (byte_off_2_2 >> 7 & 7) << 4; + int a_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_2; + unsigned int d32_frag[4]; + unsigned int c32_frag[4]; + float dc32_acc[8]; + unsigned int dc32_bf16[4]; + unsigned int a32_frag[4]; + float o32_acc[8]; + unsigned int o32_bf16[4]; + unsigned int zero32_bf16[4]; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(d32_frag[0]), "=r"(d32_frag[1]), "=r"(d32_frag[2]), "=r"(d32_frag[3]) + : "r"(d_addr_1) + : "memory"); + int d_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_11 = static_cast((unsigned long long)d_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_11), "r"(*reinterpret_cast(&d32_frag[0])), "r"(*reinterpret_cast(&d32_frag[1])), "r"(*reinterpret_cast(&d32_frag[2])), "r"(*reinterpret_cast(&d32_frag[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(c32_frag[0]), "=r"(c32_frag[1]), "=r"(c32_frag[2]), "=r"(c32_frag[3]) + : "r"(c_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[0]), "=f"(dc32_acc[1]), "=f"(dc32_acc[2]), "=f"(dc32_acc[3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[0]), "r"(c32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[4]), "=f"(dc32_acc[(4) + 1]), "=f"(dc32_acc[(4) + 2]), "=f"(dc32_acc[(4) + 3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[2]), "r"(c32_frag[(2) + 1])); + const float2 _scale2_12 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc32_acc)[_ls], _scale2_12); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc32_acc[_lp*2 + 0], dc32_acc[_lp*2+1 + 0])); + dc32_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a32_frag[0]), "=r"(a32_frag[1]), "=r"(a32_frag[2]), "=r"(a32_frag[3]) + : "r"(a_addr_1) + : "memory"); + int a_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_13 = static_cast((unsigned long long)a_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_13), "r"(*reinterpret_cast(&a32_frag[0])), "r"(*reinterpret_cast(&a32_frag[1])), "r"(*reinterpret_cast(&a32_frag[2])), "r"(*reinterpret_cast(&a32_frag[3])) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[0]), "=f"(o32_acc[1]), "=f"(o32_acc[2]), "=f"(o32_acc[3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[0]), "r"(a32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[4]), "=f"(o32_acc[(4) + 1]), "=f"(o32_acc[(4) + 2]), "=f"(o32_acc[(4) + 3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[2]), "r"(a32_frag[(2) + 1])); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o32_acc[_lp*2 + 0], o32_acc[_lp*2+1 + 0])); + o32_bf16[_lp] = *(uint32_t*)&_bf2; + } + int o_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_14 = static_cast((unsigned long long)o_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_14), "r"(*reinterpret_cast(&o32_bf16[0])), "r"(*reinterpret_cast(&o32_bf16[1])), "r"(*reinterpret_cast(&o32_bf16[2])), "r"(*reinterpret_cast(&o32_bf16[3])) + : "memory"); + #pragma unroll + for (int zero_word = 0; zero_word < 4; zero_word++) { + zero32_bf16[zero_word] = 0; + } + int zero_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_15 = static_cast((unsigned long long)zero_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_15), "r"(*reinterpret_cast(&zero32_bf16[0])), "r"(*reinterpret_cast(&zero32_bf16[1])), "r"(*reinterpret_cast(&zero32_bf16[2])), "r"(*reinterpret_cast(&zero32_bf16[3])) + : "memory"); + } else if (prep_local_warp == 1) { + int stage_f32_0_1 = prep_stage * 10496; + float restore_scale_1 = smem_restore_factor_all[stage_f32_0_1 + 128]; + float restore_factor_1[8]; + int restore_segment_1 = lane & 15; + #pragma unroll + for (int restore_elem_2 = 0; restore_elem_2 < 8; restore_elem_2++) { + int restore_col_1 = restore_segment_1 * 8 + restore_elem_2; + restore_factor_1[restore_elem_2] = smem_restore_factor_all[stage_f32_0_1 + restore_col_1]; + } + #pragma unroll 1 + for (int restore_pass_1 = 0; restore_pass_1 < 4; restore_pass_1++) { + int restore_row_1 = restore_pass_1 * 2 + (lane >> 4); + float restore_qd_values_1[8]; + float restore_kd_values_1[8]; + float restore_ki_values_1[8]; + unsigned int packed_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_6[0])), "=r"(*reinterpret_cast(&packed_6[(0) + 1])), "=r"(*reinterpret_cast(&packed_6[(0) + 2])), "=r"(*reinterpret_cast(&packed_6[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_3[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_3[_pair * 2])[0]), "=f"((&packed_fp32_3[_pair * 2])[1]) + : "r"(packed_6[_pair + 0])); + } + #pragma unroll + for (int value_idx_6 = 0; value_idx_6 < 8; value_idx_6++) { + restore_qd_values_1[value_idx_6] = packed_fp32_3[value_idx_6]; + } + unsigned int packed_0_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_4[0])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_2[_pair * 2])[0]), "=f"((&packed_0_fp32_2[_pair * 2])[1]) + : "r"(packed_0_4[_pair + 0])); + } + #pragma unroll + for (int value_idx_7 = 0; value_idx_7 < 8; value_idx_7++) { + restore_kd_values_1[value_idx_7] = packed_0_fp32_2[value_idx_7]; + } + unsigned int packed_1_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_3[0])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32_1[_pair * 2])[0]), "=f"((&packed_1_fp32_1[_pair * 2])[1]) + : "r"(packed_1_3[_pair + 0])); + } + #pragma unroll + for (int value_idx_8 = 0; value_idx_8 < 8; value_idx_8++) { + restore_ki_values_1[value_idx_8] = packed_1_fp32_1[value_idx_8]; + } + float restore_kr_values_1[8]; + #pragma unroll + for (int restore_elem_3 = 0; restore_elem_3 < 8; restore_elem_3++) { + restore_kr_values_1[restore_elem_3] = restore_ki_values_1[restore_elem_3] * restore_factor_1[restore_elem_3]; + } + const float2 _scale2_16 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values_1)[_ls], _scale2_16); + const float2 _scale2_17 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values_1)[_ls], _scale2_17); + unsigned int packed_2_2[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values_1[_lp*2 + 0], restore_qd_values_1[_lp*2+1 + 0])); + packed_2_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_7 = 0; word_7 < 4; word_7++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_7 * 4)), "r"(packed_2_2[word_7])); + } + unsigned int packed_3_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values_1[_lp*2 + 0], restore_kd_values_1[_lp*2+1 + 0])); + packed_3_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_8 = 0; word_8 < 4; word_8++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_8 * 4)), "r"(packed_3_1[word_8])); + } + unsigned int packed_4_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values_1[_lp*2 + 0], restore_kr_values_1[_lp*2+1 + 0])); + packed_4_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_9 = 0; word_9 < 4; word_9++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_9 * 4)), "r"(packed_4_1[word_9])); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(qk_full_addr + (prep_stage) * 8); + } + } + for (int _advance = 0; _advance < 5; _advance++) { + prep_stage += 1; + if (prep_stage == 5) { prep_stage = 0; _phase_raw_inputs_free ^= 1; _phase_gate_raw_full ^= 1; _phase_smem_free ^= 1; _phase_qk_raw_full ^= 1; _phase_prep_diag_ready ^= 1; _phase_prep_inv16_ready ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + +// clang-format on diff --git a/csrc/kda/flashkda_bf16_fused_m128_binding.cu b/csrc/kda/flashkda_bf16_fused_m128_binding.cu new file mode 100644 index 00000000000..8fe63fe35af --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m128_binding.cu @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "flashkda_binding_common.cuh" + +// See the M64 binding for why the frozen standalone typedefs are isolated. +#define uint8_t flashkda_generated_uint8_t +#define uint16_t flashkda_generated_uint16_t +#define uint32_t flashkda_generated_uint32_t +#define uint64_t flashkda_generated_uint64_t +#define int32_t flashkda_generated_int32_t +#define int16_t flashkda_generated_int16_t +#include "flashkda_bf16_fused_m128.cu" +#undef uint8_t +#undef uint16_t +#undef uint32_t +#undef uint64_t +#undef int32_t +#undef int16_t + +namespace flashinfer { +namespace flash_kda { + +static_assert(THREADS == 1024); +static_assert(SMEM_TOTAL == 227328); + +void RunM128(TensorView q, TensorView k, TensorView v, TensorView g, TensorView beta, + TensorView beta_tma, TensorView A_log, TensorView dt_bias, TensorView cu_seqlens, + TensorView seq_order, TensorView initial_state, TensorView out, TensorView final_state, + TensorView descriptor_storage, int64_t prepare_descriptors, int64_t num_heads, + int64_t use_initial_state, int64_t store_final_state, double scale, double lower_bound, + int64_t cuda_stream) { + TVM_FFI_ICHECK(cuda_stream >= 0) << "cuda_stream must be a non-negative stream handle"; + TVM_FFI_ICHECK(q.device().device_type == kDLCUDA) << "q must be a CUDA tensor"; + const int32_t device_id = q.device().device_id; + ffi::CUDADeviceGuard device_guard(device_id); + CheckExactSm100a(device_id); + + const int64_t num_seqs = + CheckCommonInputs(q, k, v, g, beta, beta_tma, A_log, dt_bias, cu_seqlens, seq_order, + initial_state, out, final_state, descriptor_storage, prepare_descriptors, + num_heads, use_initial_state, store_final_state, scale, lower_bound); + + constexpr int32_t kSmemBytes = SMEM_TOTAL; + CheckDynamicSmemCapacity(device_id, kSmemBytes); + + CheckCuda(cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m128, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes), + "cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m128)"); + + const int64_t grid_x_i64 = num_seqs * num_heads; + TVM_FFI_ICHECK(grid_x_i64 > 0 && grid_x_i64 <= std::numeric_limits::max()) + << "M128 FlashKDA grid.x is out of range: " << grid_x_i64; + const dim3 grid(static_cast(grid_x_i64), 1, 1); + const dim3 block(THREADS, 1, 1); + const cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream)); + const TmaPointers tma = EncodeTmaPointers<128>(q, k, v, g, beta_tma, out, descriptor_storage, + prepare_descriptors, stream); + + kernel_flashkda_bf16_fused_m128<<>>( + reinterpret_cast<__nv_bfloat16*>(q.data_ptr()), tma.q, + reinterpret_cast<__nv_bfloat16*>(k.data_ptr()), tma.k, + reinterpret_cast<__nv_bfloat16*>(v.data_ptr()), tma.v, + reinterpret_cast<__nv_bfloat16*>(g.data_ptr()), tma.g, + reinterpret_cast<__nv_bfloat16*>(beta.data_ptr()), tma.beta, + reinterpret_cast(A_log.data_ptr()), reinterpret_cast(dt_bias.data_ptr()), + reinterpret_cast(cu_seqlens.data_ptr()), + reinterpret_cast(seq_order.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(initial_state.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), tma.out, + reinterpret_cast<__nv_bfloat16*>(final_state.data_ptr()), static_cast(num_heads), + static_cast(use_initial_state), static_cast(store_final_state), + static_cast(scale), static_cast(lower_bound)); + CheckCuda(cudaGetLastError(), "kernel_flashkda_bf16_fused_m128 launch"); +} + +} // namespace flash_kda +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, flashinfer::flash_kda::RunM128); diff --git a/csrc/kda/flashkda_bf16_fused_m64.cu b/csrc/kda/flashkda_bf16_fused_m64.cu new file mode 100644 index 00000000000..5f8541897e4 --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m64.cu @@ -0,0 +1,2566 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +// clang-format off +// Generated by tools/export-generated-programs (device kernel). +// Provenance: generated Loom schedule 'flashkda_bf16_fused_m64'; embedded in the host TU as flashkda_bf16_fused_m64_f0217be48b. +typedef unsigned char uint8_t; +typedef unsigned short uint16_t; +typedef unsigned int uint32_t; +typedef unsigned long long uint64_t; +typedef signed int int32_t; +typedef short int int16_t; + +#include + +#define LOOM_INF CUDART_INF_F +#define TMEM_NCOLS 256 +#define TMEM_TMEM_STATE_OFFSET 64 +#define TMEM_TMEM_STATE_INP_OFFSET 0 +#define TMEM_TMEM_U_ACC_OFFSET 224 +#define TMEM_TMEM_U2_INP_OFFSET 224 +#define TMEM_TMEM_U2_ACC_OFFSET 0 +#define TMEM_TMEM_OUT_OFFSET 192 +#define TMEM_TMEM_STATE_OUT_OFFSET 64 +#define NUM_CHUNK_PIPE_STAGES 5 +#define SMEM_SMEM_QD_OFF 1024 +#define SMEM_SMEM_QD_STAGE_BYTES 8192 +#define SMEM_SMEM_QD_STRIDE 41984 +#define SMEM_SMEM_G_RAW_OFF 1024 +#define SMEM_SMEM_G_RAW_STAGE_BYTES 8192 +#define SMEM_SMEM_G_RAW_STRIDE 41984 +#define SMEM_SMEM_G_RAW_ALL_OFF 1024 +#define SMEM_SMEM_G_RAW_ALL_STAGE_BYTES 176128 +#define SMEM_SMEM_G_RAW_ALL_STRIDE 176128 +#define SMEM_SMEM_KD_OFF 9216 +#define SMEM_SMEM_KD_STAGE_BYTES 8192 +#define SMEM_SMEM_KD_STRIDE 41984 +#define SMEM_SMEM_Q_RAW_PREFETCH_OFF 17408 +#define SMEM_SMEM_Q_RAW_PREFETCH_STAGE_BYTES 8192 +#define SMEM_SMEM_Q_RAW_PREFETCH_STRIDE 41984 +#define SMEM_SMEM_FINAL_TRANS_OFF 17408 +#define SMEM_SMEM_FINAL_TRANS_STAGE_BYTES 12288 +#define SMEM_SMEM_FINAL_TRANS_STRIDE 41984 +#define SMEM_SMEM_KR_TRANS_OFF 17408 +#define SMEM_SMEM_KR_TRANS_STAGE_BYTES 8192 +#define SMEM_SMEM_KR_TRANS_STRIDE 41984 +#define SMEM_SMEM_MQK_TRANS_OFF 25600 +#define SMEM_SMEM_MQK_TRANS_STAGE_BYTES 2048 +#define SMEM_SMEM_MQK_TRANS_STRIDE 41984 +#define SMEM_SMEM_INV_OFF 29696 +#define SMEM_SMEM_INV_STAGE_BYTES 2048 +#define SMEM_SMEM_INV_STRIDE 41984 +#define SMEM_SMEM_V_OFF 32384 +#define SMEM_SMEM_V_STAGE_BYTES 4096 +#define SMEM_SMEM_V_STRIDE 41984 +#define SMEM_SMEM_KI_OFF 17408 +#define SMEM_SMEM_KI_STAGE_BYTES 8192 +#define SMEM_SMEM_KI_STRIDE 41984 +#define SMEM_SMEM_GATE_OFF 25600 +#define SMEM_SMEM_GATE_STAGE_BYTES 16384 +#define SMEM_SMEM_GATE_STRIDE 41984 +#define SMEM_SMEM_BETA_RAW_OFF 41984 +#define SMEM_SMEM_BETA_RAW_STAGE_BYTES 512 +#define SMEM_SMEM_BETA_RAW_STRIDE 41984 +#define SMEM_SMEM_INV_WORK_OFF 32384 +#define SMEM_SMEM_INV_WORK_STAGE_BYTES 4096 +#define SMEM_SMEM_INV_WORK_STRIDE 41984 +#define SMEM_SMEM_OUT_OFF 210944 +#define SMEM_SMEM_OUT_STAGE_BYTES 4096 +#define SMEM_SMEM_OUT_STRIDE 4096 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_OFF 41984 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STAGE_BYTES 168452 +#define SMEM_SMEM_RESTORE_FACTOR_ALL_STRIDE 168452 +#define SMEM_SMEM_GT_PREFIX_ALL_OFF 41472 +#define SMEM_SMEM_GT_PREFIX_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_PREFIX_ALL_STRIDE 168448 +#define SMEM_SMEM_GT_ALL_OFF 31744 +#define SMEM_SMEM_GT_ALL_STAGE_BYTES 168448 +#define SMEM_SMEM_GT_ALL_STRIDE 168448 +#define SMEM_SMEM_PREP_BETA_ALL_OFF 42500 +#define SMEM_SMEM_PREP_BETA_ALL_STAGE_BYTES 168064 +#define SMEM_SMEM_PREP_BETA_ALL_STRIDE 168064 +#define SMEM_SMEM_GATE_RATE_ALL_OFF 42628 +#define SMEM_SMEM_GATE_RATE_ALL_STAGE_BYTES 167940 +#define SMEM_SMEM_GATE_RATE_ALL_STRIDE 167940 +#define SMEM_SMEM_GATE_ALL_OFF 25600 +#define SMEM_SMEM_GATE_ALL_STAGE_BYTES 184320 +#define SMEM_SMEM_GATE_ALL_STRIDE 184320 +#define SMEM_TOTAL 219136 +#define THREADS 1024 + +#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)); +} + + +__device__ __forceinline__ uint32_t mbarrier_try_wait(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ uint32_t mbarrier_try_wait_cluster(int mbar_addr, int phase) { + uint32_t token; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, P1;\n\t" + "}\n" + : "=r"(token) + : "r"(mbar_addr), "r"(phase) : "memory"); + return token; +} + +__device__ __forceinline__ void mbarrier_wait(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + 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, %2;\n\t" + "@P1 bra.uni DONE;\n\t" + "bra.uni LAB_WAIT;\n\t" + "DONE:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_cluster(int mbar_addr, int phase) { + uint32_t ticks = 0x989680; + asm volatile( + "{\n\t" + ".reg .pred P1;\n\t" + "LAB_WAIT_CLUSTER:\n\t" + "mbarrier.try_wait.parity.acquire.cluster.shared::cta.b64" + " P1, [%0], %1, %2;\n\t" + "@P1 bra.uni DONE_CLUSTER;\n\t" + "bra.uni LAB_WAIT_CLUSTER;\n\t" + "DONE_CLUSTER:\n\t" + "}\n" + :: "r"(mbar_addr), "r"(phase), "r"(ticks) : "memory"); +} + +__device__ __forceinline__ void mbarrier_wait_token(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait(mbar_addr, phase); + } +} + +__device__ __forceinline__ void mbarrier_wait_token_cluster(int mbar_addr, int phase, uint32_t token) { + if (token == 0) { + mbarrier_wait_cluster(mbar_addr, phase); + } +} + + +__device__ __forceinline__ void tcgen05_mma_f16( + int taddr, uint64_t a_desc, uint64_t b_desc, + uint32_t i_desc, int enable_input_d) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" + "tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p;\n\t" + "}\n" + :: "r"(taddr), "l"(a_desc), "l"(b_desc), + "r"(i_desc), "r"(enable_input_d)); +} + + +__device__ __forceinline__ uint64_t desc_encode(uint64_t x) { + return (x & 0x3FFFFULL) >> 4ULL; +} + + +__device__ __forceinline__ void mma_ts_step( + int taddr_out, int taddr_a, int b_lo, uint32_t b_dhi, + uint32_t i_desc, int enable_d) { + asm volatile( + "{\n\t" + ".reg .pred leader, p;\n\t" + ".reg .b32 dhi;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p, %5, 0;\n\t" + "mov.b32 dhi, %3;\n\t" + "mov.b64 db, {%2, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%1], db, %4, p;\n\t" + "}\n" + :: "r"(taddr_out), "r"(taddr_a), "r"(b_lo), "r"(b_dhi), + "r"(i_desc), "r"(enable_d)); +} + + +__device__ __forceinline__ void elect_commit(int mbar_addr) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "}\n" + :: "r"(mbar_addr)); +} + + +__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__ void tmem_ld_x32(float* dst, int tmem_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"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]), + "=f"(dst[16]), "=f"(dst[17]), "=f"(dst[18]), "=f"(dst[19]), + "=f"(dst[20]), "=f"(dst[21]), "=f"(dst[22]), "=f"(dst[23]), + "=f"(dst[24]), "=f"(dst[25]), "=f"(dst[26]), "=f"(dst[27]), + "=f"(dst[28]), "=f"(dst[29]), "=f"(dst[30]), "=f"(dst[31]) + : "r"(tmem_addr)); +} + + +__device__ __forceinline__ void tmem_ld_x16(float* dst, int tmem_addr) { + asm volatile( + "tcgen05.ld.sync.aligned.32x32b.x16.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7," + " %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=f"(dst[0]), "=f"(dst[1]), "=f"(dst[2]), "=f"(dst[3]), + "=f"(dst[4]), "=f"(dst[5]), "=f"(dst[6]), "=f"(dst[7]), + "=f"(dst[8]), "=f"(dst[9]), "=f"(dst[10]), "=f"(dst[11]), + "=f"(dst[12]), "=f"(dst[13]), "=f"(dst[14]), "=f"(dst[15]) + : "r"(tmem_addr)); +} + + +__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__ void mbarrier_init_pred(int mbar_addr, uint32_t count, uint32_t pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.init.shared::cta.b64 [%0], %1;\n\t" + "}\n" :: "r"(mbar_addr), "r"(count), "r"(pred)); +} + + +__device__ __forceinline__ float approx_exp2(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} + + +__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_inplace(float2* a, float2 b) { + asm("mul.rn.ftz.f32x2 %0, %0, %1;" + : "+l"(*(unsigned long long*)a) : "l"(*(unsigned long long*)&b)); +} + +__device__ __forceinline__ void add_f32x2_inplace(float2* a, float2 b) { + asm("add.rn.ftz.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)); +} + +__device__ __forceinline__ float2 add_f32x2(float2 a, float2 b) { + float2 r; + asm("add.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ float2 sub_f32x2(float2 a, float2 b) { + float2 r; + asm("sub.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +__device__ __forceinline__ void fma_scale_x32( + float* sv, const float2* scale2, const float2* neg_max2) +{ + float2* sv_2 = reinterpret_cast(sv); + #pragma unroll + for (int j = 0; j < 16; j++) + fma_f32x2_inplace(&sv_2[j], *scale2, *neg_max2); +} + +__device__ __forceinline__ float2 fma_f32x2(float2 a, float2 b, float2 c) { + float2 r; + asm("fma.rn.ftz.f32x2 %0, %1, %2, %3;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b), + "l"(*(unsigned long long*)&c)); + return r; +} + +__device__ __forceinline__ float2 mul_f32x2(float2 a, float2 b) { + float2 r; + asm("mul.rn.ftz.f32x2 %0, %1, %2;" + : "=l"(*(unsigned long long*)&r) + : "l"(*(unsigned long long*)&a), "l"(*(unsigned long long*)&b)); + return r; +} + +// ex2_emulation_f32x2 defined in softmax_frag_exp2_cast helper (or standalone) + + +__device__ __forceinline__ void elect_commit2(int mbar_addr0, int mbar_addr1) { + asm volatile( + "{\n\t" + ".reg .pred leader;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];\n\t" + "@leader tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%1];\n\t" + "}\n" + :: "r"(mbar_addr0), "r"(mbar_addr1) : "memory"); +} + + +__device__ __forceinline__ void fence_async_shared() { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); +} + + +__device__ __forceinline__ uint64_t make_smem_desc(int addr) { + const int SBO = 1024; + return desc_encode(addr) + | (desc_encode(SBO) << 32ULL) + | (1ULL << 46ULL) + | (2ULL << 61ULL); +} + + +__device__ __forceinline__ void tma_3d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.3d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3, %4}], [%5];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_2d_gmem2smem( + int dst, const void *tmap_ptr, int x, int y, int mbar_addr) { + asm volatile( + "cp.async.bulk.tensor.2d.shared::cta.global" + ".mbarrier::complete_tx::bytes" + " [%0], [%1, {%2, %3}], [%4];" + :: "r"(dst), "l"(tmap_ptr), "r"(x), "r"(y), + "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tma_4d_gmem2smem( + 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::cta.global" + ".mbarrier::complete_tx::bytes" + " [%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 tma_store_4d( + const void *tmap, int x, int y, int z, int w, unsigned smem_addr) { + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.tile.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" + :: "l"(tmap), "r"(x), "r"(y), "r"(z), "r"(w), "r"(smem_addr) : "memory"); +} + + +__device__ __forceinline__ void tcgen05_commit(int mbar_addr) { + asm volatile( + "tcgen05.commit.cta_group::1.mbarrier::arrive::one" + ".shared::cluster.b64 [%0];" + :: "r"(mbar_addr) : "memory"); +} + + +__device__ __forceinline__ void tmem_st_x8_u32(int addr, const uint32_t* src) { + asm volatile( + "tcgen05.st.sync.aligned.32x32b.x8.b32" + " [%0], {%1,%2,%3,%4,%5,%6,%7,%8};" + :: "r"(addr), + "r"(src[0]), "r"(src[1]), "r"(src[2]), "r"(src[3]), + "r"(src[4]), "r"(src[5]), "r"(src[6]), "r"(src[7])); +} + + +__device__ __forceinline__ uint32_t make_warp_uniform(uint32_t val) { + uint32_t result; + asm volatile("shfl.sync.idx.b32 %0, %1, 0, 0x1f, 0xffffffff;" + : "=r"(result) : "r"(val)); + return result; +} + +extern "C" { + +__global__ __launch_bounds__(1024) void +// FLASHINFER INTEGRATION BEGIN: allow exact state alias +kernel_flashkda_bf16_fused_m64(__nv_bfloat16* __restrict__ q, const void* __restrict__ q_tma, __nv_bfloat16* __restrict__ k, const void* __restrict__ k_tma, __nv_bfloat16* __restrict__ v, const void* __restrict__ v_tma, __nv_bfloat16* __restrict__ g, const void* __restrict__ g_tma, __nv_bfloat16* __restrict__ beta, const void* __restrict__ beta_tma, float* __restrict__ A_log, float* __restrict__ dt_bias, long long* __restrict__ cu_seqlens, int* __restrict__ seq_order, __nv_bfloat16* initial_state, __nv_bfloat16* __restrict__ out, const void* __restrict__ out_tma, __nv_bfloat16* final_state, int num_heads, int use_initial_state, int store_final_state, float scale, float lower_bound) +// FLASHINFER INTEGRATION END: allow exact state alias +{ + // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps + // CUDA kernel-start ordering does not acquire the tensor-map proxy. + // One thread acquires each 128-byte map; the CTA barrier publishes those + // acquires to every thread before any TMA instruction can use a map. + if (threadIdx.x == 0) { + asm volatile( + "fence.proxy.tensormap::generic.acquire.gpu [%0], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%1], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%2], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%3], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%4], 128;\n" + "fence.proxy.tensormap::generic.acquire.gpu [%5], 128;\n" + :: "l"(q_tma), "l"(k_tma), "l"(v_tma), "l"(g_tma), + "l"(beta_tma), "l"(out_tma) + : "memory"); + } + __syncthreads(); + // FLASHINFER INTEGRATION END: acquire global tensor maps + 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; + smem = (int)(unsigned long long)__cvta_generic_to_shared(smem_raw); + + const int bid = blockIdx.x; + const int num_bids = gridDim.x; + + // Kernel setup ops + __nv_bfloat16* smem_qd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_qd_addr = smem + 1024; + __nv_bfloat16* smem_g_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_addr = smem + 1024; + __nv_bfloat16* smem_g_raw_all = reinterpret_cast<__nv_bfloat16*>(smem_raw + 1024); + const int smem_g_raw_all_addr = smem + 1024; + __nv_bfloat16* smem_kd = reinterpret_cast<__nv_bfloat16*>(smem_raw + 9216); + const int smem_kd_addr = smem + 9216; + __nv_bfloat16* smem_q_raw_prefetch = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_q_raw_prefetch_addr = smem + 17408; + __nv_bfloat16* smem_final_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_final_trans_addr = smem + 17408; + __nv_bfloat16* smem_kr_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_kr_trans_addr = smem + 17408; + __nv_bfloat16* smem_mqk_trans = reinterpret_cast<__nv_bfloat16*>(smem_raw + 25600); + const int smem_mqk_trans_addr = smem + 25600; + __nv_bfloat16* smem_inv = reinterpret_cast<__nv_bfloat16*>(smem_raw + 29696); + const int smem_inv_addr = smem + 29696; + __nv_bfloat16* smem_v = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_v_addr = smem + 32384; + __nv_bfloat16* smem_ki = reinterpret_cast<__nv_bfloat16*>(smem_raw + 17408); + const int smem_ki_addr = smem + 17408; + float* smem_gate = reinterpret_cast(smem_raw + 25600); + const int smem_gate_addr = smem + 25600; + __nv_bfloat16* smem_beta_raw = reinterpret_cast<__nv_bfloat16*>(smem_raw + 41984); + const int smem_beta_raw_addr = smem + 41984; + __nv_bfloat16* smem_inv_work = reinterpret_cast<__nv_bfloat16*>(smem_raw + 32384); + const int smem_inv_work_addr = smem + 32384; + __nv_bfloat16* smem_out = reinterpret_cast<__nv_bfloat16*>(smem_raw + 210944); + const int smem_out_addr = smem + 210944; + float* smem_restore_factor_all = reinterpret_cast(smem_raw + 41984); + const int smem_restore_factor_all_addr = smem + 41984; + float* smem_gt_prefix_all = reinterpret_cast(smem_raw + 41472); + const int smem_gt_prefix_all_addr = smem + 41472; + float* smem_gt_all = reinterpret_cast(smem_raw + 31744); + const int smem_gt_all_addr = smem + 31744; + float* smem_prep_beta_all = reinterpret_cast(smem_raw + 42500); + const int smem_prep_beta_all_addr = smem + 42500; + float* smem_gate_rate_all = reinterpret_cast(smem_raw + 42628); + const int smem_gate_rate_all_addr = smem + 42628; + float* smem_gate_all = reinterpret_cast(smem_raw + 25600); + const int smem_gate_all_addr = smem + 25600; + + // Mbarrier init (17 groups, 77 barriers) + // Mbarriers at smem_raw[0..616) + + if (warp == 0) { + uint32_t leader = elect_sync(); + // --- pipeline 'chunk_pipe' --- + // qk_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 0, 1, leader); + mbarrier_init_pred(smem + 8, 1, leader); + mbarrier_init_pred(smem + 16, 1, leader); + mbarrier_init_pred(smem + 24, 1, leader); + mbarrier_init_pred(smem + 32, 1, leader); + // gate_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 40, 1, leader); + mbarrier_init_pred(smem + 48, 1, leader); + mbarrier_init_pred(smem + 56, 1, leader); + mbarrier_init_pred(smem + 64, 1, leader); + mbarrier_init_pred(smem + 72, 1, leader); + // qk_raw_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 80, 1, leader); + mbarrier_init_pred(smem + 88, 1, leader); + mbarrier_init_pred(smem + 96, 1, leader); + mbarrier_init_pred(smem + 104, 1, leader); + mbarrier_init_pred(smem + 112, 1, leader); + // v_full: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 120, 1, leader); + mbarrier_init_pred(smem + 128, 1, leader); + mbarrier_init_pred(smem + 136, 1, leader); + mbarrier_init_pred(smem + 144, 1, leader); + mbarrier_init_pred(smem + 152, 1, leader); + // v_free: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 160, 4, leader); + mbarrier_init_pred(smem + 168, 4, leader); + mbarrier_init_pred(smem + 176, 4, leader); + mbarrier_init_pred(smem + 184, 4, leader); + mbarrier_init_pred(smem + 192, 4, leader); + // smem_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 200, 1, leader); + mbarrier_init_pred(smem + 208, 1, leader); + mbarrier_init_pred(smem + 216, 1, leader); + mbarrier_init_pred(smem + 224, 1, leader); + mbarrier_init_pred(smem + 232, 1, leader); + // raw_inputs_free: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 240, 1, leader); + mbarrier_init_pred(smem + 248, 1, leader); + mbarrier_init_pred(smem + 256, 1, leader); + mbarrier_init_pred(smem + 264, 1, leader); + mbarrier_init_pred(smem + 272, 1, leader); + // state_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 280, 4, leader); + mbarrier_init_pred(smem + 288, 4, leader); + mbarrier_init_pred(smem + 296, 4, leader); + mbarrier_init_pred(smem + 304, 4, leader); + mbarrier_init_pred(smem + 312, 4, leader); + // old_out_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 320, 1, leader); + mbarrier_init_pred(smem + 328, 1, leader); + mbarrier_init_pred(smem + 336, 1, leader); + mbarrier_init_pred(smem + 344, 1, leader); + mbarrier_init_pred(smem + 352, 1, leader); + // u_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 360, 4, leader); + mbarrier_init_pred(smem + 368, 4, leader); + mbarrier_init_pred(smem + 376, 4, leader); + mbarrier_init_pred(smem + 384, 4, leader); + mbarrier_init_pred(smem + 392, 4, leader); + // u2_acc_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 400, 1, leader); + mbarrier_init_pred(smem + 408, 1, leader); + mbarrier_init_pred(smem + 416, 1, leader); + mbarrier_init_pred(smem + 424, 1, leader); + mbarrier_init_pred(smem + 432, 1, leader); + // u2_inp_ready: 5 barriers, init_count=4 + mbarrier_init_pred(smem + 440, 4, leader); + mbarrier_init_pred(smem + 448, 4, leader); + mbarrier_init_pred(smem + 456, 4, leader); + mbarrier_init_pred(smem + 464, 4, leader); + mbarrier_init_pred(smem + 472, 4, leader); + // final_ready: 5 barriers, init_count=1 + mbarrier_init_pred(smem + 480, 1, leader); + mbarrier_init_pred(smem + 488, 1, leader); + mbarrier_init_pred(smem + 496, 1, leader); + mbarrier_init_pred(smem + 504, 1, leader); + mbarrier_init_pred(smem + 512, 1, leader); + // out_empty: 1 barriers, init_count=4 + mbarrier_init_pred(smem + 520, 4, leader); + // tmem_dealloc_ready: 1 barriers, init_count=2 + mbarrier_init_pred(smem + 528, 2, leader); + // prep_diag_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 536, 2, leader); + mbarrier_init_pred(smem + 544, 2, leader); + mbarrier_init_pred(smem + 552, 2, leader); + mbarrier_init_pred(smem + 560, 2, leader); + mbarrier_init_pred(smem + 568, 2, leader); + // prep_inv16_ready: 5 barriers, init_count=2 + mbarrier_init_pred(smem + 576, 2, leader); + mbarrier_init_pred(smem + 584, 2, leader); + mbarrier_init_pred(smem + 592, 2, leader); + mbarrier_init_pred(smem + 600, 2, leader); + mbarrier_init_pred(smem + 608, 2, leader); + asm volatile("fence.mbarrier_init.release.cluster;"); + } + + __syncthreads(); + + // TMEM alloc (256 columns, 256 used) + volatile int* tmem_addr_storage = (volatile int*)(smem_raw + 616); + if (warp == 0) { + int _tmem_hold = smem + 616; + asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;" :: "r"(_tmem_hold), "r"(256) : "memory"); + } + + __syncthreads(); + asm volatile("tcgen05.fence::after_thread_sync;"); + + const int mbar_base = smem; + #define qk_full_addr (mbar_base + 0) + #define gate_raw_full_addr (mbar_base + 40) + #define qk_raw_full_addr (mbar_base + 80) + #define v_full_addr (mbar_base + 120) + #define v_free_addr (mbar_base + 160) + #define smem_free_addr (mbar_base + 200) + #define raw_inputs_free_addr (mbar_base + 240) + #define state_inp_ready_addr (mbar_base + 280) + #define old_out_ready_addr (mbar_base + 320) + #define u_inp_ready_addr (mbar_base + 360) + #define u2_acc_ready_addr (mbar_base + 400) + #define u2_inp_ready_addr (mbar_base + 440) + #define final_ready_addr (mbar_base + 480) + #define out_empty_addr (mbar_base + 520) + #define tmem_dealloc_ready_addr (mbar_base + 528) + #define prep_diag_ready_addr (mbar_base + 536) + #define prep_inv16_ready_addr (mbar_base + 576) + const int taddr = tmem_addr_storage[0]; + + // Kernel post-init ops + const int tmem_tmem_state = taddr + 64; + const int tmem_tmem_state_inp = taddr; + const int tmem_tmem_u_acc = taddr + 224; + const int tmem_tmem_u2_inp = taddr + 224; + const int tmem_tmem_u2_acc = taddr; + const int tmem_tmem_out = taddr + 192; + const int tmem_tmem_state_out = taddr + 64; + + // ---- Register redistribution for WGs split across roles ---- + // Dec phase frees registers before any WG attempts inc. + if (warp >= 8 && warp <= 11) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + } + + // ---- Role: compute ---- + if (warp <= 3) { + asm volatile("setmaxnreg.inc.sync.aligned.u32 168;"); + { // compute_main + int split_task_idx = blockIdx.x; + int task_idx = split_task_idx / 2; + int value_split_idx = split_task_idx % 2; + int value_row_offset = value_split_idx * 64; + int seq_idx = seq_order[task_idx / num_heads]; + int head_idx = task_idx % num_heads; + long long bos = cu_seqlens[seq_idx]; + long long eos = cu_seqlens[seq_idx + 1]; + int seq_len = (int)(eos - bos); + int num_chunks = (seq_len + 32 - 1) / 32; + int warp_in_wg = warp % 4; + const int tmem_row_base = warp_in_wg * 32 << 16; + int lane_quad = lane & 3; + int local_row_top = warp_in_wg * 16 + lane / 4; + int local_row_bot = local_row_top + 8; + int state_row_top = value_row_offset + local_row_top; + int state_row_bot = value_row_offset + local_row_bot; + int warp_id_in_role = (warp - 0); + int compute_local_warp = warp_id_in_role; + long long state_head_base = ((long long)seq_idx * (long long)num_heads + (long long)head_idx) * 128 * 128; + long long state_base_top = state_head_base + (long long)state_row_top * 128; + long long state_base_bot = state_head_base + (long long)state_row_bot * 128; + #pragma unroll + for (int state_col_half = 0; state_col_half < 2; state_col_half++) { + float state_init[32]; + state_init[0] = 0.0f; + state_init[1] = 0.0f; + state_init[2] = 0.0f; + state_init[3] = 0.0f; + state_init[4] = 0.0f; + state_init[5] = 0.0f; + state_init[6] = 0.0f; + state_init[7] = 0.0f; + state_init[8] = 0.0f; + state_init[9] = 0.0f; + state_init[10] = 0.0f; + state_init[11] = 0.0f; + state_init[12] = 0.0f; + state_init[13] = 0.0f; + state_init[14] = 0.0f; + state_init[15] = 0.0f; + state_init[16] = 0.0f; + state_init[17] = 0.0f; + state_init[18] = 0.0f; + state_init[19] = 0.0f; + state_init[20] = 0.0f; + state_init[21] = 0.0f; + state_init[22] = 0.0f; + state_init[23] = 0.0f; + state_init[24] = 0.0f; + state_init[25] = 0.0f; + state_init[26] = 0.0f; + state_init[27] = 0.0f; + state_init[28] = 0.0f; + state_init[29] = 0.0f; + state_init[30] = 0.0f; + state_init[31] = 0.0f; + if (use_initial_state != 0) { + #pragma unroll + for (int state_col_group = 0; state_col_group < 8; state_col_group++) { + int state_col_pair = state_col_half * 64 + state_col_group * 8 + lane_quad * 2; + const int state_reg_base = state_col_group * 4; + state_init[state_reg_base] = (float)initial_state[state_base_top + (long long)state_col_pair]; + state_init[state_reg_base + 1] = (float)initial_state[state_base_top + (long long)state_col_pair + 1]; + state_init[state_reg_base + 2] = (float)initial_state[state_base_bot + (long long)state_col_pair]; + state_init[state_reg_base + 3] = (float)initial_state[state_base_bot + (long long)state_col_pair + 1]; + } + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x8.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"(taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_half * 64)), "r"(*reinterpret_cast(&state_init[0])), "r"(*reinterpret_cast(&state_init[1])), "r"(*reinterpret_cast(&state_init[2])), "r"(*reinterpret_cast(&state_init[3])), "r"(*reinterpret_cast(&state_init[4])), "r"(*reinterpret_cast(&state_init[5])), "r"(*reinterpret_cast(&state_init[6])), "r"(*reinterpret_cast(&state_init[7])), "r"(*reinterpret_cast(&state_init[8])), "r"(*reinterpret_cast(&state_init[9])), "r"(*reinterpret_cast(&state_init[10])), "r"(*reinterpret_cast(&state_init[11])), "r"(*reinterpret_cast(&state_init[12])), "r"(*reinterpret_cast(&state_init[13])), "r"(*reinterpret_cast(&state_init[14])), "r"(*reinterpret_cast(&state_init[15])), "r"(*reinterpret_cast(&state_init[16])), "r"(*reinterpret_cast(&state_init[17])), "r"(*reinterpret_cast(&state_init[18])), "r"(*reinterpret_cast(&state_init[19])), "r"(*reinterpret_cast(&state_init[20])), "r"(*reinterpret_cast(&state_init[21])), "r"(*reinterpret_cast(&state_init[22])), "r"(*reinterpret_cast(&state_init[23])), "r"(*reinterpret_cast(&state_init[24])), "r"(*reinterpret_cast(&state_init[25])), "r"(*reinterpret_cast(&state_init[26])), "r"(*reinterpret_cast(&state_init[27])), "r"(*reinterpret_cast(&state_init[28])), "r"(*reinterpret_cast(&state_init[29])), "r"(*reinterpret_cast(&state_init[30])), "r"(*reinterpret_cast(&state_init[31])) + : "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + unsigned int compute_stage = 0; + unsigned int _phase_qk_full = 0; + unsigned int _phase_v_full = 0; + unsigned int _phase_old_out_ready = 0; + unsigned int _phase_u2_acc_ready = 0; + unsigned int _phase_final_ready = 0; + #pragma unroll 1 + for (int chunk_idx = 0; chunk_idx < num_chunks; chunk_idx++) { + mbarrier_wait(qk_full_addr + (compute_stage) * 8, _phase_qk_full); + #pragma unroll + for (int state_col_half_1 = 0; state_col_half_1 < 2; state_col_half_1++) { + int state_addr = taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_half_1 * 64); + float _tmem_load_0[32]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x8.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"(*reinterpret_cast(&_tmem_load_0[0])), "=r"(*reinterpret_cast(&_tmem_load_0[1])), "=r"(*reinterpret_cast(&_tmem_load_0[2])), "=r"(*reinterpret_cast(&_tmem_load_0[3])), "=r"(*reinterpret_cast(&_tmem_load_0[4])), "=r"(*reinterpret_cast(&_tmem_load_0[5])), "=r"(*reinterpret_cast(&_tmem_load_0[6])), "=r"(*reinterpret_cast(&_tmem_load_0[7])), "=r"(*reinterpret_cast(&_tmem_load_0[8])), "=r"(*reinterpret_cast(&_tmem_load_0[9])), "=r"(*reinterpret_cast(&_tmem_load_0[10])), "=r"(*reinterpret_cast(&_tmem_load_0[11])), "=r"(*reinterpret_cast(&_tmem_load_0[12])), "=r"(*reinterpret_cast(&_tmem_load_0[13])), "=r"(*reinterpret_cast(&_tmem_load_0[14])), "=r"(*reinterpret_cast(&_tmem_load_0[15])), "=r"(*reinterpret_cast(&_tmem_load_0[16])), "=r"(*reinterpret_cast(&_tmem_load_0[17])), "=r"(*reinterpret_cast(&_tmem_load_0[18])), "=r"(*reinterpret_cast(&_tmem_load_0[19])), "=r"(*reinterpret_cast(&_tmem_load_0[20])), "=r"(*reinterpret_cast(&_tmem_load_0[21])), "=r"(*reinterpret_cast(&_tmem_load_0[22])), "=r"(*reinterpret_cast(&_tmem_load_0[23])), "=r"(*reinterpret_cast(&_tmem_load_0[24])), "=r"(*reinterpret_cast(&_tmem_load_0[25])), "=r"(*reinterpret_cast(&_tmem_load_0[26])), "=r"(*reinterpret_cast(&_tmem_load_0[27])), "=r"(*reinterpret_cast(&_tmem_load_0[28])), "=r"(*reinterpret_cast(&_tmem_load_0[29])), "=r"(*reinterpret_cast(&_tmem_load_0[30])), "=r"(*reinterpret_cast(&_tmem_load_0[31])) + : "r"(state_addr) + : "memory"); + uint32_t _tmem_load_0_bf16[16]; + #pragma unroll + for (int _lp = 0; _lp < 16; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_0[_lp*2 + 0], _tmem_load_0[_lp*2+1 + 0])); + _tmem_load_0_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x8.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, %16};" + :: "r"(taddr + (unsigned int)tmem_row_base + (unsigned int)(state_col_half_1 * 32)), "r"(*reinterpret_cast(&_tmem_load_0_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[7])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[8])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[9])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[10])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[11])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[12])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[13])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[14])), "r"(*reinterpret_cast(&_tmem_load_0_bf16[15])) + : "memory"); + #pragma unroll + for (int state_col_group_1 = 0; state_col_group_1 < 8; state_col_group_1++) { + int state_col_pair_1 = state_col_half_1 * 64 + state_col_group_1 * 8 + lane_quad * 2; + const int state_reg_base_1 = state_col_group_1 * 4; + float state_scale_0 = smem_gt_all[compute_stage * 10496 + (unsigned int)state_col_pair_1]; + float state_scale_1 = smem_gt_all[compute_stage * 10496 + (unsigned int)state_col_pair_1 + 1]; + _tmem_load_0[state_reg_base_1] = _tmem_load_0[state_reg_base_1] * state_scale_0; + _tmem_load_0[state_reg_base_1 + 1] = _tmem_load_0[state_reg_base_1 + 1] * state_scale_1; + _tmem_load_0[state_reg_base_1 + 2] = _tmem_load_0[state_reg_base_1 + 2] * state_scale_0; + _tmem_load_0[state_reg_base_1 + 3] = _tmem_load_0[state_reg_base_1 + 3] * state_scale_1; + } + asm volatile( + "tcgen05.st.sync.aligned.16x256b.x8.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"(state_addr), "r"(*reinterpret_cast(&_tmem_load_0[0])), "r"(*reinterpret_cast(&_tmem_load_0[1])), "r"(*reinterpret_cast(&_tmem_load_0[2])), "r"(*reinterpret_cast(&_tmem_load_0[3])), "r"(*reinterpret_cast(&_tmem_load_0[4])), "r"(*reinterpret_cast(&_tmem_load_0[5])), "r"(*reinterpret_cast(&_tmem_load_0[6])), "r"(*reinterpret_cast(&_tmem_load_0[7])), "r"(*reinterpret_cast(&_tmem_load_0[8])), "r"(*reinterpret_cast(&_tmem_load_0[9])), "r"(*reinterpret_cast(&_tmem_load_0[10])), "r"(*reinterpret_cast(&_tmem_load_0[11])), "r"(*reinterpret_cast(&_tmem_load_0[12])), "r"(*reinterpret_cast(&_tmem_load_0[13])), "r"(*reinterpret_cast(&_tmem_load_0[14])), "r"(*reinterpret_cast(&_tmem_load_0[15])), "r"(*reinterpret_cast(&_tmem_load_0[16])), "r"(*reinterpret_cast(&_tmem_load_0[17])), "r"(*reinterpret_cast(&_tmem_load_0[18])), "r"(*reinterpret_cast(&_tmem_load_0[19])), "r"(*reinterpret_cast(&_tmem_load_0[20])), "r"(*reinterpret_cast(&_tmem_load_0[21])), "r"(*reinterpret_cast(&_tmem_load_0[22])), "r"(*reinterpret_cast(&_tmem_load_0[23])), "r"(*reinterpret_cast(&_tmem_load_0[24])), "r"(*reinterpret_cast(&_tmem_load_0[25])), "r"(*reinterpret_cast(&_tmem_load_0[26])), "r"(*reinterpret_cast(&_tmem_load_0[27])), "r"(*reinterpret_cast(&_tmem_load_0[28])), "r"(*reinterpret_cast(&_tmem_load_0[29])), "r"(*reinterpret_cast(&_tmem_load_0[30])), "r"(*reinterpret_cast(&_tmem_load_0[31])) + : "memory"); + } + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(state_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(v_full_addr + (compute_stage) * 8, _phase_v_full); + mbarrier_wait(old_out_ready_addr + (compute_stage) * 8, _phase_old_out_ready); + float _tmem_load_1[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_1[0])), "=r"(*reinterpret_cast(&_tmem_load_1[1])), "=r"(*reinterpret_cast(&_tmem_load_1[2])), "=r"(*reinterpret_cast(&_tmem_load_1[3])), "=r"(*reinterpret_cast(&_tmem_load_1[4])), "=r"(*reinterpret_cast(&_tmem_load_1[5])), "=r"(*reinterpret_cast(&_tmem_load_1[6])), "=r"(*reinterpret_cast(&_tmem_load_1[7])), "=r"(*reinterpret_cast(&_tmem_load_1[8])), "=r"(*reinterpret_cast(&_tmem_load_1[9])), "=r"(*reinterpret_cast(&_tmem_load_1[10])), "=r"(*reinterpret_cast(&_tmem_load_1[11])), "=r"(*reinterpret_cast(&_tmem_load_1[12])), "=r"(*reinterpret_cast(&_tmem_load_1[13])), "=r"(*reinterpret_cast(&_tmem_load_1[14])), "=r"(*reinterpret_cast(&_tmem_load_1[15])) + : "r"(taddr + 224 + (unsigned int)tmem_row_base) + : "memory"); + float residual_values[16]; + int v_stage_addr = smem_v_addr + compute_stage * 41984; + unsigned int v_ld_bits[2]; + #pragma unroll + for (int token_group = 0; token_group < 4; token_group++) { + int token_pair = token_group * 8 + lane_quad * 2; + const int residual_reg_base = token_group * 4; + float beta_0 = smem_prep_beta_all[compute_stage * 10496 + (unsigned int)token_pair]; + float beta_1 = smem_prep_beta_all[compute_stage * 10496 + (unsigned int)token_pair + 1]; + int v_ld_matrix = lane / 8 & 1; + int v_ld_token = token_group * 8 + (lane & 7); + int v_ld_row = warp_in_wg * 16 + v_ld_matrix * 8; + int v_ld_row_addr = v_stage_addr + v_ld_token * 64 * 2; + int v_ld_addr = (v_ld_row_addr + (v_ld_row * 2 ^ (v_ld_row_addr >> 7 & 7) << 4)); + asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n" + : "=r"(v_ld_bits[0]), "=r"(v_ld_bits[1]) + : "r"(v_ld_addr) + : "memory"); + float v_ld_bits_fp32[4]; + #pragma unroll + for (int _pair = 0; _pair < 2; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&v_ld_bits_fp32[_pair * 2])[0]), "=f"((&v_ld_bits_fp32[_pair * 2])[1]) + : "r"(v_ld_bits[_pair + 0])); + } + residual_values[residual_reg_base] = (v_ld_bits_fp32[0] - _tmem_load_1[residual_reg_base]) * beta_0; + residual_values[residual_reg_base + 1] = (v_ld_bits_fp32[1] - _tmem_load_1[residual_reg_base + 1]) * beta_1; + residual_values[residual_reg_base + 2] = (v_ld_bits_fp32[2] - _tmem_load_1[residual_reg_base + 2]) * beta_0; + residual_values[residual_reg_base + 3] = (v_ld_bits_fp32[3] - _tmem_load_1[residual_reg_base + 3]) * beta_1; + } + uint32_t residual_values_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(residual_values[_lp*2 + 0], residual_values[_lp*2+1 + 0])); + residual_values_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x4.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 224 + (unsigned int)tmem_row_base), "r"(*reinterpret_cast(&residual_values_bf16[0])), "r"(*reinterpret_cast(&residual_values_bf16[1])), "r"(*reinterpret_cast(&residual_values_bf16[2])), "r"(*reinterpret_cast(&residual_values_bf16[3])), "r"(*reinterpret_cast(&residual_values_bf16[4])), "r"(*reinterpret_cast(&residual_values_bf16[5])), "r"(*reinterpret_cast(&residual_values_bf16[6])), "r"(*reinterpret_cast(&residual_values_bf16[7])) + : "memory"); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(v_free_addr + (compute_stage) * 8); + mbarrier_arrive(u_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(u2_acc_ready_addr + (compute_stage) * 8, _phase_u2_acc_ready); + float _tmem_load_2[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_2[0])), "=r"(*reinterpret_cast(&_tmem_load_2[1])), "=r"(*reinterpret_cast(&_tmem_load_2[2])), "=r"(*reinterpret_cast(&_tmem_load_2[3])), "=r"(*reinterpret_cast(&_tmem_load_2[4])), "=r"(*reinterpret_cast(&_tmem_load_2[5])), "=r"(*reinterpret_cast(&_tmem_load_2[6])), "=r"(*reinterpret_cast(&_tmem_load_2[7])), "=r"(*reinterpret_cast(&_tmem_load_2[8])), "=r"(*reinterpret_cast(&_tmem_load_2[9])), "=r"(*reinterpret_cast(&_tmem_load_2[10])), "=r"(*reinterpret_cast(&_tmem_load_2[11])), "=r"(*reinterpret_cast(&_tmem_load_2[12])), "=r"(*reinterpret_cast(&_tmem_load_2[13])), "=r"(*reinterpret_cast(&_tmem_load_2[14])), "=r"(*reinterpret_cast(&_tmem_load_2[15])) + : "r"(taddr + (unsigned int)tmem_row_base) + : "memory"); + uint32_t _tmem_load_2_bf16[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_2[_lp*2 + 0], _tmem_load_2[_lp*2+1 + 0])); + _tmem_load_2_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile( + "tcgen05.st.sync.aligned.16x128b.x4.b32" + " [%0], {%1, %2, %3, %4, %5, %6, %7, %8};" + :: "r"(taddr + 224 + (unsigned int)tmem_row_base), "r"(*reinterpret_cast(&_tmem_load_2_bf16[0])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[1])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[2])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[3])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[4])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[5])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[6])), "r"(*reinterpret_cast(&_tmem_load_2_bf16[7])) + : "memory"); + asm volatile("tcgen05.wait::st.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(u2_inp_ready_addr + (compute_stage) * 8); + } + mbarrier_wait(final_ready_addr + (compute_stage) * 8, _phase_final_ready); + compute_stage += 1; + if (compute_stage == 5) { compute_stage = 0; _phase_qk_full ^= 1; _phase_v_full ^= 1; _phase_old_out_ready ^= 1; _phase_u2_acc_ready ^= 1; _phase_final_ready ^= 1; } + } + if (store_final_state != 0) { + #pragma unroll + for (int state_col_half_2 = 0; state_col_half_2 < 2; state_col_half_2++) { + float _tmem_load_3[32]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x8.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"(*reinterpret_cast(&_tmem_load_3[0])), "=r"(*reinterpret_cast(&_tmem_load_3[1])), "=r"(*reinterpret_cast(&_tmem_load_3[2])), "=r"(*reinterpret_cast(&_tmem_load_3[3])), "=r"(*reinterpret_cast(&_tmem_load_3[4])), "=r"(*reinterpret_cast(&_tmem_load_3[5])), "=r"(*reinterpret_cast(&_tmem_load_3[6])), "=r"(*reinterpret_cast(&_tmem_load_3[7])), "=r"(*reinterpret_cast(&_tmem_load_3[8])), "=r"(*reinterpret_cast(&_tmem_load_3[9])), "=r"(*reinterpret_cast(&_tmem_load_3[10])), "=r"(*reinterpret_cast(&_tmem_load_3[11])), "=r"(*reinterpret_cast(&_tmem_load_3[12])), "=r"(*reinterpret_cast(&_tmem_load_3[13])), "=r"(*reinterpret_cast(&_tmem_load_3[14])), "=r"(*reinterpret_cast(&_tmem_load_3[15])), "=r"(*reinterpret_cast(&_tmem_load_3[16])), "=r"(*reinterpret_cast(&_tmem_load_3[17])), "=r"(*reinterpret_cast(&_tmem_load_3[18])), "=r"(*reinterpret_cast(&_tmem_load_3[19])), "=r"(*reinterpret_cast(&_tmem_load_3[20])), "=r"(*reinterpret_cast(&_tmem_load_3[21])), "=r"(*reinterpret_cast(&_tmem_load_3[22])), "=r"(*reinterpret_cast(&_tmem_load_3[23])), "=r"(*reinterpret_cast(&_tmem_load_3[24])), "=r"(*reinterpret_cast(&_tmem_load_3[25])), "=r"(*reinterpret_cast(&_tmem_load_3[26])), "=r"(*reinterpret_cast(&_tmem_load_3[27])), "=r"(*reinterpret_cast(&_tmem_load_3[28])), "=r"(*reinterpret_cast(&_tmem_load_3[29])), "=r"(*reinterpret_cast(&_tmem_load_3[30])), "=r"(*reinterpret_cast(&_tmem_load_3[31])) + : "r"(taddr + 64 + (unsigned int)tmem_row_base + (unsigned int)(state_col_half_2 * 64)) + : "memory"); + #pragma unroll + for (int state_col_group_2 = 0; state_col_group_2 < 8; state_col_group_2++) { + int state_col_pair_2 = state_col_half_2 * 64 + state_col_group_2 * 8 + lane_quad * 2; + const int state_reg_base_2 = state_col_group_2 * 4; + final_state[state_base_top + (long long)state_col_pair_2] = _tmem_load_3[state_reg_base_2]; + final_state[state_base_top + (long long)state_col_pair_2 + 1] = _tmem_load_3[state_reg_base_2 + 1]; + final_state[state_base_bot + (long long)state_col_pair_2] = _tmem_load_3[state_reg_base_2 + 2]; + final_state[state_base_bot + (long long)state_col_pair_2 + 1] = _tmem_load_3[state_reg_base_2 + 3]; + } + } + } + asm volatile("barrier.sync 10, 128;" ::: "memory"); + if (compute_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: epilogue ---- + } else if (warp >= 4 && warp <= 7) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // epilogue_main + int split_task_idx_1 = blockIdx.x; + int task_idx_1 = split_task_idx_1 / 2; + int value_split_idx_1 = split_task_idx_1 % 2; + int value_row_offset_1 = value_split_idx_1 * 64; + int seq_idx_1 = seq_order[task_idx_1 / num_heads]; + int head_idx_1 = task_idx_1 % num_heads; + long long bos_1 = cu_seqlens[seq_idx_1]; + long long eos_1 = cu_seqlens[seq_idx_1 + 1]; + int seq_len_1 = (int)(eos_1 - bos_1); + int num_chunks_1 = (seq_len_1 + 32 - 1) / 32; + int warp_id_in_role_1 = (warp - 4); + int epilogue_local_warp = warp_id_in_role_1; + int warp_in_wg_1 = warp % 4; + const int tmem_row_base_1 = warp_in_wg_1 * 32 << 16; + int lane_quad_1 = lane & 3; + int local_row_top_1 = warp_in_wg_1 * 16 + lane / 4; + int local_row_bot_1 = local_row_top_1 + 8; + int state_row_top_1 = value_row_offset_1 + local_row_top_1; + int state_row_bot_1 = value_row_offset_1 + local_row_bot_1; + unsigned int epilogue_stage = 0; + unsigned int output_stage = 0; + unsigned int _phase_final_ready_1 = 0; + #pragma unroll 1 + for (int chunk_idx_1 = 0; chunk_idx_1 < num_chunks_1; chunk_idx_1++) { + mbarrier_wait(final_ready_addr + (epilogue_stage) * 8, _phase_final_ready_1); + int chunk_is_full = ((seq_len_1 >= (chunk_idx_1 + 1) * 32) ? 1 : 0); + if (chunk_is_full != 0) { + float _tmem_load_4[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_4[0])), "=r"(*reinterpret_cast(&_tmem_load_4[1])), "=r"(*reinterpret_cast(&_tmem_load_4[2])), "=r"(*reinterpret_cast(&_tmem_load_4[3])), "=r"(*reinterpret_cast(&_tmem_load_4[4])), "=r"(*reinterpret_cast(&_tmem_load_4[5])), "=r"(*reinterpret_cast(&_tmem_load_4[6])), "=r"(*reinterpret_cast(&_tmem_load_4[7])), "=r"(*reinterpret_cast(&_tmem_load_4[8])), "=r"(*reinterpret_cast(&_tmem_load_4[9])), "=r"(*reinterpret_cast(&_tmem_load_4[10])), "=r"(*reinterpret_cast(&_tmem_load_4[11])), "=r"(*reinterpret_cast(&_tmem_load_4[12])), "=r"(*reinterpret_cast(&_tmem_load_4[13])), "=r"(*reinterpret_cast(&_tmem_load_4[14])), "=r"(*reinterpret_cast(&_tmem_load_4[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + if (epilogue_local_warp == 0) { + if (chunk_idx_1 >= 2) { + asm volatile("cp.async.bulk.wait_group.read 1;"); + } + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + int out_stage_addr = smem_out_addr + output_stage * 4096; + unsigned int out_packed[8]; + #pragma unroll + for (int _lp = 0; _lp < 8; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(_tmem_load_4[_lp*2 + 0], _tmem_load_4[_lp*2+1 + 0])); + out_packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int token_group_1 = 0; token_group_1 < 2; token_group_1++) { + int mtx_idx = lane / 8; + int row_addr = lane & 7; + int dim_base = epilogue_local_warp * 16 + (mtx_idx & 1) * 8; + int token_base = token_group_1 * 16 + mtx_idx / 2 * 8; + int token_addr = token_base + row_addr; + int token_pair_1 = token_addr / 2; + int token_parity = token_addr & 1; + int raw_row = token_pair_1; + int raw_col = (dim_base & 63 ^ (token_pair_1 & 3) << 4 ^ token_parity << 3) + token_parity * 64; + int stsm_offset = (raw_row * 128 + raw_col) * 2; + const int pack_base = token_group_1 * 4; + uint32_t _stmatrix_addr_0 = static_cast((unsigned long long)(out_stage_addr + stsm_offset)); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_0), "r"(*reinterpret_cast(&out_packed[pack_base])), "r"(*reinterpret_cast(&out_packed[pack_base + 1])), "r"(*reinterpret_cast(&out_packed[pack_base + 2])), "r"(*reinterpret_cast(&out_packed[pack_base + 3])) + : "memory"); + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + tma_store_4d(out_tma, 0, (int)(bos_1 + (long long)(chunk_idx_1 * 32)), head_idx_1, value_split_idx_1, smem_out_addr + output_stage * 4096); + } + asm volatile("cp.async.bulk.commit_group;"); + } + output_stage = output_stage ^ 1; + } else { + float _tmem_load_5[16]; + asm volatile( + "tcgen05.ld.sync.aligned.16x256b.x4.b32" + " {%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15}, [%16];" + : "=r"(*reinterpret_cast(&_tmem_load_5[0])), "=r"(*reinterpret_cast(&_tmem_load_5[1])), "=r"(*reinterpret_cast(&_tmem_load_5[2])), "=r"(*reinterpret_cast(&_tmem_load_5[3])), "=r"(*reinterpret_cast(&_tmem_load_5[4])), "=r"(*reinterpret_cast(&_tmem_load_5[5])), "=r"(*reinterpret_cast(&_tmem_load_5[6])), "=r"(*reinterpret_cast(&_tmem_load_5[7])), "=r"(*reinterpret_cast(&_tmem_load_5[8])), "=r"(*reinterpret_cast(&_tmem_load_5[9])), "=r"(*reinterpret_cast(&_tmem_load_5[10])), "=r"(*reinterpret_cast(&_tmem_load_5[11])), "=r"(*reinterpret_cast(&_tmem_load_5[12])), "=r"(*reinterpret_cast(&_tmem_load_5[13])), "=r"(*reinterpret_cast(&_tmem_load_5[14])), "=r"(*reinterpret_cast(&_tmem_load_5[15])) + : "r"(taddr + 192 + (unsigned int)tmem_row_base_1) + : "memory"); + asm volatile("tcgen05.wait::ld.sync.aligned;" ::: "memory"); + if (elect_sync()) { + mbarrier_arrive(out_empty_addr); + } + #pragma unroll + for (int token_group_2 = 0; token_group_2 < 4; token_group_2++) { + int token_pair_2 = token_group_2 * 8 + lane_quad_1 * 2; + const int out_reg_base = token_group_2 * 4; + long long out_token_0 = bos_1 + (long long)(chunk_idx_1 * 32 + token_pair_2); + long long out_token_1 = out_token_0 + 1; + if (out_token_0 < eos_1) { + long long out_idx_top_0 = (out_token_0 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_top_1; + long long out_idx_bot_0 = (out_token_0 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_bot_1; + out[out_idx_top_0] = _tmem_load_5[out_reg_base]; + out[out_idx_bot_0] = _tmem_load_5[out_reg_base + 2]; + } + if (out_token_1 < eos_1) { + long long out_idx_top_1 = (out_token_1 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_top_1; + long long out_idx_bot_1 = (out_token_1 * (long long)num_heads + (long long)head_idx_1) * 128 + (long long)state_row_bot_1; + out[out_idx_top_1] = _tmem_load_5[out_reg_base + 1]; + out[out_idx_bot_1] = _tmem_load_5[out_reg_base + 3]; + } + } + } + epilogue_stage += 1; + if (epilogue_stage == 5) { epilogue_stage = 0; _phase_final_ready_1 ^= 1; } + } + if (epilogue_local_warp == 0) { + asm volatile("cp.async.bulk.wait_group 0;"); + } + asm volatile("barrier.sync 9, 128;" ::: "memory"); + if (epilogue_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(tmem_dealloc_ready_addr); + } + } + } + // ---- Role: mma ---- + } else if (warp == 9) { + { // mma_main + int split_task_idx_2 = blockIdx.x; + int task_idx_2 = split_task_idx_2 / 2; + int seq_idx_2 = seq_order[task_idx_2 / num_heads]; + long long bos_2 = cu_seqlens[seq_idx_2]; + long long eos_2 = cu_seqlens[seq_idx_2 + 1]; + int seq_len_2 = (int)(eos_2 - bos_2); + int num_chunks_2 = (seq_len_2 + 32 - 1) / 32; + unsigned int mma_stage = 0; + unsigned int _phase_qk_full_1 = 0; + unsigned int _phase_state_inp_ready = 0; + unsigned int _phase_out_empty_0 = 1; + unsigned int _phase_u_inp_ready = 0; + unsigned int _phase_u2_inp_ready = 0; + #pragma unroll 1 + for (int _chunk_idx = 0; _chunk_idx < num_chunks_2; _chunk_idx++) { + mbarrier_wait(qk_full_addr + (mma_stage) * 8, _phase_qk_full_1); + mbarrier_wait(state_inp_ready_addr + (mma_stage) * 8, _phase_state_inp_ready); + mbarrier_wait(out_empty_addr, _phase_out_empty_0); + _phase_out_empty_0 ^= 1; + int _mma_b_addr_0 = smem_qd_addr + mma_stage * 41984; + int _mma_b_lo_0 = make_warp_uniform((_mma_b_addr_0 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 67634320;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_out), "r"(_mma_b_lo_0), "r"(tmem_tmem_state_inp), "r"(0)); + int _mma_b_addr_1 = smem_kd_addr + mma_stage * 41984; + int _mma_b_lo_1 = make_warp_uniform((_mma_b_addr_1 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 67634320;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 16], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 24], db, id, p1;\n\t" + "add.u32 blo, blo, 250;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 32], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 40], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 48], db, id, p1;\n\t" + "add.u32 blo, blo, 2;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 56], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u_acc), "r"(_mma_b_lo_1), "r"(tmem_tmem_state_inp), "r"(0)); + elect_commit2(old_out_ready_addr + (mma_stage) * 8, raw_inputs_free_addr + (mma_stage) * 8); + mbarrier_wait(u_inp_ready_addr + (mma_stage) * 8, _phase_u_inp_ready); + int _mma_b_addr_2 = smem_inv_addr + mma_stage * 41984; + int _mma_b_lo_2 = make_warp_uniform((_mma_b_addr_2 >> 4) & 0x3FFF); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0xC0004010;\n\t" + "mov.b32 id, 67634320;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 64;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_u2_acc), "r"(_mma_b_lo_2), "r"(tmem_tmem_u2_inp), "r"(0)); + elect_commit(u2_acc_ready_addr + (mma_stage) * 8); + mbarrier_wait(u2_inp_ready_addr + (mma_stage) * 8, _phase_u2_inp_ready); + int _mma_b_addr_3 = smem_final_trans_addr + mma_stage * 41984; + int _mma_b_lo_3 = make_warp_uniform(((_mma_b_addr_3 >> 4) & 0x3FFF) | 0x1000000); + asm volatile( + "{\n\t" + ".reg .pred leader, p0, p1;\n\t" + ".reg .b32 dhi, blo, id;\n\t" + ".reg .b64 db;\n\t" + "elect.sync _|leader, 0xFFFFFFFF;\n\t" + "setp.ne.b32 p0, %3, 0;\n\t" + "setp.ne.b32 p1, 1, 0;\n\t" + "" + "mov.b32 dhi, 0x40004040;\n\t" + "mov.b32 id, 69797008;\n\t" + "mov.b32 blo, %1;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2], db, id, p0;\n\t" + "add.u32 blo, blo, 128;\n\t" + "mov.b64 db, {blo, dhi};\n\t" + "@leader tcgen05.mma.cta_group::1.kind::f16 [%0], [%2 + 8], db, id, p1;\n\t" + "}\n" + :: "r"(tmem_tmem_state_out), "r"(_mma_b_lo_3), "r"(tmem_tmem_u2_inp), "r"(1)); + elect_commit2(final_ready_addr + (mma_stage) * 8, smem_free_addr + (mma_stage) * 8); + mma_stage += 1; + if (mma_stage == 5) { mma_stage = 0; _phase_qk_full_1 ^= 1; _phase_state_inp_ready ^= 1; _phase_u_inp_ready ^= 1; _phase_u2_inp_ready ^= 1; } + } + unsigned int _phase_tmem_dealloc_ready_0 = 0; + mbarrier_wait(tmem_dealloc_ready_addr, _phase_tmem_dealloc_ready_0); + _phase_tmem_dealloc_ready_0 ^= 1; + int _tmem_dealloc_addr = *((volatile int*)tmem_addr_storage); + asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(_tmem_dealloc_addr), "r"(256)); + asm volatile("tcgen05.relinquish_alloc_permit.cta_group::1.sync.aligned;"); + } + // ---- Role: load ---- + } else if (warp == 10) { + { // load_main + int split_task_idx_3 = blockIdx.x; + int task_idx_3 = split_task_idx_3 / 2; + int value_split_idx_2 = split_task_idx_3 % 2; + int value_row_offset_2 = value_split_idx_2 * 64; + int seq_idx_3 = seq_order[task_idx_3 / num_heads]; + int head_idx_2 = task_idx_3 % num_heads; + long long bos_3 = cu_seqlens[seq_idx_3]; + long long eos_3 = cu_seqlens[seq_idx_3 + 1]; + int seq_len_3 = (int)(eos_3 - bos_3); + int num_chunks_3 = (seq_len_3 + 32 - 1) / 32; + unsigned int load_stage = 0; + unsigned int _phase_v_free = 1; + unsigned int _phase_qk_full_2 = 0; + #pragma unroll 1 + for (int chunk_idx_2 = 0; chunk_idx_2 < num_chunks_3; chunk_idx_2++) { + mbarrier_wait(v_free_addr + (load_stage) * 8, _phase_v_free); + mbarrier_wait(qk_full_addr + (load_stage) * 8, _phase_qk_full_2); + int chunk_is_full_1 = ((seq_len_3 >= (chunk_idx_2 + 1) * 32) ? 1 : 0); + if (elect_sync()) { + if (chunk_is_full_1 != 0) { + mbarrier_arrive_expect_tx(v_full_addr + (load_stage) * 8, 4096); + tma_3d_gmem2smem(smem_v_addr + load_stage * 41984, v_tma, value_row_offset_2, head_idx_2, (int)(bos_3 + (long long)(chunk_idx_2 * 32)), v_full_addr + (load_stage) * 8); + } + } + if (chunk_is_full_1 == 0) { + #pragma unroll + for (int v_load_iter = 0; v_load_iter < 8; v_load_iter++) { + int v_item = v_load_iter * 32 + lane; + int row = v_item / 8; + int segment = v_item % 8; + long long token = bos_3 + (long long)(chunk_idx_2 * 32 + row); + int token_valid = ((token < eos_3) ? 1 : 0); + long long v_src = (token * (long long)num_heads + (long long)head_idx_2) * 128 + (long long)value_row_offset_2 + (long long)(segment * 8); + int v_dst_row_addr = smem_v_addr + load_stage * 41984 + (unsigned int)(row * 64 * 2); + int v_dst_addr = (v_dst_row_addr + (segment * 8 * 2 ^ (v_dst_row_addr >> 7 & 7) << 4)); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(v_dst_addr), "l"(v + v_src), "r"((token_valid != 0) ? 16 : 0)); + } + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + } + asm volatile("barrier.sync 8, 32;" ::: "memory"); + if (elect_sync()) { + if (chunk_is_full_1 == 0) { + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + mbarrier_arrive(v_full_addr + (load_stage) * 8); + } + } + load_stage += 1; + if (load_stage == 5) { load_stage = 0; _phase_v_free ^= 1; _phase_qk_full_2 ^= 1; } + } + } + // ---- Role: prep ---- + } else if (warp >= 12 && warp <= 31) { + asm volatile("setmaxnreg.dec.sync.aligned.u32 48;"); + { // prep_main + int split_task_idx_4 = blockIdx.x; + int task_idx_4 = split_task_idx_4 / 2; + int seq_idx_4 = seq_order[task_idx_4 / num_heads]; + int head_idx_3 = task_idx_4 % num_heads; + long long bos_4 = cu_seqlens[seq_idx_4]; + long long eos_4 = cu_seqlens[seq_idx_4 + 1]; + int seq_len_4 = (int)(eos_4 - bos_4); + int num_chunks_4 = (seq_len_4 + 32 - 1) / 32; + int instance_id = (warp - 12) / 4; + int prep_instance = instance_id; + int warp_id_in_role_2 = (warp - 12); + int prep_local_warp = warp_id_in_role_2 - prep_instance * 4; + int prep_tid = prep_local_warp * 32 + lane; + int num_prep_iters = (num_chunks_4 + 4 - prep_instance) / 5; + unsigned int prep_stage = (unsigned int)prep_instance; + int gate_rate_stage_f32 = prep_instance * 10496; + if (prep_tid == 0) { + float _expf_0 = __expf(A_log[head_idx_3]); + smem_gate_rate_all[gate_rate_stage_f32] = _expf_0; + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + unsigned int _phase_raw_inputs_free = 1; + unsigned int _phase_gate_raw_full = 0; + unsigned int _phase_smem_free = 1; + unsigned int _phase_qk_raw_full = 0; + unsigned int _phase_prep_diag_ready = 0; + unsigned int _phase_prep_inv16_ready = 0; + #pragma unroll 1 + for (int prep_iter = 0; prep_iter < num_prep_iters; prep_iter++) { + int chunk_idx_3 = prep_iter * 5 + prep_instance; + int stage_f32 = prep_stage * 10496; + int stage_bf16 = prep_stage * 20992; + int chunk_is_full_2 = ((seq_len_4 >= (chunk_idx_3 + 1) * 32) ? 1 : 0); + float early_beta_value = 0.0f; + float early_gate0 = 0.0f; + if (chunk_is_full_2 != 0) { + mbarrier_wait(raw_inputs_free_addr + (prep_stage) * 8, _phase_raw_inputs_free); + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive_expect_tx(gate_raw_full_addr + (prep_stage) * 8, 8704); + tma_3d_gmem2smem(smem_g_raw_addr + prep_stage * 41984, g_tma, 0, head_idx_3, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + tma_2d_gmem2smem(smem_beta_raw_addr + prep_stage * 41984, beta_tma, head_idx_3 / 8 * 8, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), gate_raw_full_addr + (prep_stage) * 8); + mbarrier_arrive_expect_tx(qk_raw_full_addr + (prep_stage) * 8, 16384); + tma_4d_gmem2smem(smem_kd_addr + prep_stage * 41984, k_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + mbarrier_wait(gate_raw_full_addr + (prep_stage) * 8, _phase_gate_raw_full); + if (prep_local_warp == 2 && lane < 32) { + unsigned int beta_raw_pair[1]; + asm volatile("ld.shared.b32 %0, [%1];" : "=r"(*reinterpret_cast(&beta_raw_pair[0])) : "r"(smem_beta_raw_addr + prep_stage * 41984 + (unsigned int)(lane * 16) + (unsigned int)(head_idx_3 % 8 / 2 * 4))); + float beta_raw_pair_fp32[2]; + #pragma unroll + for (int _pair = 0; _pair < 1; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&beta_raw_pair_fp32[_pair * 2])[0]), "=f"((&beta_raw_pair_fp32[_pair * 2])[1]) + : "r"(beta_raw_pair[_pair + 0])); + } + float beta_logit = beta_raw_pair_fp32[0]; + if (head_idx_3 % 2 != 0) { + beta_logit = beta_raw_pair_fp32[1]; + } + float _tanh_approx_0; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_0) : "f"(beta_logit * 0.5f)); + early_beta_value = _tanh_approx_0 * 0.5f + 0.5f; + } + if (prep_tid < 128) { + float early_gate_rate = smem_gate_rate_all[stage_f32]; + float early_gate_bias = dt_bias[head_idx_3 * 128 + prep_tid]; + __nv_bfloat16 early_gate_raw = smem_g_raw_all[stage_bf16 + prep_tid]; + float _cvt_f32_0 = __bfloat162float(early_gate_raw); + float early_gate_arg = early_gate_rate * (_cvt_f32_0 + early_gate_bias); + float _tanh_approx_1; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_1) : "f"(early_gate_arg * 0.5f)); + float early_gate_sigmoid = _tanh_approx_1 * 0.5f + 0.5f; + early_gate0 = lower_bound * 1.4426950408889634f * early_gate_sigmoid; + } + } + mbarrier_wait(smem_free_addr + (prep_stage) * 8, _phase_smem_free); + if (chunk_is_full_2 != 0) { + if (prep_local_warp == 0) { + if (elect_sync()) { + tma_4d_gmem2smem(smem_q_raw_prefetch_addr + prep_stage * 41984, q_tma, 0, (int)(bos_4 + (long long)(chunk_idx_3 * 32)), head_idx_3, 0, qk_raw_full_addr + (prep_stage) * 8); + } + } + } + if (chunk_is_full_2 == 0) { + #pragma unroll + for (int gate_load_pass = 0; gate_load_pass < 4; gate_load_pass++) { + int gate_load_item = gate_load_pass * 128 + prep_tid; + int gate_load_row = gate_load_item / 16; + int gate_load_segment = gate_load_item % 16; + long long gate_load_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_load_row); + long long gate_load_base = (gate_load_token * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(gate_load_segment * 8); + asm volatile("cp.async.cg.shared::cta.global [%0], [%1], 16, %2;" + :: "r"(smem_g_raw_addr + prep_stage * 41984 + (unsigned int)(gate_load_item * 16)), "l"(g + gate_load_base), "r"((gate_load_token < eos_4) ? 16 : 0)); + } + } + if (chunk_is_full_2 == 0) { + asm volatile("cp.async.commit_group;"); + asm volatile("cp.async.wait_group 0;"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + } + if (prep_local_warp == 2 && lane < 32) { + float beta_value = early_beta_value; + if (chunk_is_full_2 == 0) { + long long beta_token = bos_4 + (long long)(chunk_idx_3 * 32 + lane); + if (beta_token < eos_4) { + float beta_logit_1 = (float)beta[beta_token * (long long)num_heads + (long long)head_idx_3]; + float _tanh_approx_2; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_2) : "f"(beta_logit_1 * 0.5f)); + beta_value = _tanh_approx_2 * 0.5f + 0.5f; + } + } + smem_prep_beta_all[stage_f32 + lane] = beta_value; + } + if (prep_tid < 128) { + int gate_col = prep_tid; + float gate_rate = smem_gate_rate_all[stage_f32]; + float gate_bias = dt_bias[head_idx_3 * 128 + gate_col]; + float prefix_log2 = 0.0f; + for (int gate_row = 0; gate_row < 32; gate_row++) { + long long gate_token = bos_4 + (long long)(chunk_idx_3 * 32 + gate_row); + float gate_log2 = 0.0f; + int gate_needs_compute = 1; + if (gate_row == 0) { + if (chunk_is_full_2 != 0) { + gate_log2 = early_gate0; + gate_needs_compute = 0; + } + } + if (gate_needs_compute != 0) { + if (gate_token < eos_4) { + __nv_bfloat16 gate_raw = smem_g_raw_all[stage_bf16 + gate_row * 128 + gate_col]; + float _cvt_f32_1 = __bfloat162float(gate_raw); + float gate_arg = gate_rate * (_cvt_f32_1 + gate_bias); + float _tanh_approx_3; + asm volatile("tanh.approx.f32 %0, %1;" : "=f"(_tanh_approx_3) : "f"(gate_arg * 0.5f)); + float gate_sigmoid = _tanh_approx_3 * 0.5f + 0.5f; + gate_log2 = lower_bound * 1.4426950408889634f * gate_sigmoid; + } + } + prefix_log2 += gate_log2; + smem_gate_all[stage_f32 + gate_row * 128 + gate_col] = prefix_log2; + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (chunk_is_full_2 != 0) { + mbarrier_wait(qk_raw_full_addr + (prep_stage) * 8, _phase_qk_raw_full); + } + if (prep_tid < 128) { + float total_log2 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_0 = approx_exp2(total_log2 - lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + prep_tid] = _exp2_0; + } + if (prep_tid == 0) { + float _exp2_1 = approx_exp2(lower_bound * 1.4426950408889634f * 16.0f); + smem_restore_factor_all[stage_f32 + 128] = _exp2_1; + } + #pragma unroll 1 + for (int work_pass = 0; work_pass < 4; work_pass++) { + int work_item = work_pass * 128 + prep_tid; + int row_1 = work_item / 16; + int segment_1 = work_item % 16; + long long token_1 = bos_4 + (long long)(chunk_idx_3 * 32 + row_1); + int token_valid_1 = ((token_1 < eos_4) ? 1 : 0); + long long gmem_base = (token_1 * (long long)num_heads + (long long)head_idx_3) * 128 + (long long)(segment_1 * 8); + float q_raw_vec[8]; + float k_raw_vec[8]; + q_raw_vec[0] = 0.0f; + q_raw_vec[1] = 0.0f; + q_raw_vec[2] = 0.0f; + q_raw_vec[3] = 0.0f; + q_raw_vec[4] = 0.0f; + q_raw_vec[5] = 0.0f; + q_raw_vec[6] = 0.0f; + q_raw_vec[7] = 0.0f; + k_raw_vec[0] = 0.0f; + k_raw_vec[1] = 0.0f; + k_raw_vec[2] = 0.0f; + k_raw_vec[3] = 0.0f; + k_raw_vec[4] = 0.0f; + k_raw_vec[5] = 0.0f; + k_raw_vec[6] = 0.0f; + k_raw_vec[7] = 0.0f; + if (chunk_is_full_2 != 0) { + unsigned int packed[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed[0])), "=r"(*reinterpret_cast(&packed[(0) + 1])), "=r"(*reinterpret_cast(&packed[(0) + 2])), "=r"(*reinterpret_cast(&packed[(0) + 3])) + : "r"((smem_q_raw_prefetch_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32[_pair * 2])[0]), "=f"((&packed_fp32[_pair * 2])[1]) + : "r"(packed[_pair + 0])); + } + #pragma unroll + for (int value_idx = 0; value_idx < 8; value_idx++) { + q_raw_vec[value_idx] = packed_fp32[value_idx]; + } + unsigned int packed_0[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0[0])), "=r"(*reinterpret_cast(&packed_0[(0) + 1])), "=r"(*reinterpret_cast(&packed_0[(0) + 2])), "=r"(*reinterpret_cast(&packed_0[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32[_pair * 2])[0]), "=f"((&packed_0_fp32[_pair * 2])[1]) + : "r"(packed_0[_pair + 0])); + } + #pragma unroll + for (int value_idx_1 = 0; value_idx_1 < 8; value_idx_1++) { + k_raw_vec[value_idx_1] = packed_0_fp32[value_idx_1]; + } + } else if (token_valid_1 != 0) { + { + const uint4* _vptr_0 = reinterpret_cast(q + gmem_base); + uint4 _vld_0[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_0[_blk] = _vptr_0[_blk]; + uint32_t* _vpairs_0 = reinterpret_cast(&_vld_0[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&q_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_0[_pair])); + } + } + } + { + const uint4* _vptr_1 = reinterpret_cast(k + gmem_base); + uint4 _vld_1[1]; + #pragma unroll + for (int _blk = 0; _blk < 1; _blk++) { + _vld_1[_blk] = _vptr_1[_blk]; + uint32_t* _vpairs_1 = reinterpret_cast(&_vld_1[_blk]); + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[0]), "=f"((&k_raw_vec[0 + _blk * 8 + _pair * 2])[1]) + : "r"(_vpairs_1[_pair])); + } + } + } + } + float q_sum = 0.0f; + float k_sum = 0.0f; + for (int elem_in_segment = 0; elem_in_segment < 8; elem_in_segment++) { + float q_raw = q_raw_vec[elem_in_segment]; + float k_raw = k_raw_vec[elem_in_segment]; + float _fma_0 = __fmaf_rn(q_raw, q_raw, q_sum); + q_sum = _fma_0; + float _fma_1 = __fmaf_rn(k_raw, k_raw, k_sum); + k_sum = _fma_1; + } + float _shfl_xor_0 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 8); + q_sum += _shfl_xor_0; + float _shfl_xor_1 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 8); + k_sum += _shfl_xor_1; + float _shfl_xor_2 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 4); + q_sum += _shfl_xor_2; + float _shfl_xor_3 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 4); + k_sum += _shfl_xor_3; + float _shfl_xor_4 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 2); + q_sum += _shfl_xor_4; + float _shfl_xor_5 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 2); + k_sum += _shfl_xor_5; + float _shfl_xor_6 = __shfl_xor_sync(0xFFFFFFFF, q_sum, 1); + q_sum += _shfl_xor_6; + float _shfl_xor_7 = __shfl_xor_sync(0xFFFFFFFF, k_sum, 1); + k_sum += _shfl_xor_7; + float _rsqrt_0 = rsqrtf(q_sum + 1e-06f); + float q_inv = _rsqrt_0; + float _rsqrt_1 = rsqrtf(k_sum + 1e-06f); + float k_inv = _rsqrt_1; + const float2 _scale2_2 = {q_inv, q_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(q_raw_vec)[_ls], _scale2_2); + const float2 _scale2_3 = {k_inv, k_inv}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(k_raw_vec)[_ls], _scale2_3); + float qd_vec[8]; + float kd_vec[8]; + float ki_vec[8]; + for (int elem_in_segment_1 = 0; elem_in_segment_1 < 8; elem_in_segment_1++) { + int col = segment_1 * 8 + elem_in_segment_1; + float prefix = smem_gate_all[stage_f32 + row_1 * 128 + col]; + float common_log2 = lower_bound * 1.4426950408889634f * 16.0f; + float _exp2_2 = approx_exp2(prefix - common_log2); + float decay = _exp2_2; + qd_vec[elem_in_segment_1] = decay; + kd_vec[elem_in_segment_1] = decay; + ki_vec[elem_in_segment_1] = k_raw_vec[elem_in_segment_1] / decay; + } + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], reinterpret_cast(q_raw_vec)[_ls]); + const float2 _scale2_4 = {scale, scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(qd_vec)[_ls], _scale2_4); + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(kd_vec)[_ls], reinterpret_cast(k_raw_vec)[_ls]); + unsigned int packed_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(qd_vec[_lp*2 + 0], qd_vec[_lp*2+1 + 0])); + packed_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word = 0; word < 4; word++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word * 4)), "r"(packed_1[word])); + } + unsigned int packed_0_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(kd_vec[_lp*2 + 0], kd_vec[_lp*2+1 + 0])); + packed_0_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_1 = 0; word_1 < 4; word_1++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_1 * 4)), "r"(packed_0_1[word_1])); + } + unsigned int packed_1_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(ki_vec[_lp*2 + 0], ki_vec[_lp*2+1 + 0])); + packed_1_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_2 = 0; word_2 < 4; word_2++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 ^ (segment_1 * 8 / 64 * 4096 + row_1 * 128 + segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_2 * 4)), "r"(packed_1_1[word_2])); + } + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + int pair_row_base = prep_local_warp / 2 * 16; + int pair_col_base = prep_local_warp % 2 * 16; + unsigned int a_frag[4]; + unsigned int b_frag[4]; + float acc[8]; + if (pair_row_base >= pair_col_base) { + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_kd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + int row0 = pair_row_base + lane / 4; + int row1 = row0 + 8; + int col0 = pair_col_base + lane % 4 * 2; + float beta0 = smem_prep_beta_all[stage_f32 + row0]; + float beta1 = smem_prep_beta_all[stage_f32 + row1]; + float seed[8]; + seed[0] = 0.0f; + seed[1] = 0.0f; + seed[2] = 0.0f; + seed[3] = 0.0f; + seed[4] = 0.0f; + seed[5] = 0.0f; + seed[6] = 0.0f; + seed[7] = 0.0f; + if (row0 > col0) { + seed[0] = acc[0] * beta0; + } + if (row0 > col0 + 1) { + seed[1] = acc[1] * beta0; + } + if (row1 > col0) { + seed[2] = acc[2] * beta1; + } + if (row1 > col0 + 1) { + seed[3] = acc[3] * beta1; + } + if (row0 > col0 + 8) { + seed[4] = acc[4] * beta0; + } + if (row0 > col0 + 9) { + seed[5] = acc[5] * beta0; + } + if (row1 > col0 + 8) { + seed[6] = acc[6] * beta1; + } + if (row1 > col0 + 9) { + seed[7] = acc[7] * beta1; + } + unsigned int seed_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(seed[_lp*2 + 0], seed[_lp*2+1 + 0])); + seed_packed[_lp] = *(uint32_t*)&_bf2; + } + int seed_lane_row = lane % 16; + int seed_lane_col = lane / 16 * 8; + int byte_off = (pair_row_base + seed_lane_row) * 128 + (pair_col_base + seed_lane_col) * 2; + int swizzled_off = byte_off ^ (byte_off >> 7 & 7) << 4; + int seed_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off; + uint32_t _stmatrix_addr_5 = static_cast((unsigned long long)seed_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_5), "r"(*reinterpret_cast(&seed_packed[0])), "r"(*reinterpret_cast(&seed_packed[1])), "r"(*reinterpret_cast(&seed_packed[2])), "r"(*reinterpret_cast(&seed_packed[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[0]), "=f"(acc[1]), "=f"(acc[2]), "=f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(acc[4]), "=f"(acc[(4) + 1]), "=f"(acc[(4) + 2]), "=f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a_frag[0]), "=r"(a_frag[1]), "=r"(a_frag[2]), "=r"(a_frag[3]) + : "r"(smem_qd_addr + prep_stage * 41984 + (unsigned int)(((lane / 16 / 8 * 256 + (pair_row_base + lane % 16) * 8 + (lane / 16 % 8 * 16 ^ (pair_row_base + lane % 16 & 7) << 4) / 16 ^ 2 ^ 6 ^ 2 ^ 6) + 256 ^ 2 ^ 6 ^ 2) * 16)) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(b_frag[0]), "=r"(b_frag[1]), "=r"(b_frag[2]), "=r"(b_frag[3]) + : "r"(smem_ki_addr + prep_stage * 41984 + (unsigned int)(((((((((lane % 16 / 8 / 8 * 256 + (pair_col_base + 8 * (lane / 16) + lane % 8) * 8 + (lane % 16 / 8 % 8 * 16 ^ (pair_col_base + 8 * (lane / 16) + lane % 8 & 7) << 4) / 16 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256 + 256 ^ 6) + 256 - 256 + 256 ^ 2) - 256 + 256 ^ 6) - 256 + 256 ^ 2) - 256) * 16)) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[0]), "+f"(acc[1]), "+f"(acc[2]), "+f"(acc[3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[0]), "r"(b_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};\n" + : "+f"(acc[4]), "+f"(acc[(4) + 1]), "+f"(acc[(4) + 2]), "+f"(acc[(4) + 3]) + : "r"(a_frag[0]), "r"(a_frag[1]), "r"(a_frag[2]), "r"(a_frag[3]), "r"(b_frag[2]), "r"(b_frag[(2) + 1])); + } else { + acc[0] = 0.0f; + acc[1] = 0.0f; + acc[2] = 0.0f; + acc[3] = 0.0f; + acc[4] = 0.0f; + acc[5] = 0.0f; + acc[6] = 0.0f; + acc[7] = 0.0f; + } + int row0_1 = pair_row_base + lane / 4; + int row1_1 = row0_1 + 8; + int col0_1 = pair_col_base + lane % 4 * 2; + float mqk[8]; + mqk[0] = 0.0f; + mqk[1] = 0.0f; + mqk[2] = 0.0f; + mqk[3] = 0.0f; + mqk[4] = 0.0f; + mqk[5] = 0.0f; + mqk[6] = 0.0f; + mqk[7] = 0.0f; + if (row0_1 >= col0_1) { + mqk[0] = acc[0]; + } + if (row0_1 >= col0_1 + 1) { + mqk[1] = acc[1]; + } + if (row1_1 >= col0_1) { + mqk[2] = acc[2]; + } + if (row1_1 >= col0_1 + 1) { + mqk[3] = acc[3]; + } + if (row0_1 >= col0_1 + 8) { + mqk[4] = acc[4]; + } + if (row0_1 >= col0_1 + 9) { + mqk[5] = acc[5]; + } + if (row1_1 >= col0_1 + 8) { + mqk[6] = acc[6]; + } + if (row1_1 >= col0_1 + 9) { + mqk[7] = acc[7]; + } + unsigned int mqk_packed[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(mqk[_lp*2 + 0], mqk[_lp*2+1 + 0])); + mqk_packed[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int publish_pair = 0; publish_pair < 2; publish_pair++) { + int publish_row = pair_col_base + publish_pair * 8 + (lane & 7); + int publish_col = 128 + pair_row_base + lane / 8 * 8; + uint32_t _stmatrix_addr_6 = static_cast((unsigned long long)(smem_final_trans_addr + prep_stage * 41984 + (unsigned int)(publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 ^ (publish_col / 64 * 4096 + publish_row * 128 + publish_col % 64 * 2 >> 7 & 7) << 4))); + asm volatile("stmatrix.sync.aligned.m8n8.x2.trans.shared.b16 [%0], {%1, %2};\n" + :: "r"(_stmatrix_addr_6), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2])), "r"(*reinterpret_cast(&mqk_packed[publish_pair * 2 + 1])) + : "memory"); + } + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_tid < 128) { + float total_log2_1 = smem_gt_prefix_all[stage_f32 + prep_tid]; + float _exp2_3 = approx_exp2(total_log2_1); + smem_gt_all[stage_f32 + prep_tid] = _exp2_3; + } + if (prep_local_warp >= 2) { + int stage_f32_0 = prep_stage * 10496; + float restore_scale = smem_restore_factor_all[stage_f32_0 + 128]; + float restore_factor[8]; + int restore_segment = lane & 15; + #pragma unroll + for (int restore_elem = 0; restore_elem < 8; restore_elem++) { + int restore_col = restore_segment * 8 + restore_elem; + restore_factor[restore_elem] = smem_restore_factor_all[stage_f32_0 + restore_col]; + } + #pragma unroll 1 + for (int restore_pass = 0; restore_pass < 6; restore_pass++) { + int restore_row = 8 + (prep_local_warp - 2) * 12 + restore_pass * 2 + (lane >> 4); + float restore_qd_values[8]; + float restore_kd_values[8]; + float restore_ki_values[8]; + unsigned int packed_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_2[0])), "=r"(*reinterpret_cast(&packed_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_2[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_1[_pair * 2])[0]), "=f"((&packed_fp32_1[_pair * 2])[1]) + : "r"(packed_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_2 = 0; value_idx_2 < 8; value_idx_2++) { + restore_qd_values[value_idx_2] = packed_fp32_1[value_idx_2]; + } + unsigned int packed_0_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_2[0])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_2[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_1[_pair * 2])[0]), "=f"((&packed_0_fp32_1[_pair * 2])[1]) + : "r"(packed_0_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_3 = 0; value_idx_3 < 8; value_idx_3++) { + restore_kd_values[value_idx_3] = packed_0_fp32_1[value_idx_3]; + } + unsigned int packed_1_2[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_2[0])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_2[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32[_pair * 2])[0]), "=f"((&packed_1_fp32[_pair * 2])[1]) + : "r"(packed_1_2[_pair + 0])); + } + #pragma unroll + for (int value_idx_4 = 0; value_idx_4 < 8; value_idx_4++) { + restore_ki_values[value_idx_4] = packed_1_fp32[value_idx_4]; + } + float restore_kr_values[8]; + #pragma unroll + for (int restore_elem_1 = 0; restore_elem_1 < 8; restore_elem_1++) { + restore_kr_values[restore_elem_1] = restore_ki_values[restore_elem_1] * restore_factor[restore_elem_1]; + } + const float2 _scale2_7 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values)[_ls], _scale2_7); + const float2 _scale2_8 = {restore_scale, restore_scale}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values)[_ls], _scale2_8); + unsigned int packed_2_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values[_lp*2 + 0], restore_qd_values[_lp*2+1 + 0])); + packed_2_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_3 = 0; word_3 < 4; word_3++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_3 * 4)), "r"(packed_2_1[word_3])); + } + unsigned int packed_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values[_lp*2 + 0], restore_kd_values[_lp*2+1 + 0])); + packed_3[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_4 = 0; word_4 < 4; word_4++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_4 * 4)), "r"(packed_3[word_4])); + } + unsigned int packed_4[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values[_lp*2 + 0], restore_kr_values[_lp*2+1 + 0])); + packed_4[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_5 = 0; word_5 < 4; word_5++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 ^ (restore_segment * 8 / 64 * 4096 + restore_row * 128 + restore_segment * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_5 * 4)), "r"(packed_4[word_5])); + } + } + } + if (prep_local_warp == 0) { + int inverse_row = lane; + int diag_block = inverse_row / 8; + int lane_in_diag = lane & 7; + float inv_row[8]; + unsigned int packed_5[4]; + int byte_off_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_1 = byte_off_1 ^ (byte_off_1 >> 7 & 7) << 4; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_5[0])), "=r"(*reinterpret_cast(&packed_5[(0) + 1])), "=r"(*reinterpret_cast(&packed_5[(0) + 2])), "=r"(*reinterpret_cast(&packed_5[(0) + 3])) + : "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1)); + float packed_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_2[_pair * 2])[0]), "=f"((&packed_fp32_2[_pair * 2])[1]) + : "r"(packed_5[_pair + 0])); + } + #pragma unroll + for (int value_idx_5 = 0; value_idx_5 < 8; value_idx_5++) { + inv_row[value_idx_5] = packed_fp32_2[value_idx_5]; + } + #pragma unroll + for (int diag_elem = 0; diag_elem < 8; diag_elem++) { + if (lane_in_diag == diag_elem) { + inv_row[diag_elem] = 1.0f; + } + } + int diag_group_base = lane - lane_in_diag; + #pragma unroll + for (int src_row = 0; src_row < 7; src_row++) { + float row_scale = -inv_row[src_row]; + #pragma unroll + for (int prev_col = 0; prev_col < src_row; prev_col++) { + int pivot_lane = diag_group_base + src_row; + float _shfl_0 = __shfl_sync(0xFFFFFFFF, inv_row[prev_col], pivot_lane); + float pivot = _shfl_0; + if (lane_in_diag > src_row) { + float _fma_2 = __fmaf_rn(row_scale, pivot, inv_row[prev_col]); + inv_row[prev_col] = _fma_2; + } + } + if (lane_in_diag > src_row) { + inv_row[src_row] = row_scale; + } + } + unsigned int packed_0_3[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(inv_row[_lp*2 + 0], inv_row[_lp*2+1 + 0])); + packed_0_3[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_1_1 = inverse_row * 128 + diag_block * 8 * 2; + int swizzled_off_2 = byte_off_1_1 ^ (byte_off_1_1 >> 7 & 7) << 4; + #pragma unroll + for (int word_6 = 0; word_6 < 4; word_6++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"(smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_2 + (unsigned int)(word_6 * 4)), "r"(packed_0_3[word_6])); + } + } + if (prep_local_warp < 2) { + if (elect_sync()) { + mbarrier_arrive(prep_diag_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_diag_ready_addr + (prep_stage) * 8, _phase_prep_diag_ready); + } + if (prep_local_warp < 2) { + int lane_row = lane & 7; + int byte_off_2 = (prep_local_warp * 16 + 8 + lane_row) * 128 + (prep_local_warp * 16 + 8) * 2; + int swizzled_off_3 = byte_off_2 ^ (byte_off_2 >> 7 & 7) << 4; + int d_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3; + int byte_off_0 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_1_1 = byte_off_0 ^ (byte_off_0 >> 7 & 7) << 4; + int c_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_1; + int byte_off_2_1 = (prep_local_warp * 16 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_3_1 = byte_off_2_1 ^ (byte_off_2_1 >> 7 & 7) << 4; + int a_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_1; + unsigned int d_frag[2]; + unsigned int c_frag[1]; + float dc_acc[4]; + unsigned int dc_bf16[2]; + unsigned int inv_a_frag[1]; + float o_acc[4]; + unsigned int o_bf16[2]; + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[0]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" + : "=r"(d_frag[1]) + : "r"(d_addr) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(c_frag[0]) + : "r"(c_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc_acc[0]), "=f"(dc_acc[1]), "=f"(dc_acc[2]), "=f"(dc_acc[3]) + : "r"(d_frag[0]), "r"(d_frag[1]), "r"(c_frag[0])); + const float2 _scale2_9 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 2; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc_acc)[_ls], _scale2_9); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc_acc[_lp*2 + 0], dc_acc[_lp*2+1 + 0])); + dc_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n" + : "=r"(inv_a_frag[0]) + : "r"(a_addr) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k8.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o_acc[0]), "=f"(o_acc[1]), "=f"(o_acc[2]), "=f"(o_acc[3]) + : "r"(dc_bf16[0]), "r"(dc_bf16[1]), "r"(inv_a_frag[0])); + #pragma unroll + for (int _lp = 0; _lp < 2; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o_acc[_lp*2 + 0], o_acc[_lp*2+1 + 0])); + o_bf16[_lp] = *(uint32_t*)&_bf2; + } + int byte_off_4 = (prep_local_warp * 16 + 8 + lane_row) * 128 + prep_local_warp * 16 * 2; + int swizzled_off_5 = byte_off_4 ^ (byte_off_4 >> 7 & 7) << 4; + int o_addr = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_5; + uint32_t _stmatrix_addr_10 = static_cast((unsigned long long)o_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x1.shared.b16 [%0], {%1};\n" + :: "r"(_stmatrix_addr_10), "r"(*reinterpret_cast(&o_bf16[0])) + : "memory"); + if (elect_sync()) { + mbarrier_arrive(prep_inv16_ready_addr + (prep_stage) * 8); + } + mbarrier_wait(prep_inv16_ready_addr + (prep_stage) * 8, _phase_prep_inv16_ready); + } + if (prep_local_warp == 0) { + int lane_row_1 = lane % 16; + int lane_col = lane / 16 * 8; + int byte_off_3 = (16 + lane_row_1) * 128 + (16 + lane_col) * 2; + int swizzled_off_4 = byte_off_3 ^ (byte_off_3 >> 7 & 7) << 4; + int d_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_4; + int byte_off_0_1 = (16 + lane_row_1) * 128 + lane_col * 2; + int swizzled_off_1_2 = byte_off_0_1 ^ (byte_off_0_1 >> 7 & 7) << 4; + int c_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_1_2; + int byte_off_2_2 = lane_row_1 * 128 + lane_col * 2; + int swizzled_off_3_2 = byte_off_2_2 ^ (byte_off_2_2 >> 7 & 7) << 4; + int a_addr_1 = smem_inv_work_addr + prep_stage * 41984 + (unsigned int)swizzled_off_3_2; + unsigned int d32_frag[4]; + unsigned int c32_frag[4]; + float dc32_acc[8]; + unsigned int dc32_bf16[4]; + unsigned int a32_frag[4]; + float o32_acc[8]; + unsigned int o32_bf16[4]; + unsigned int zero32_bf16[4]; + asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(d32_frag[0]), "=r"(d32_frag[1]), "=r"(d32_frag[2]), "=r"(d32_frag[3]) + : "r"(d_addr_1) + : "memory"); + int d_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + (16 + lane_row_1) * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_11 = static_cast((unsigned long long)d_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_11), "r"(*reinterpret_cast(&d32_frag[0])), "r"(*reinterpret_cast(&d32_frag[1])), "r"(*reinterpret_cast(&d32_frag[2])), "r"(*reinterpret_cast(&d32_frag[3])) + : "memory"); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(c32_frag[0]), "=r"(c32_frag[1]), "=r"(c32_frag[2]), "=r"(c32_frag[3]) + : "r"(c_addr_1) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[0]), "=f"(dc32_acc[1]), "=f"(dc32_acc[2]), "=f"(dc32_acc[3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[0]), "r"(c32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(dc32_acc[4]), "=f"(dc32_acc[(4) + 1]), "=f"(dc32_acc[(4) + 2]), "=f"(dc32_acc[(4) + 3]) + : "r"(d32_frag[0]), "r"(d32_frag[1]), "r"(d32_frag[2]), "r"(d32_frag[3]), "r"(c32_frag[2]), "r"(c32_frag[(2) + 1])); + const float2 _scale2_12 = {-1.0f, -1.0f}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(dc32_acc)[_ls], _scale2_12); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(dc32_acc[_lp*2 + 0], dc32_acc[_lp*2+1 + 0])); + dc32_bf16[_lp] = *(uint32_t*)&_bf2; + } + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n" + : "=r"(a32_frag[0]), "=r"(a32_frag[1]), "=r"(a32_frag[2]), "=r"(a32_frag[3]) + : "r"(a_addr_1) + : "memory"); + int a_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + lane_row_1 * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_13 = static_cast((unsigned long long)a_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.trans.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_13), "r"(*reinterpret_cast(&a32_frag[0])), "r"(*reinterpret_cast(&a32_frag[1])), "r"(*reinterpret_cast(&a32_frag[2])), "r"(*reinterpret_cast(&a32_frag[3])) + : "memory"); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[0]), "=f"(o32_acc[1]), "=f"(o32_acc[2]), "=f"(o32_acc[3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[0]), "r"(a32_frag[1])); + asm volatile("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {0f00000000, 0f00000000, 0f00000000, 0f00000000};\n" + : "=f"(o32_acc[4]), "=f"(o32_acc[(4) + 1]), "=f"(o32_acc[(4) + 2]), "=f"(o32_acc[(4) + 3]) + : "r"(dc32_bf16[0]), "r"(dc32_bf16[1]), "r"(dc32_bf16[2]), "r"(dc32_bf16[3]), "r"(a32_frag[2]), "r"(a32_frag[(2) + 1])); + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(o32_acc[_lp*2 + 0], o32_acc[_lp*2+1 + 0])); + o32_bf16[_lp] = *(uint32_t*)&_bf2; + } + int o_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)(lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 ^ (lane_col / 16 * 1024 + (16 + lane_row_1) * 32 + lane_col % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_14 = static_cast((unsigned long long)o_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_14), "r"(*reinterpret_cast(&o32_bf16[0])), "r"(*reinterpret_cast(&o32_bf16[1])), "r"(*reinterpret_cast(&o32_bf16[2])), "r"(*reinterpret_cast(&o32_bf16[3])) + : "memory"); + #pragma unroll + for (int zero_word = 0; zero_word < 4; zero_word++) { + zero32_bf16[zero_word] = 0; + } + int zero_publish_addr = (smem_inv_addr + prep_stage * 41984 + (unsigned int)((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 ^ ((16 + lane_col) / 16 * 1024 + lane_row_1 * 32 + (16 + lane_col) % 16 * 2 >> 7 & 1) << 4)); + uint32_t _stmatrix_addr_15 = static_cast((unsigned long long)zero_publish_addr); + asm volatile("stmatrix.sync.aligned.m8n8.x4.shared.b16 [%0], {%1, %2, %3, %4};\n" + :: "r"(_stmatrix_addr_15), "r"(*reinterpret_cast(&zero32_bf16[0])), "r"(*reinterpret_cast(&zero32_bf16[1])), "r"(*reinterpret_cast(&zero32_bf16[2])), "r"(*reinterpret_cast(&zero32_bf16[3])) + : "memory"); + } else if (prep_local_warp == 1) { + int stage_f32_0_1 = prep_stage * 10496; + float restore_scale_1 = smem_restore_factor_all[stage_f32_0_1 + 128]; + float restore_factor_1[8]; + int restore_segment_1 = lane & 15; + #pragma unroll + for (int restore_elem_2 = 0; restore_elem_2 < 8; restore_elem_2++) { + int restore_col_1 = restore_segment_1 * 8 + restore_elem_2; + restore_factor_1[restore_elem_2] = smem_restore_factor_all[stage_f32_0_1 + restore_col_1]; + } + #pragma unroll 1 + for (int restore_pass_1 = 0; restore_pass_1 < 4; restore_pass_1++) { + int restore_row_1 = restore_pass_1 * 2 + (lane >> 4); + float restore_qd_values_1[8]; + float restore_kd_values_1[8]; + float restore_ki_values_1[8]; + unsigned int packed_6[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_6[0])), "=r"(*reinterpret_cast(&packed_6[(0) + 1])), "=r"(*reinterpret_cast(&packed_6[(0) + 2])), "=r"(*reinterpret_cast(&packed_6[(0) + 3])) + : "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_fp32_3[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_fp32_3[_pair * 2])[0]), "=f"((&packed_fp32_3[_pair * 2])[1]) + : "r"(packed_6[_pair + 0])); + } + #pragma unroll + for (int value_idx_6 = 0; value_idx_6 < 8; value_idx_6++) { + restore_qd_values_1[value_idx_6] = packed_fp32_3[value_idx_6]; + } + unsigned int packed_0_4[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_0_4[0])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 1])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 2])), "=r"(*reinterpret_cast(&packed_0_4[(0) + 3])) + : "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_0_fp32_2[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_0_fp32_2[_pair * 2])[0]), "=f"((&packed_0_fp32_2[_pair * 2])[1]) + : "r"(packed_0_4[_pair + 0])); + } + #pragma unroll + for (int value_idx_7 = 0; value_idx_7 < 8; value_idx_7++) { + restore_kd_values_1[value_idx_7] = packed_0_fp32_2[value_idx_7]; + } + unsigned int packed_1_3[4]; + asm volatile("ld.shared.v4.b32 {%0,%1,%2,%3}, [%4];" + : "=r"(*reinterpret_cast(&packed_1_3[0])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 1])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 2])), "=r"(*reinterpret_cast(&packed_1_3[(0) + 3])) + : "r"((smem_ki_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)))); + float packed_1_fp32_1[8]; + #pragma unroll + for (int _pair = 0; _pair < 4; _pair++) { + asm volatile( + "{\n\t" + "shl.b32 %0, %2, 16;\n\t" + "and.b32 %1, %2, 0xffff0000;\n\t" + "}\n" + : "=f"((&packed_1_fp32_1[_pair * 2])[0]), "=f"((&packed_1_fp32_1[_pair * 2])[1]) + : "r"(packed_1_3[_pair + 0])); + } + #pragma unroll + for (int value_idx_8 = 0; value_idx_8 < 8; value_idx_8++) { + restore_ki_values_1[value_idx_8] = packed_1_fp32_1[value_idx_8]; + } + float restore_kr_values_1[8]; + #pragma unroll + for (int restore_elem_3 = 0; restore_elem_3 < 8; restore_elem_3++) { + restore_kr_values_1[restore_elem_3] = restore_ki_values_1[restore_elem_3] * restore_factor_1[restore_elem_3]; + } + const float2 _scale2_16 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_qd_values_1)[_ls], _scale2_16); + const float2 _scale2_17 = {restore_scale_1, restore_scale_1}; + #pragma unroll + for (int _ls = 0; _ls < 4; _ls++) + mul_f32x2_inplace(&reinterpret_cast(restore_kd_values_1)[_ls], _scale2_17); + unsigned int packed_2_2[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_qd_values_1[_lp*2 + 0], restore_qd_values_1[_lp*2+1 + 0])); + packed_2_2[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_7 = 0; word_7 < 4; word_7++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_qd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_7 * 4)), "r"(packed_2_2[word_7])); + } + unsigned int packed_3_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kd_values_1[_lp*2 + 0], restore_kd_values_1[_lp*2+1 + 0])); + packed_3_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_8 = 0; word_8 < 4; word_8++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kd_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_8 * 4)), "r"(packed_3_1[word_8])); + } + unsigned int packed_4_1[4]; + #pragma unroll + for (int _lp = 0; _lp < 4; _lp++) { + __nv_bfloat162 _bf2 = __float22bfloat162_rn(make_float2(restore_kr_values_1[_lp*2 + 0], restore_kr_values_1[_lp*2+1 + 0])); + packed_4_1[_lp] = *(uint32_t*)&_bf2; + } + #pragma unroll + for (int word_9 = 0; word_9 < 4; word_9++) { + asm volatile("st.shared.b32 [%0], %1;" :: "r"((smem_kr_trans_addr + prep_stage * 41984 + (unsigned int)(restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 ^ (restore_segment_1 * 8 / 64 * 4096 + restore_row_1 * 128 + restore_segment_1 * 8 % 64 * 2 >> 7 & 7) << 4)) + (unsigned int)(word_9 * 4)), "r"(packed_4_1[word_9])); + } + } + } + asm volatile("fence.proxy.async.shared::cta;" ::: "memory"); + if (prep_instance == 0) { + asm volatile("barrier.sync 11, 128;" ::: "memory"); + } else if (prep_instance == 1) { + asm volatile("barrier.sync 12, 128;" ::: "memory"); + } else { + if (prep_instance == 2) { + asm volatile("barrier.sync 13, 128;" ::: "memory"); + } else if (prep_instance == 3) { + asm volatile("barrier.sync 14, 128;" ::: "memory"); + } else { + asm volatile("barrier.sync 15, 128;" ::: "memory"); + } + } + if (prep_local_warp == 0) { + if (elect_sync()) { + mbarrier_arrive(qk_full_addr + (prep_stage) * 8); + } + } + for (int _advance = 0; _advance < 5; _advance++) { + prep_stage += 1; + if (prep_stage == 5) { prep_stage = 0; _phase_raw_inputs_free ^= 1; _phase_gate_raw_full ^= 1; _phase_smem_free ^= 1; _phase_qk_raw_full ^= 1; _phase_prep_diag_ready ^= 1; _phase_prep_inv16_ready ^= 1; } + } + } + } + } + + // Cleanup +} + +} // extern "C" + +// clang-format on diff --git a/csrc/kda/flashkda_bf16_fused_m64_binding.cu b/csrc/kda/flashkda_bf16_fused_m64_binding.cu new file mode 100644 index 00000000000..e3ef52f3445 --- /dev/null +++ b/csrc/kda/flashkda_bf16_fused_m64_binding.cu @@ -0,0 +1,100 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#include "flashkda_binding_common.cuh" + +// The standalone exporter deliberately emits its own fixed-width typedefs so +// the frozen device TU has no host-library dependency. Isolate those names +// when the source is included into FlashInfer's C++ binding, whose CUDA/TVM +// headers already define the standard typedefs. +#define uint8_t flashkda_generated_uint8_t +#define uint16_t flashkda_generated_uint16_t +#define uint32_t flashkda_generated_uint32_t +#define uint64_t flashkda_generated_uint64_t +#define int32_t flashkda_generated_int32_t +#define int16_t flashkda_generated_int16_t +#include "flashkda_bf16_fused_m64.cu" +#undef uint8_t +#undef uint16_t +#undef uint32_t +#undef uint64_t +#undef int32_t +#undef int16_t + +namespace flashinfer { +namespace flash_kda { + +static_assert(THREADS == 1024); +static_assert(SMEM_TOTAL == 219136); + +void RunM64(TensorView q, TensorView k, TensorView v, TensorView g, TensorView beta, + TensorView beta_tma, TensorView A_log, TensorView dt_bias, TensorView cu_seqlens, + TensorView seq_order, TensorView initial_state, TensorView out, TensorView final_state, + TensorView descriptor_storage, int64_t prepare_descriptors, int64_t num_heads, + int64_t use_initial_state, int64_t store_final_state, double scale, double lower_bound, + int64_t cuda_stream) { + TVM_FFI_ICHECK(cuda_stream >= 0) << "cuda_stream must be a non-negative stream handle"; + TVM_FFI_ICHECK(q.device().device_type == kDLCUDA) << "q must be a CUDA tensor"; + const int32_t device_id = q.device().device_id; + ffi::CUDADeviceGuard device_guard(device_id); + CheckExactSm100a(device_id); + + const int64_t num_seqs = + CheckCommonInputs(q, k, v, g, beta, beta_tma, A_log, dt_bias, cu_seqlens, seq_order, + initial_state, out, final_state, descriptor_storage, prepare_descriptors, + num_heads, use_initial_state, store_final_state, scale, lower_bound); + TVM_FFI_ICHECK(num_seqs == 1 && num_heads == 64) + << "the M64 FlashKDA variant is specialized for fixed N=1, H=64; got " + "N=" + << num_seqs << ", H=" << num_heads; + + constexpr int32_t kSmemBytes = SMEM_TOTAL; + CheckDynamicSmemCapacity(device_id, kSmemBytes); + + CheckCuda(cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m64, + cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemBytes), + "cudaFuncSetAttribute(kernel_flashkda_bf16_fused_m64)"); + + const int64_t grid_x_i64 = 2 * num_seqs * num_heads; + TVM_FFI_ICHECK(grid_x_i64 > 0 && grid_x_i64 <= std::numeric_limits::max()) + << "M64 FlashKDA grid.x is out of range: " << grid_x_i64; + const dim3 grid(static_cast(grid_x_i64), 1, 1); + const dim3 block(THREADS, 1, 1); + const cudaStream_t stream = reinterpret_cast(static_cast(cuda_stream)); + const TmaPointers tma = EncodeTmaPointers<64>(q, k, v, g, beta_tma, out, descriptor_storage, + prepare_descriptors, stream); + + kernel_flashkda_bf16_fused_m64<<>>( + reinterpret_cast<__nv_bfloat16*>(q.data_ptr()), tma.q, + reinterpret_cast<__nv_bfloat16*>(k.data_ptr()), tma.k, + reinterpret_cast<__nv_bfloat16*>(v.data_ptr()), tma.v, + reinterpret_cast<__nv_bfloat16*>(g.data_ptr()), tma.g, + reinterpret_cast<__nv_bfloat16*>(beta.data_ptr()), tma.beta, + reinterpret_cast(A_log.data_ptr()), reinterpret_cast(dt_bias.data_ptr()), + reinterpret_cast(cu_seqlens.data_ptr()), + reinterpret_cast(seq_order.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(initial_state.data_ptr()), + reinterpret_cast<__nv_bfloat16*>(out.data_ptr()), tma.out, + reinterpret_cast<__nv_bfloat16*>(final_state.data_ptr()), static_cast(num_heads), + static_cast(use_initial_state), static_cast(store_final_state), + static_cast(scale), static_cast(lower_bound)); + CheckCuda(cudaGetLastError(), "kernel_flashkda_bf16_fused_m64 launch"); +} + +} // namespace flash_kda +} // namespace flashinfer + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(run, flashinfer::flash_kda::RunM64); diff --git a/csrc/kda/flashkda_binding_common.cuh b/csrc/kda/flashkda_binding_common.cuh new file mode 100644 index 00000000000..0e4bc0c8152 --- /dev/null +++ b/csrc/kda/flashkda_binding_common.cuh @@ -0,0 +1,453 @@ +/* + * Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + * + * 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. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "tvm_ffi_utils.h" + +namespace flashinfer { +namespace flash_kda { + +constexpr int64_t kHeadDim = 128; +constexpr size_t kTensorMapCount = 6; +constexpr size_t kTensorMapAlignment = 64; +static_assert(sizeof(CUtensorMap) == 128); +constexpr size_t kDescriptorStorageBytes = kTensorMapCount * sizeof(CUtensorMap); + +inline void CheckCuda(cudaError_t status, const char* operation) { + TVM_FFI_ICHECK(status == cudaSuccess) << operation << " failed: " << cudaGetErrorString(status); +} + +inline void CheckCudaTensor(const TensorView& tensor, const char* name, int32_t device_id) { + TVM_FFI_ICHECK(tensor.device().device_type == kDLCUDA) << name << " must be a CUDA tensor"; + TVM_FFI_ICHECK(tensor.device().device_id == device_id) + << name << " must be on CUDA device " << device_id << ", got " << tensor.device().device_id; + TVM_FFI_ICHECK(tensor.IsContiguous()) << name << " must be contiguous"; +} + +inline void CheckDtype(const TensorView& tensor, const char* name, DLDataType expected) { + const DLDataType actual = tensor.dtype(); + TVM_FFI_ICHECK(actual.code == expected.code && actual.bits == expected.bits && + actual.lanes == expected.lanes) + << name << " has wrong dtype: expected (code=" << int(expected.code) + << ", bits=" << int(expected.bits) << ", lanes=" << int(expected.lanes) + << "), got (code=" << int(actual.code) << ", bits=" << int(actual.bits) + << ", lanes=" << int(actual.lanes) << ")"; +} + +struct TensorByteRange { + uintptr_t begin; + uintptr_t end; +}; + +inline TensorByteRange GetTensorByteRange(const TensorView& tensor, const char* name) { + const DLDataType dtype = tensor.dtype(); + const uint64_t bits = static_cast(dtype.bits) * static_cast(dtype.lanes); + TVM_FFI_ICHECK(bits > 0 && bits % 8 == 0) << name << " has a non-byte-addressable dtype"; + const uint64_t bytes = static_cast(tensor.numel()) * (bits / 8); + const uintptr_t begin = reinterpret_cast(tensor.data_ptr()); + TVM_FFI_ICHECK(bytes <= std::numeric_limits::max() - begin) + << name << " byte range overflows uintptr_t"; + return {begin, begin + static_cast(bytes)}; +} + +inline void CheckNoOverlap(const TensorView& lhs, const char* lhs_name, const TensorView& rhs, + const char* rhs_name) { + const TensorByteRange lhs_range = GetTensorByteRange(lhs, lhs_name); + const TensorByteRange rhs_range = GetTensorByteRange(rhs, rhs_name); + const bool overlaps = lhs_range.begin < rhs_range.end && rhs_range.begin < lhs_range.end; + TVM_FFI_ICHECK(!overlaps) << lhs_name << " must not overlap " << rhs_name + << ": the frozen kernel uses __restrict__ pointers"; +} + +inline void CheckNoPartialOverlapOrExactAlias(const TensorView& lhs, const char* lhs_name, + const TensorView& rhs, const char* rhs_name) { + const TensorByteRange lhs_range = GetTensorByteRange(lhs, lhs_name); + const TensorByteRange rhs_range = GetTensorByteRange(rhs, rhs_name); + const bool overlaps = lhs_range.begin < rhs_range.end && rhs_range.begin < lhs_range.end; + const bool exact_alias = lhs_range.begin == rhs_range.begin && lhs_range.end == rhs_range.end; + TVM_FFI_ICHECK(!overlaps || exact_alias) + << lhs_name << " and " << rhs_name + << " must either be disjoint or exactly alias the same state storage"; +} + +inline void CheckExactSm100a(int32_t device_id) { + int major = 0; + int minor = 0; + CheckCuda(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device_id), + "cudaDeviceGetAttribute(major)"); + CheckCuda(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, device_id), + "cudaDeviceGetAttribute(minor)"); + TVM_FFI_ICHECK(major == 10 && minor == 0) + << "FlashKDA frozen kernels require exact compute capability 10.0 " + "(sm_100a), got " + << major << "." << minor; +} + +inline int64_t CheckCommonInputs(const TensorView& q, const TensorView& k, const TensorView& v, + const TensorView& g, const TensorView& beta, + const TensorView& beta_tma, const TensorView& A_log, + const TensorView& dt_bias, const TensorView& cu_seqlens, + const TensorView& seq_order, const TensorView& initial_state, + const TensorView& out, const TensorView& final_state, + const TensorView& descriptor_storage, int64_t prepare_descriptors, + int64_t num_heads, int64_t use_initial_state, + int64_t store_final_state, double scale, double lower_bound) { + TVM_FFI_ICHECK(prepare_descriptors == 0 || prepare_descriptors == 1) + << "prepare_descriptors must be 0 or 1, got " << prepare_descriptors; + TVM_FFI_ICHECK(num_heads > 0 && num_heads <= std::numeric_limits::max()) + << "num_heads must be in the positive int32 range, got " << num_heads; + TVM_FFI_ICHECK(use_initial_state == 0 || use_initial_state == 1) + << "use_initial_state must be 0 or 1, got " << use_initial_state; + TVM_FFI_ICHECK(store_final_state == 0 || store_final_state == 1) + << "store_final_state must be 0 or 1, got " << store_final_state; + TVM_FFI_ICHECK(std::isfinite(scale) && std::isfinite(static_cast(scale))) + << "scale must be finite and representable as float32, got " << scale; + TVM_FFI_ICHECK(std::isfinite(lower_bound) && lower_bound < 0.0 && + std::isfinite(static_cast(lower_bound))) + << "lower_bound must be finite, negative, and representable as " + "float32, got " + << lower_bound; + + const int32_t device_id = q.device().device_id; + CheckCudaTensor(q, "q", device_id); + CheckCudaTensor(k, "k", device_id); + CheckCudaTensor(v, "v", device_id); + CheckCudaTensor(g, "g", device_id); + CheckCudaTensor(beta, "beta", device_id); + CheckCudaTensor(beta_tma, "beta_tma", device_id); + CheckCudaTensor(A_log, "A_log", device_id); + CheckCudaTensor(dt_bias, "dt_bias", device_id); + CheckCudaTensor(cu_seqlens, "cu_seqlens", device_id); + CheckCudaTensor(seq_order, "seq_order", device_id); + CheckCudaTensor(initial_state, "initial_state", device_id); + CheckCudaTensor(out, "out", device_id); + CheckCudaTensor(final_state, "final_state", device_id); + CheckCudaTensor(descriptor_storage, "descriptor_storage", device_id); + + CheckDtype(q, "q", dl_bfloat16); + CheckDtype(k, "k", dl_bfloat16); + CheckDtype(v, "v", dl_bfloat16); + CheckDtype(g, "g", dl_bfloat16); + CheckDtype(beta, "beta", dl_bfloat16); + CheckDtype(beta_tma, "beta_tma", dl_bfloat16); + CheckDtype(A_log, "A_log", dl_float32); + CheckDtype(dt_bias, "dt_bias", dl_float32); + CheckDtype(cu_seqlens, "cu_seqlens", dl_int64); + CheckDtype(seq_order, "seq_order", dl_int32); + CheckDtype(initial_state, "initial_state", dl_bfloat16); + CheckDtype(out, "out", dl_bfloat16); + CheckDtype(final_state, "final_state", dl_bfloat16); + CheckDtype(descriptor_storage, "descriptor_storage", dl_uint8); + + TVM_FFI_ICHECK(descriptor_storage.numel() >= static_cast(kDescriptorStorageBytes)) + << "descriptor_storage must contain at least " << kDescriptorStorageBytes << " bytes"; + TVM_FFI_ICHECK(reinterpret_cast(descriptor_storage.data_ptr()) % kTensorMapAlignment == + 0) + << "descriptor_storage must be aligned to " << kTensorMapAlignment << " bytes"; + + TVM_FFI_ICHECK(q.ndim() >= 3) << "q must have trailing [H, 128] dimensions"; + TVM_FFI_ICHECK(q.size(q.ndim() - 1) == kHeadDim && q.size(q.ndim() - 2) == num_heads) + << "q must have trailing shape [" << num_heads << ", 128]"; + const int64_t token_count = q.numel() / (num_heads * kHeadDim); + TVM_FFI_ICHECK(token_count > 0) << "q must contain at least one token"; + + for (const auto& named : {std::pair(&k, "k"), + std::pair(&v, "v"), + std::pair(&g, "g"), + std::pair(&out, "out")}) { + const TensorView& tensor = *named.first; + TVM_FFI_ICHECK(tensor.ndim() >= 3 && tensor.size(tensor.ndim() - 1) == kHeadDim && + tensor.size(tensor.ndim() - 2) == num_heads && tensor.numel() == q.numel()) + << named.second << " must match q's flattened [tokens, H, 128] shape"; + } + + TVM_FFI_ICHECK(beta.ndim() >= 2 && beta.size(beta.ndim() - 1) == num_heads && + beta.numel() == token_count * num_heads) + << "beta must match flattened [tokens, H]"; + const int64_t beta_tma_heads = std::max(num_heads, 8); + TVM_FFI_ICHECK(beta_tma.ndim() >= 2 && beta_tma.size(beta_tma.ndim() - 1) == beta_tma_heads && + beta_tma.numel() % beta_tma_heads == 0 && + beta_tma.numel() / beta_tma_heads >= std::max(token_count, 32)) + << "beta_tma must have at least [max(tokens, 32), max(H, 8)] " + "storage"; + TVM_FFI_ICHECK(A_log.numel() == num_heads) << "A_log must contain H elements"; + TVM_FFI_ICHECK(dt_bias.numel() == num_heads * kHeadDim) + << "dt_bias must contain H * 128 elements"; + TVM_FFI_ICHECK(cu_seqlens.ndim() == 1 && cu_seqlens.numel() >= 2) + << "cu_seqlens must be a one-dimensional tensor with at least two " + "elements"; + + const int64_t num_seqs = cu_seqlens.numel() - 1; + TVM_FFI_ICHECK(seq_order.ndim() == 1 && seq_order.numel() == num_seqs) + << "seq_order must contain one int32 entry per sequence"; + const int64_t state_numel = num_seqs * num_heads * kHeadDim * kHeadDim; + if (use_initial_state != 0) { + TVM_FFI_ICHECK(initial_state.numel() == state_numel) + << "initial_state must have flattened [N, H, 128, 128] size"; + } + if (store_final_state != 0) { + TVM_FFI_ICHECK(final_state.numel() == state_numel) + << "final_state must have flattened [N, H, 128, 128] size"; + } + CheckNoOverlap(out, "out", q, "q"); + CheckNoOverlap(out, "out", k, "k"); + CheckNoOverlap(out, "out", v, "v"); + CheckNoOverlap(out, "out", g, "g"); + CheckNoOverlap(out, "out", beta, "beta"); + CheckNoOverlap(out, "out", beta_tma, "beta_tma"); + CheckNoOverlap(out, "out", A_log, "A_log"); + CheckNoOverlap(out, "out", dt_bias, "dt_bias"); + CheckNoOverlap(out, "out", cu_seqlens, "cu_seqlens"); + CheckNoOverlap(out, "out", seq_order, "seq_order"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", q, "q"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", k, "k"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", v, "v"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", g, "g"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", beta, "beta"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", beta_tma, "beta_tma"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", A_log, "A_log"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", dt_bias, "dt_bias"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", cu_seqlens, "cu_seqlens"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", seq_order, "seq_order"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", out, "out"); + if (use_initial_state != 0) { + CheckNoOverlap(out, "out", initial_state, "initial_state"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", initial_state, "initial_state"); + } + if (store_final_state != 0) { + CheckNoOverlap(out, "out", final_state, "final_state"); + CheckNoOverlap(descriptor_storage, "descriptor_storage", final_state, "final_state"); + CheckNoOverlap(final_state, "final_state", q, "q"); + CheckNoOverlap(final_state, "final_state", k, "k"); + CheckNoOverlap(final_state, "final_state", v, "v"); + CheckNoOverlap(final_state, "final_state", g, "g"); + CheckNoOverlap(final_state, "final_state", beta, "beta"); + CheckNoOverlap(final_state, "final_state", beta_tma, "beta_tma"); + CheckNoOverlap(final_state, "final_state", A_log, "A_log"); + CheckNoOverlap(final_state, "final_state", dt_bias, "dt_bias"); + CheckNoOverlap(final_state, "final_state", cu_seqlens, "cu_seqlens"); + CheckNoOverlap(final_state, "final_state", seq_order, "seq_order"); + if (use_initial_state != 0) { + CheckNoPartialOverlapOrExactAlias(initial_state, "initial_state", final_state, "final_state"); + } + } + return num_seqs; +} + +inline CUtensorMap EncodeQkTma(const TensorView& tensor, const char* name) { + TVM_FFI_ICHECK(tensor.ndim() >= 2) << name << " must have at least two dimensions"; + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + TVM_FFI_ICHECK(d1 > 0 && d2 > 0 && d1 % 64 == 0) << name << " has invalid trailing dimensions"; + const int64_t outer2 = tensor.numel() / (d1 * d2); + uint64_t global_dim[4] = {64, static_cast(outer2), static_cast(d2), + static_cast(d1 / 64)}; + TVM_FFI_ICHECK(global_dim[0] > 0 && global_dim[1] > 0 && global_dim[2] >= 1 && global_dim[3] >= 2) + << name << " cannot encode the (64, 32, 1, 2) TMA box"; + uint64_t global_strides[3] = {static_cast(d2 * d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(64 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[4] = {64, 32, 1, 2}; + uint32_t elem_strides[4] = {1, 1, 1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, tensor.data_ptr(), + 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); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for " << name << " with CUresult=" << int(result); + return tensor_map; +} + +template +inline CUtensorMap EncodeValueTma(const TensorView& tensor) { + static_assert(ValueRows == 64 || ValueRows == 128); + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + const int64_t outer2 = tensor.numel() / (d1 * d2); + uint64_t global_dim[3] = {static_cast(d1), static_cast(d2), + static_cast(outer2)}; + TVM_FFI_ICHECK(global_dim[0] >= ValueRows && global_dim[1] >= 1 && global_dim[2] > 0) + << "v cannot encode the (" << ValueRows << ", 1, 32) TMA box"; + uint64_t global_strides[2] = {static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * d2 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[3] = {ValueRows, 1, 32}; + uint32_t elem_strides[3] = {1, 1, 1}; + CUtensorMap tensor_map{}; + constexpr CUtensorMapSwizzle swizzle = + ValueRows == 64 ? CU_TENSOR_MAP_SWIZZLE_128B : CU_TENSOR_MAP_SWIZZLE_NONE; + const CUresult result = cuTensorMapEncodeTiled( + &tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, tensor.data_ptr(), global_dim, + global_strides, box_dim, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, swizzle, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for v with CUresult=" << int(result); + return tensor_map; +} + +inline CUtensorMap EncodeGateTma(const TensorView& tensor) { + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + const int64_t outer2 = tensor.numel() / (d1 * d2); + uint64_t global_dim[3] = {static_cast(d1), static_cast(d2), + static_cast(outer2)}; + TVM_FFI_ICHECK(global_dim[0] >= 128 && global_dim[1] >= 1 && global_dim[2] > 0) + << "g cannot encode the (128, 1, 32) TMA box"; + uint64_t global_strides[2] = {static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * d2 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[3] = {128, 1, 32}; + uint32_t elem_strides[3] = {1, 1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 3, tensor.data_ptr(), + global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for g with CUresult=" << int(result); + return tensor_map; +} + +inline CUtensorMap EncodeBetaTma(const TensorView& tensor) { + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t outer1 = tensor.numel() / d1; + uint64_t global_dim[2] = {static_cast(d1), static_cast(outer1)}; + TVM_FFI_ICHECK(global_dim[0] >= 8 && global_dim[1] >= 32) + << "beta_tma cannot encode the (8, 32) TMA box"; + uint64_t global_strides[1] = {static_cast(d1 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[2] = {8, 32}; + uint32_t elem_strides[2] = {1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 2, tensor.data_ptr(), + global_dim, global_strides, box_dim, elem_strides, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, + CU_TENSOR_MAP_L2_PROMOTION_NONE, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for beta_tma with CUresult=" << int(result); + return tensor_map; +} + +template +inline CUtensorMap EncodeOutputTma(const TensorView& tensor) { + static_assert(ValueRows == 64 || ValueRows == 128); + const int64_t d1 = tensor.size(tensor.ndim() - 1); + const int64_t d2 = tensor.size(tensor.ndim() - 2); + const int64_t outer2 = tensor.numel() / (d1 * d2); + TVM_FFI_ICHECK(d1 > 0 && d2 > 0 && d1 % 64 == 0) << "out has invalid trailing dimensions"; + uint64_t global_dim[4] = {64, static_cast(outer2), static_cast(d2), + static_cast(d1 / 64)}; + constexpr uint32_t value_splits = ValueRows / 64; + TVM_FFI_ICHECK(global_dim[0] >= 64 && global_dim[1] > 0 && global_dim[2] >= 1 && + global_dim[3] >= value_splits) + << "out cannot encode the (64, 32, 1, " << value_splits << ") TMA box"; + uint64_t global_strides[3] = {static_cast(d2 * d1 * sizeof(__nv_bfloat16)), + static_cast(d1 * sizeof(__nv_bfloat16)), + static_cast(64 * sizeof(__nv_bfloat16))}; + uint32_t box_dim[4] = {64, 32, 1, value_splits}; + uint32_t elem_strides[4] = {1, 1, 1, 1}; + CUtensorMap tensor_map{}; + const CUresult result = + cuTensorMapEncodeTiled(&tensor_map, CU_TENSOR_MAP_DATA_TYPE_BFLOAT16, 4, tensor.data_ptr(), + 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); + TVM_FFI_ICHECK(result == CUDA_SUCCESS) + << "cuTensorMapEncodeTiled failed for out with CUresult=" << int(result); + return tensor_map; +} + +struct TmaPointers { + void* q; + void* k; + void* v; + void* g; + void* beta; + void* out; +}; + +struct TensorMapWords { + static constexpr size_t kWordCount = kDescriptorStorageBytes / sizeof(uint64_t); + uint64_t words[kWordCount]; +}; + +static __global__ void PublishTensorMaps(uint64_t* destination, TensorMapWords source) { + const uint32_t index = threadIdx.x; + if (index < TensorMapWords::kWordCount) { + destination[index] = source.words[index]; + } +} + +template +inline TmaPointers EncodeTmaPointers(const TensorView& q, const TensorView& k, const TensorView& v, + const TensorView& g, const TensorView& beta_tma, + const TensorView& out, const TensorView& descriptor_storage, + int64_t prepare_descriptors, cudaStream_t stream) { + if (prepare_descriptors != 0) { + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone; + CheckCuda(cudaStreamIsCapturing(stream, &capture_status), + "cudaStreamIsCapturing(TMA descriptor preparation)"); + TVM_FFI_ICHECK(capture_status == cudaStreamCaptureStatusNone) + << "prepare_descriptors must be 0 during CUDA graph capture; warm " + "this exact workspace and tensor signature before capture"; + + const std::array host_maps = { + EncodeQkTma(q, "q"), EncodeQkTma(k, "k"), EncodeValueTma(v), + EncodeGateTma(g), EncodeBetaTma(beta_tma), EncodeOutputTma(out), + }; + static_assert(sizeof(host_maps) == kDescriptorStorageBytes); + TensorMapWords words{}; + std::memcpy(words.words, host_maps.data(), sizeof(host_maps)); + PublishTensorMaps<<<1, 128, 0, stream>>>( + reinterpret_cast(descriptor_storage.data_ptr()), words); + CheckCuda(cudaGetLastError(), "PublishTensorMaps launch"); + } + + auto* bytes = static_cast(descriptor_storage.data_ptr()); + constexpr size_t stride = sizeof(CUtensorMap); + return { + bytes + 0 * stride, bytes + 1 * stride, bytes + 2 * stride, + bytes + 3 * stride, bytes + 4 * stride, bytes + 5 * stride, + }; +} + +inline void CheckDynamicSmemCapacity(int32_t device_id, int32_t smem_bytes) { + int max_optin = 0; + CheckCuda(cudaDeviceGetAttribute(&max_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id), + "cudaDeviceGetAttribute(max dynamic shared memory)"); + TVM_FFI_ICHECK(max_optin >= smem_bytes) + << "device exposes only " << max_optin << " bytes of opt-in shared memory; FlashKDA requires " + << smem_bytes; +} + +} // namespace flash_kda +} // namespace flashinfer diff --git a/docs/api/kda.rst b/docs/api/kda.rst new file mode 100644 index 00000000000..43f173f10cb --- /dev/null +++ b/docs/api/kda.rst @@ -0,0 +1,16 @@ +.. _apikda: + +flashinfer.kda +============== + +Phase-neutral recurrent Kimi Delta Attention (KDA) facade. The public +``recurrent_kda`` entry point keeps decode and speculative decode on +``flashinfer.kda_decode`` while dispatching eligible ordinary multi-token +prefill to the optimized backend described in :ref:`apikda_prefill`. + +.. currentmodule:: flashinfer.kda + +.. autosummary:: + :toctree: ../generated + + recurrent_kda diff --git a/docs/api/kda_prefill.rst b/docs/api/kda_prefill.rst new file mode 100644 index 00000000000..30992059bf0 --- /dev/null +++ b/docs/api/kda_prefill.rst @@ -0,0 +1,103 @@ +.. _apikda_prefill: + +flashinfer.kda_prefill +====================== + +Optimized recurrent Kimi Delta Attention (KDA) prefill support. The +:func:`flashinfer.kda.recurrent_kda` facade dispatches a strict ordinary +multi-token prefill subset to frozen FlashKDA-compatible SM100a kernels. + +.. currentmodule:: flashinfer.kda_prefill + +.. autosummary:: + :toctree: ../generated + + RecurrentKDAPrefillWorkspace + +Optimized B200 prefill subset +----------------------------- + +``flashinfer.kda.recurrent_kda`` uses the frozen prefill backend only when +every condition below holds: + +* the device has compute capability 10.0; +* input is ordinary multi-token prefill: fixed ``T > 1``, or packed input + whose total token count is greater than its number of sequences; +* Q, K, V, and G are contiguous BF16 ``[B,T,H,128]`` tensors with one shared + head count, and beta is contiguous BF16 ``[B,T,H]``; +* ``A_log`` is contiguous FP32 ``[H]`` and ``dt_bias`` is contiguous FP32 + ``[H,128]`` or flattened ``[H*128]``; +* ``use_qk_l2norm_in_kernel=True``, ``use_gate_in_kernel=True``, + ``beta_is_logit=True``, and ``lower_bound`` is a finite negative value; +* speculative decode, GQA, state indices, committed-state sources, and + accepted-token/checkpoint features are not enabled. + +Calls outside that subset retain the existing CuTe-DSL path. In particular, +T=1 decode and speculative decode are not rerouted. + +Fixed input omits ``cu_seqlens``. Packed input has ``B=1`` and accepts a +contiguous CUDA int32 or int64 ``cu_seqlens``. The frozen binding consumes +int64 offsets; pass int64 directly for CUDA graph capture to avoid an +in-capture conversion allocation. Offset values are a caller contract: +``cu_seqlens[0] == 0``, entries are strictly increasing (every sequence is +non-empty), and ``cu_seqlens[-1] == total_tokens``. FlashInfer does not +synchronize the device to inspect these values; invalid offsets may cause +out-of-bounds device access. + +Packed scheduling +----------------- + +Packed prefill optionally accepts ``seq_order``, a contiguous CUDA int32 +tensor with one entry per sequence. It is a caller contract that this tensor +is a permutation of ``[0, N)``. Ordering sequences by decreasing length +reduces the final partial wave. FlashInfer validates dtype, device, rank, and +size without synchronizing the device to inspect permutation values. + +When ``seq_order=None``, a cached identity order is used. Fixed ``B=1,H=64`` +selects the two-CTA M64 value-split kernel; every packed input and every other +head count selects M128. + +State and graph semantics +------------------------- + +The BF16 state layout remains ``[N,H,V,K]`` and an explicitly supplied +``initial_state`` is still updated in place, even when +``output_final_state=False``. The frozen kernels load each CTA's disjoint +state rows before writing the final rows back to the same storage, so no +separate state scratch or copy-back is required. If no initial state is +supplied, a final state is allocated only when ``output_final_state=True``. + +The frozen kernel uses restricted output storage. A preallocated ``output`` +must not overlap Q, K, V, G, beta, or ``initial_state``. + +Eager calls without ``prefill_workspace`` use an internal serialized workspace +for the current CUDA stream. This default workspace is eager-only and cannot +be used during CUDA graph capture. + +CUDA graph capture requires a caller-owned +``RecurrentKDAPrefillWorkspace(device)`` and a preallocated ``output``. The +workspace owns optional final-state scratch for calls without an initial +state, beta padding, and separate 768-byte M64 and M128 TMA descriptor blocks. +It binds to the device and CUDA stream of its first ``recurrent_kda`` call. +Warm it eagerly on the intended capture stream with the exact Q, K, V, G, +beta, and output tensors, then synchronize that stream before capture. Packed +graphs must also pass preallocated int64 ``cu_seqlens`` and int32 +``seq_order``. The warm call prepares descriptors; capture accepts only the +exact warmed pointer, shape, stride, and dtype signature and performs no +descriptor preparation. + +The workspace must outlive its graph and every replay. Use one distinct +workspace for each captured ``recurrent_kda`` invocation, including two KDA +invocations in the same graph. Once a workspace participates in capture, any +later Python use through ``recurrent_kda``—eager or another capture—is +rejected. ``graph.replay()`` does not re-enter Python and remains valid. +Sequential replay launches may be issued while a different PyTorch stream is +current, but the caller must provide normal stream ordering. The Python stream +binding applies to eager warmup and capture calls, which must use the same +stream. + +When an explicit workspace is used with ``initial_state=None`` and +``output_final_state=True``, the returned final state is workspace-owned +stable scratch. Otherwise an explicitly supplied ``initial_state`` is updated +directly in place by the frozen kernel. The small-head ``H < 8`` path captures +the beta copy into workspace-owned padded storage before the frozen launch. diff --git a/docs/index.rst b/docs/index.rst index 3e2f7d25709..5789e238619 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -53,7 +53,9 @@ FlashInfer is a library and kernel generator for Large Language Models that prov api/activation api/gdn_decode api/gdn_prefill + api/kda api/kda_decode + api/kda_prefill api/mamba api/mhc api/quantization diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 41cbed50b4b..754a2405154 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -124,7 +124,10 @@ from .grouped_mm import grouped_mm_fp8 as grouped_mm_fp8 from .grouped_mm import grouped_mm_mxfp8 as grouped_mm_mxfp8 from .grouped_mm import grouped_mm_fp4 as grouped_mm_fp4 -from .kda_decode import recurrent_kda as recurrent_kda +from .kda_prefill import ( + RecurrentKDAPrefillWorkspace as RecurrentKDAPrefillWorkspace, +) +from .kda import recurrent_kda as recurrent_kda from .mla import BatchMLAPagedAttentionWrapper as BatchMLAPagedAttentionWrapper from . import mhc as mhc from . import msa_ops as msa_ops diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 04e71dd11af..0349646288a 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -60,6 +60,10 @@ ) from .jit.fp4_kv_dequantization import gen_fp4_kv_dequantization_module from .jit.fp4_kv_quantization import gen_fp4_kv_quantization_module +from .jit.flash_kda import ( + gen_flash_kda_m64_module, + gen_flash_kda_m128_module, +) from .jit.nvfp4_attention_sm120 import gen_nvfp4_attention_sm120_module from .jit.fp8_quantization import gen_mxfp8_quantization_sm100_module from .jit.fused_moe import ( @@ -498,6 +502,7 @@ def gen_all_modules( has_sm80 = sm_capabilities.get("sm80", False) has_sm90 = sm_capabilities.get("sm90", False) has_sm100 = sm_capabilities.get("sm100", False) + has_sm100a_exact = sm_capabilities.get("sm100a_exact", False) has_sm100f = sm_capabilities.get("sm100f", False) has_sm103 = sm_capabilities.get("sm103", False) has_sm110 = sm_capabilities.get("sm110", False) @@ -521,6 +526,16 @@ def gen_all_modules( ) if has_sm120 or has_sm121: jit_specs.append(gen_nvfp4_attention_sm120_module()) + if has_sm100a_exact: + # Frozen FlashKDA sources use the SM100a-only tcgen05/TMEM surface. + # Do not package them for SM100f, SM103, or later architectures until + # those exact cubins have been independently validated. + jit_specs.extend( + [ + gen_flash_kda_m64_module(), + gen_flash_kda_m128_module(), + ] + ) if add_act: for act_name in act_func_def_str: @@ -962,6 +977,8 @@ def has_sm(compute: str, version: str) -> bool: "sm80": has_any_sm8x and get_cuda_version() >= Version("11.0"), "sm90": has_sm("compute_90", "12.3"), "sm100": has_sm("compute_100", "12.8"), + "sm100a_exact": (10, "0a") in compilation_context.TARGET_CUDA_ARCHS + and get_cuda_version() >= Version("12.8"), "sm100f": has_sm("compute_100", "12.9"), "sm103": has_sm("compute_103", "12.9"), "sm110": has_sm("compute_110", "13.0"), diff --git a/flashinfer/jit/__init__.py b/flashinfer/jit/__init__.py index 75192911c78..cbfb10f3945 100644 --- a/flashinfer/jit/__init__.py +++ b/flashinfer/jit/__init__.py @@ -100,6 +100,21 @@ from .fp4_kv_quantization import ( gen_fp4_kv_quantization_module as gen_fp4_kv_quantization_module, ) +from .flash_kda import ( + gen_flash_kda_m64_module as gen_flash_kda_m64_module, +) +from .flash_kda import ( + gen_flash_kda_m128_module as gen_flash_kda_m128_module, +) +from .flash_kda import ( + get_flash_kda_prefill_module as get_flash_kda_prefill_module, +) +from .flash_kda import ( + load_flash_kda_m64_module as load_flash_kda_m64_module, +) +from .flash_kda import ( + load_flash_kda_m128_module as load_flash_kda_m128_module, +) from .nvfp4_attention_sm120 import ( gen_nvfp4_attention_sm120_module as gen_nvfp4_attention_sm120_module, ) diff --git a/flashinfer/jit/flash_kda.py b/flashinfer/jit/flash_kda.py new file mode 100644 index 00000000000..13d9d53a21b --- /dev/null +++ b/flashinfer/jit/flash_kda.py @@ -0,0 +1,148 @@ +""" +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. +""" + +import functools +from pathlib import Path +from typing import Literal + +from . import env as jit_env +from .core import JitSpec, gen_jit_spec, logger, sm100a_nvcc_flags + +FlashKDAVariant = Literal["m64", "m128"] + + +def _get_flash_kda_csrc_dir() -> Path: + """Locate frozen FlashKDA sources in installed and source checkouts.""" + + installed = jit_env.FLASHINFER_CSRC_DIR / "kda" + if installed.exists(): + return installed + + checkout = Path(__file__).resolve().parents[2] / "csrc" / "kda" + if checkout.exists(): + return checkout + + raise FileNotFoundError( + "FlashKDA CUDA sources were not found. Checked:\n" + f" - {installed}\n" + f" - {checkout}" + ) + + +def _get_flash_kda_include_dir() -> Path: + """Locate FlashInfer headers in installed and source checkouts.""" + + if jit_env.FLASHINFER_INCLUDE_DIR.exists(): + return jit_env.FLASHINFER_INCLUDE_DIR + checkout = Path(__file__).resolve().parents[2] / "include" + if checkout.exists(): + return checkout + raise FileNotFoundError( + "FlashInfer headers were not found. Checked:\n" + f" - {jit_env.FLASHINFER_INCLUDE_DIR}\n" + f" - {checkout}" + ) + + +def get_flash_kda_uri(variant: FlashKDAVariant) -> str: + """Return the stable JIT/AOT cache key for one physical schedule.""" + + if variant not in ("m64", "m128"): + raise ValueError(f"unsupported FlashKDA variant: {variant}") + return f"flash_kda_bf16_fused_{variant}_sm100a" + + +@functools.cache +def gen_flash_kda_module(variant: FlashKDAVariant) -> JitSpec: + """Generate one exact-sm_100a FlashKDA JIT module. + + Each physical schedule is compiled in its own translation unit because the + checked-in frozen sources intentionally retain generated helper names and + macros. ``gen_jit_spec`` supplies FlashInfer's standard ``-use_fast_math`` + flag; the explicit architecture flags below emit only an sm_100a cubin. + """ + + csrc_dir = _get_flash_kda_csrc_dir() + include_dir = _get_flash_kda_include_dir() + uri = get_flash_kda_uri(variant) + binding = csrc_dir / f"flashkda_bf16_fused_{variant}_binding.cu" + if not binding.exists(): + raise FileNotFoundError(f"FlashKDA binding source not found: {binding}") + + spec = gen_jit_spec( + name=uri, + sources=[binding], + extra_cuda_cflags=sm100a_nvcc_flags, + extra_include_paths=[ + csrc_dir, + csrc_dir.parent, + include_dir, + ], + ) + logger.info(f"Generated FlashKDA {variant} JIT spec: {spec.name}") + return spec + + +def gen_flash_kda_m64_module() -> JitSpec: + """Generate the fixed N=1, H=64 two-CTA M64 module.""" + + return gen_flash_kda_module("m64") + + +def gen_flash_kda_m128_module() -> JitSpec: + """Generate the general packed/fixed M128 module.""" + + return gen_flash_kda_module("m128") + + +@functools.cache +def load_flash_kda_module(variant: FlashKDAVariant): + """Build or load one physical FlashKDA module.""" + + module = gen_flash_kda_module(variant).build_and_load() + logger.info(f"Loaded FlashKDA {variant} module") + return module + + +def load_flash_kda_m64_module(): + """Load the fixed N=1, H=64 two-CTA M64 module.""" + + return load_flash_kda_module("m64") + + +def load_flash_kda_m128_module(): + """Load the general packed/fixed M128 module.""" + + return load_flash_kda_module("m128") + + +def get_flash_kda_prefill_module(variant: FlashKDAVariant): + """Return the loaded module used by the recurrent-KDA prefill dispatcher.""" + + return load_flash_kda_module(variant) + + +__all__ = [ + "FlashKDAVariant", + "gen_flash_kda_m64_module", + "gen_flash_kda_m128_module", + "gen_flash_kda_module", + "get_flash_kda_prefill_module", + "get_flash_kda_uri", + "load_flash_kda_m64_module", + "load_flash_kda_m128_module", + "load_flash_kda_module", +] diff --git a/flashinfer/kda.py b/flashinfer/kda.py new file mode 100644 index 00000000000..be806f3d934 --- /dev/null +++ b/flashinfer/kda.py @@ -0,0 +1,255 @@ +""" +Copyright (c) 2025 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. +""" + +""" +Kimi Delta Attention - Public Facade +==================================== + +This phase-neutral facade preserves the top-level recurrent KDA entry point. +Eligible ordinary multi-token prefill calls use ``flashinfer.kda_prefill``; +decode and speculative decode retain the backend exposed by +``flashinfer.kda_decode``. +""" + +from typing import Optional + +import torch + +from . import kda_decode as _kda_decode +from . import kda_prefill as _kda_prefill +from .api_logging import flashinfer_api +from .trace.templates.kda import recurrent_kda_trace + + +@flashinfer_api(trace=recurrent_kda_trace) +def recurrent_kda( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + scale: Optional[float] = None, + initial_state: Optional[torch.Tensor] = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = True, + use_gate_in_kernel: bool = False, + lower_bound: Optional[float] = None, + cu_seqlens: Optional[torch.Tensor] = None, + ssm_state_indices: Optional[torch.Tensor] = None, + num_spec_tokens: Optional[int] = None, + num_accepted_tokens: Optional[torch.Tensor] = None, + output: Optional[torch.Tensor] = None, + initial_state_source: Optional[torch.Tensor] = None, + initial_state_indices: Optional[torch.Tensor] = None, + beta_is_logit: bool = False, + seq_order: Optional[torch.Tensor] = None, + prefill_workspace: Optional[_kda_prefill.RecurrentKDAPrefillWorkspace] = None, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + r"""Recurrent KDA (Kimi Delta Attention) decode and prefill kernel. + + This is the public API layer for the CuTe DSL implementation in + ``flashinfer.kda_kernels.recurrent_kda``. It supports single-token decode, + fused speculative decode, GQA, optional cu_seqlens packing, and the same + gate modes as the backend implementation. On NVIDIA B200, the exact + FlashKDA-compatible subset of ordinary multi-token prefill is dispatched to + frozen SM100a kernels. All existing decode and speculative-decode calls + retain the CuTe DSL backend. + + Args: + q (torch.Tensor): + Query of shape ``[B, T, H, K]``, or + ``[1, total_tokens, H, K]`` when using ``cu_seqlens``. Must be + bfloat16. ``T=1`` selects decode; eligible ``T>1`` calls may select + the frozen prefill backend. + k (torch.Tensor): + Key with the same shape as ``q``. Must be bfloat16. + v (torch.Tensor): + Value of shape ``[B, T, HV, V]``, or + ``[1, total_tokens, HV, V]`` when packed. Must be bfloat16. GQA is + applied when ``HV != H``. + g (torch.Tensor): + Per-K-dimension gate of shape ``[B, T, HV, K]``, or + ``[1, total_tokens, HV, K]`` when packed. Must be bfloat16. + Log-space if pre-computed, raw input if + ``use_gate_in_kernel=True``. + beta (torch.Tensor): + Delta-rule learning rate of shape ``[B, T, HV]``, or + ``[1, total_tokens, HV]`` when packed. Must be bfloat16. + Pre-sigmoided unless ``beta_is_logit=True``. + A_log (Optional[torch.Tensor]): + Log decay parameter of shape ``[H]``. Must be float32. + Required when ``use_gate_in_kernel=True``. + dt_bias (Optional[torch.Tensor]): + Per-head-K decay bias of shape ``[H*K]`` or ``[H, K]``. Must be + float32. + scale (Optional[float]): + Scale factor for queries. If ``None``, defaults to ``1 / sqrt(K)``. + initial_state (Optional[torch.Tensor]): + Initial state of shape ``[N, HV, V, K]``. Must be bfloat16. + If ``None``, zero-initialized. Updated in-place. For batched spec + decode without ``cu_seqlens``, ``N`` is the packed checkpoint-slot + count ``B * (1 + num_spec_tokens)`` when ``ssm_state_indices`` is + omitted. + output_final_state (bool): + Whether to return the final state. Default: ``False``. + use_qk_l2norm_in_kernel (bool): + Whether to apply L2 normalization to Q and K. Default: ``True``. + use_gate_in_kernel (bool): + Whether to compute the gate inside the kernel from ``A_log`` and + ``g``. Default: ``False``. + lower_bound (Optional[float]): + If set, uses ``lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`` + gate formula instead of softplus. Must be negative. + cu_seqlens (Optional[torch.Tensor]): + Contiguous CUDA cumulative sequence lengths of shape ``[N+1]``. + May be int32 or int64. Frozen prefill converts int32 offsets to + int64 outside graph capture; graph capture requires caller-provided + int64 offsets. For frozen prefill, values must start at zero, be + strictly increasing, and end at the total token count. This value + contract is not host-validated to avoid a device synchronization. + ssm_state_indices (Optional[torch.Tensor]): + State cache indices. Shape ``[N]`` int32 for standard decode, or + ``[N, 1+S]`` int32 for spec decode (``num_spec_tokens`` must also + be set). + num_spec_tokens (Optional[int]): + Number of speculative tokens (S). When set, processes 1+S tokens in + a single fused kernel launch. Must be >= 1. + num_accepted_tokens (Optional[torch.Tensor]): + Per-sequence accepted token count from the previous spec decode + round. Shape ``[N]`` int32. If ``None``, initial state is loaded + from ``ssm_state_indices[n, 0]``. Values above ``1+S`` are clamped + to the final checkpoint slot. + output (Optional[torch.Tensor]): + Pre-allocated output tensor. Shape ``[B, T, HV, V]`` for fixed + layout, or the corresponding packed/speculative shape when using + ``cu_seqlens``. If ``None``, a new tensor is allocated. Frozen + prefill requires storage disjoint from Q, K, V, G, beta, and + ``initial_state``. + initial_state_source (Optional[torch.Tensor]): + Optional read-only committed state pool ``[N0, HV, V, K]``. When + provided, token 0 is loaded from this pool instead of + ``initial_state``. + initial_state_indices (Optional[torch.Tensor]): + Source slot per sequence, shape ``[N]`` int32. Required together + with ``initial_state_source``. + beta_is_logit (bool): + If ``True``, apply sigmoid to ``beta`` inside the recurrent kernel. + seq_order (Optional[torch.Tensor]): + Optional packed-prefill sequence order, as a contiguous CUDA int32 + permutation of shape ``[N]``. Sorting by descending sequence length + improves tail utilization. It is only consumed by the frozen + FlashKDA prefill backend; prepare it before CUDA graph capture or + timed launches. Fixed-layout prefill and decode calls must leave it + as ``None``. + prefill_workspace (Optional[RecurrentKDAPrefillWorkspace]): + Caller-owned workspace for the frozen B200 prefill backend. It is + optional for eager execution and required for CUDA graph capture. + Warm it eagerly with the exact tensors on the capture stream before + capture. Use one workspace per captured ``recurrent_kda`` + invocation. + + Returns: + Tuple of ``(output, final_state)`` where ``final_state`` is ``None`` + when ``output_final_state=False``. See + :func:`flashinfer.kda_kernels.recurrent_kda.run_recurrent_kda` for the + backend implementation. + """ + if prefill_workspace is not None and not isinstance( + prefill_workspace, _kda_prefill.RecurrentKDAPrefillWorkspace + ): + raise TypeError("prefill_workspace must be a RecurrentKDAPrefillWorkspace") + + use_flash_kda_prefill = _kda_prefill._flash_kda_prefill_is_eligible( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + initial_state=initial_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) + if use_flash_kda_prefill: + assert A_log is not None + assert dt_bias is not None + assert lower_bound is not None + return _kda_prefill._run_flash_kda_prefill( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + output=output, + seq_order=seq_order, + prefill_workspace=prefill_workspace, + ) + + if prefill_workspace is not None: + raise ValueError( + "prefill_workspace is only supported by eligible ordinary " + "prefill on the frozen B200 FlashKDA backend" + ) + if seq_order is not None: + raise ValueError( + "seq_order is only supported by eligible packed ordinary prefill " + "on the frozen B200 FlashKDA backend" + ) + if _kda_decode._run_recurrent_kda is None: + raise NotImplementedError("recurrent KDA backend is unavailable") + + return _kda_decode._run_recurrent_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + scale=scale, + initial_state=initial_state, + output_final_state=output_final_state, + use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel, + use_gate_in_kernel=use_gate_in_kernel, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + ssm_state_indices=ssm_state_indices, + num_spec_tokens=num_spec_tokens, + num_accepted_tokens=num_accepted_tokens, + output=output, + initial_state_source=initial_state_source, + initial_state_indices=initial_state_indices, + beta_is_logit=beta_is_logit, + ) diff --git a/flashinfer/kda_kernels/__init__.py b/flashinfer/kda_kernels/__init__.py index d48d87c79b8..004c8724fc5 100644 --- a/flashinfer/kda_kernels/__init__.py +++ b/flashinfer/kda_kernels/__init__.py @@ -13,7 +13,7 @@ # limitations under the License. """ -KDA (Key-Driven Attention) Kernels - CuTe DSL Implementations +KDA (Kimi Delta Attention) Kernels - CuTe DSL Implementations ============================================================== Per-K-dimension gating variant of GDN. Gate g[B,T,HV,K] applied per-lane diff --git a/flashinfer/kda_prefill.py b/flashinfer/kda_prefill.py new file mode 100644 index 00000000000..11fe6255e21 --- /dev/null +++ b/flashinfer/kda_prefill.py @@ -0,0 +1,720 @@ +""" +Copyright (c) 2025 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. +""" + +""" +Kimi Delta Attention Prefill - Backend Layer +============================================ + +This file provides workspace management, validation, and frozen-kernel launch +support for recurrent KDA prefill. The stable public dispatcher remains in +``flashinfer.kda_decode``. +""" + +import math +import threading +from typing import TYPE_CHECKING, Optional + +import torch + +from .utils import get_compute_capability + +if TYPE_CHECKING: + from .jit.flash_kda import FlashKDAVariant + +_FLASH_KDA_HEAD_DIM = 128 +_FLASH_KDA_BETA_TMA_MIN_HEADS = 8 +_FLASH_KDA_B200_COMPUTE_CAPABILITY = (10, 0) +_FLASH_KDA_DESCRIPTOR_STORAGE_BYTES = 6 * 128 +_flash_kda_tensor_cache: dict[tuple, torch.Tensor] = {} +_flash_kda_tensor_cache_lock = threading.Lock() + + +class _RecurrentKDAPrefillWorkspaceBase: + def __init__(self, device: torch.device | str) -> None: + normalized_device = torch.device(device) + if normalized_device.type != "cuda": + raise ValueError("RecurrentKDAPrefillWorkspace requires a CUDA device") + if normalized_device.index is None: + normalized_device = torch.device("cuda", torch.cuda.current_device()) + self.device = normalized_device + self._lock = threading.Lock() + self._state_scratch: Optional[torch.Tensor] = None + self._beta_padding: Optional[torch.Tensor] = None + self._descriptor_storages = { + variant: torch.empty( + _FLASH_KDA_DESCRIPTOR_STORAGE_BYTES, + dtype=torch.uint8, + device=self.device, + ) + for variant in ("m64", "m128") + } + self._descriptor_signatures: dict[str, tuple] = {} + self._bound_stream_ptr: Optional[int] = None + self._captured = False + + +class RecurrentKDAPrefillWorkspace(_RecurrentKDAPrefillWorkspaceBase): + """Caller-owned storage required for recurrent-KDA CUDA graph capture. + + Construct one workspace per captured + :func:`flashinfer.kda.recurrent_kda` invocation on the graph's CUDA + device. Warm it by invoking that function eagerly with the exact tensors + and capture stream, then synchronize that stream before capture. The + workspace owns optional final-state scratch for calls without an initial + state, beta padding, and M64/M128 TMA descriptor storage for the lifetime + of the graph. + + A workspace binds to its first stream. Once it participates in capture it + cannot be passed to Python again, either eagerly or in another capture. + Graph replay does not invoke Python and remains valid for the lifetime of + the workspace. + """ + + +class _FlashKDAStreamWorkspace(_RecurrentKDAPrefillWorkspaceBase): + """Internal eager-only workspace for one CUDA stream.""" + + +_flash_kda_stream_workspaces: dict[tuple[int, int], _FlashKDAStreamWorkspace] = {} +_flash_kda_stream_workspaces_lock = threading.Lock() + + +def _is_plain_multi_token_prefill( + q: torch.Tensor, + cu_seqlens: Optional[torch.Tensor], + num_spec_tokens: Optional[int], +) -> bool: + if num_spec_tokens is not None or not isinstance(q, torch.Tensor) or q.ndim != 4: + return False + if cu_seqlens is None: + return q.shape[1] > 1 + if not isinstance(cu_seqlens, torch.Tensor) or cu_seqlens.ndim != 1: + return False + num_sequences = cu_seqlens.numel() - 1 + return num_sequences > 0 and q.shape[1] > num_sequences + + +def _is_contiguous_cuda_tensor( + tensor: Optional[torch.Tensor], + *, + dtype: torch.dtype, + device: torch.device, +) -> bool: + return ( + isinstance(tensor, torch.Tensor) + and tensor.is_cuda + and tensor.device == device + and tensor.dtype == dtype + and tensor.is_contiguous() + ) + + +def _flash_kda_prefill_is_eligible( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: Optional[torch.Tensor], + dt_bias: Optional[torch.Tensor], + initial_state: Optional[torch.Tensor], + use_qk_l2norm_in_kernel: bool, + use_gate_in_kernel: bool, + lower_bound: Optional[float], + cu_seqlens: Optional[torch.Tensor], + ssm_state_indices: Optional[torch.Tensor], + num_spec_tokens: Optional[int], + num_accepted_tokens: Optional[torch.Tensor], + output: Optional[torch.Tensor], + initial_state_source: Optional[torch.Tensor], + initial_state_indices: Optional[torch.Tensor], + beta_is_logit: bool, +) -> bool: + """Return whether the call exactly matches the frozen FlashKDA contract.""" + + if not _is_plain_multi_token_prefill(q, cu_seqlens, num_spec_tokens): + return False + if ( + ssm_state_indices is not None + or num_accepted_tokens is not None + or initial_state_source is not None + or initial_state_indices is not None + ): + return False + if not ( + use_qk_l2norm_in_kernel + and use_gate_in_kernel + and beta_is_logit + and lower_bound is not None + and math.isfinite(float(lower_bound)) + and float(lower_bound) < 0.0 + ): + return False + if ( + not q.is_cuda + or get_compute_capability(q.device) != _FLASH_KDA_B200_COMPUTE_CAPABILITY + ): + return False + if not _is_contiguous_cuda_tensor(q, dtype=torch.bfloat16, device=q.device): + return False + if q.ndim != 4: + return False + batch_size, total_or_fixed_tokens, num_heads, head_dim = q.shape + if ( + batch_size <= 0 + or total_or_fixed_tokens <= 1 + or num_heads <= 0 + or head_dim != _FLASH_KDA_HEAD_DIM + ): + return False + for tensor in (k, v, g): + if ( + not _is_contiguous_cuda_tensor( + tensor, dtype=torch.bfloat16, device=q.device + ) + or tensor.shape != q.shape + ): + return False + if not _is_contiguous_cuda_tensor( + beta, dtype=torch.bfloat16, device=q.device + ) or beta.shape != (batch_size, total_or_fixed_tokens, num_heads): + return False + if not _is_contiguous_cuda_tensor( + A_log, dtype=torch.float32, device=q.device + ) or A_log.shape != (num_heads,): + return False + if not _is_contiguous_cuda_tensor(dt_bias, dtype=torch.float32, device=q.device): + return False + if dt_bias.numel() != num_heads * _FLASH_KDA_HEAD_DIM or dt_bias.ndim not in (1, 2): + return False + if dt_bias.ndim == 2 and dt_bias.shape != (num_heads, _FLASH_KDA_HEAD_DIM): + return False + + if cu_seqlens is None: + num_sequences = batch_size + else: + if ( + batch_size != 1 + or not cu_seqlens.is_cuda + or cu_seqlens.device != q.device + or cu_seqlens.dtype not in (torch.int32, torch.int64) + or cu_seqlens.ndim != 1 + or not cu_seqlens.is_contiguous() + ): + return False + num_sequences = cu_seqlens.numel() - 1 + if num_sequences <= 0 or total_or_fixed_tokens <= num_sequences: + return False + + if initial_state is not None: + if not _is_contiguous_cuda_tensor( + initial_state, dtype=torch.bfloat16, device=q.device + ) or initial_state.shape != ( + num_sequences, + num_heads, + _FLASH_KDA_HEAD_DIM, + _FLASH_KDA_HEAD_DIM, + ): + return False + if output is not None: + if ( + not _is_contiguous_cuda_tensor( + output, dtype=torch.bfloat16, device=q.device + ) + or output.shape != q.shape + ): + return False + return True + + +def _select_flash_kda_prefill_variant( + *, fixed_layout: bool, num_sequences: int, num_heads: int +) -> "FlashKDAVariant": + if fixed_layout and num_sequences == 1 and num_heads == 64: + return "m64" + return "m128" + + +def _cached_tensor( + key: tuple, + factory, + *, + capture_error: str, +) -> torch.Tensor: + with _flash_kda_tensor_cache_lock: + tensor = _flash_kda_tensor_cache.get(key) + if tensor is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError(capture_error) + tensor = factory() + _flash_kda_tensor_cache[key] = tensor + return tensor + + +def _fixed_cu_seqlens( + *, + device: torch.device, + batch_size: int, + seq_len: int, +) -> torch.Tensor: + key = ("fixed_cu", *_stream_cache_key(device), batch_size, seq_len) + return _cached_tensor( + key, + lambda: torch.arange( + 0, + batch_size * seq_len + 1, + seq_len, + dtype=torch.int64, + device=device, + ), + capture_error=( + "fixed-layout recurrent_kda prefill metadata is not warmed for " + "CUDA graph capture; invoke the same shape once before capture" + ), + ) + + +def _identity_seq_order( + *, + device: torch.device, + num_sequences: int, +) -> torch.Tensor: + key = ("seq_order", *_stream_cache_key(device), num_sequences) + return _cached_tensor( + key, + lambda: torch.arange(num_sequences, dtype=torch.int32, device=device), + capture_error=( + "recurrent_kda prefill seq_order is not warmed for CUDA graph " + "capture; pass a preallocated seq_order or warm the shape first" + ), + ) + + +def _dummy_bf16(device: torch.device) -> torch.Tensor: + key = ("dummy_bf16", *_stream_cache_key(device)) + return _cached_tensor( + key, + lambda: torch.empty(1, dtype=torch.bfloat16, device=device), + capture_error=( + "recurrent_kda prefill dummy state is not warmed for CUDA graph " + "capture; invoke the same device once before capture" + ), + ) + + +def _stream_cache_key(device: torch.device) -> tuple[int, int]: + stream = torch.cuda.current_stream(device) + device_index = ( + device.index if device.index is not None else torch.cuda.current_device() + ) + return device_index, int(stream.cuda_stream) + + +def _get_stream_workspace(device: torch.device) -> _FlashKDAStreamWorkspace: + key = _stream_cache_key(device) + with _flash_kda_stream_workspaces_lock: + workspace = _flash_kda_stream_workspaces.get(key) + if workspace is None: + workspace = _FlashKDAStreamWorkspace(device) + _flash_kda_stream_workspaces[key] = workspace + return workspace + + +def _workspace_buffer( + *, + workspace: _RecurrentKDAPrefillWorkspaceBase, + attribute: str, + device: torch.device, + numel: int, + capture_error: str, +) -> torch.Tensor: + buffer = getattr(workspace, attribute) + capturing = torch.cuda.is_current_stream_capturing() + if buffer is None or buffer.numel() < numel: + if capturing: + raise RuntimeError(capture_error) + buffer = torch.empty(numel, dtype=torch.bfloat16, device=device) + setattr(workspace, attribute, buffer) + return buffer[:numel] + + +def _state_scratch( + *, + workspace: _RecurrentKDAPrefillWorkspaceBase, + device: torch.device, + shape: tuple[int, int, int, int], +) -> torch.Tensor: + numel = math.prod(shape) + return _workspace_buffer( + workspace=workspace, + attribute="_state_scratch", + device=device, + numel=numel, + capture_error=( + "recurrent_kda prefill final-state workspace is not large enough " + "for CUDA graph capture; warm the largest shape on this stream " + "before capture" + ), + ).view(shape) + + +def _beta_tma_source( + beta: torch.Tensor, + workspace: _RecurrentKDAPrefillWorkspaceBase, +) -> torch.Tensor: + batch_size, seq_len, num_heads = beta.shape + total_tokens = batch_size * seq_len + beta_flat = beta.reshape(total_tokens, num_heads) + padded_tokens = max(total_tokens, 32) + padded_heads = max(num_heads, _FLASH_KDA_BETA_TMA_MIN_HEADS) + if padded_tokens == total_tokens and padded_heads == num_heads: + return beta_flat + shape = (padded_tokens, padded_heads) + padded = _workspace_buffer( + workspace=workspace, + attribute="_beta_padding", + device=beta.device, + numel=math.prod(shape), + capture_error=( + "recurrent_kda prefill beta TMA workspace is not large enough for " + "CUDA graph capture; warm the largest padded token/head shape on " + "this stream before capture" + ), + ).view(shape) + # The descriptor may fetch a full (32 token, 8 head) box at the tail. + # Clear reused padding so every legal OOB-adjacent fetch has defined data. + padded.zero_() + padded[:total_tokens, :num_heads].copy_(beta_flat) + return padded + + +def _tensor_descriptor_signature(tensor: torch.Tensor) -> tuple: + return ( + tensor.data_ptr(), + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.dtype, + ) + + +def _descriptor_signature( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta_tma: torch.Tensor, + out: torch.Tensor, +) -> tuple: + return tuple( + _tensor_descriptor_signature(tensor) for tensor in (q, k, v, g, beta_tma, out) + ) + + +def _bind_workspace( + workspace: _RecurrentKDAPrefillWorkspaceBase, + *, + device: torch.device, + stream_ptr: int, + capturing: bool, + explicit: bool, +) -> None: + if workspace.device != device: + raise ValueError( + "RecurrentKDAPrefillWorkspace is bound to " + f"{workspace.device}, but recurrent_kda inputs are on {device}" + ) + if workspace._bound_stream_ptr is None: + workspace._bound_stream_ptr = stream_ptr + elif workspace._bound_stream_ptr != stream_ptr: + raise RuntimeError( + "RecurrentKDAPrefillWorkspace is bound to a different CUDA " + "stream; warm and capture it on one stream" + ) + if explicit and workspace._captured: + reuse_kind = "captured by another CUDA graph" if capturing else "reused eagerly" + raise RuntimeError( + "RecurrentKDAPrefillWorkspace has participated in CUDA graph " + f"capture and cannot be {reuse_kind} or mutated" + ) + + +def _storage_ranges_overlap( + left: torch.Tensor, + right: torch.Tensor, +) -> bool: + if left.device != right.device or left.numel() == 0 or right.numel() == 0: + return False + left_start = left.data_ptr() + right_start = right.data_ptr() + left_end = left_start + left.numel() * left.element_size() + right_end = right_start + right.numel() * right.element_size() + return left_start < right_end and right_start < left_end + + +def _check_output_does_not_overlap_inputs( + output: torch.Tensor, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: Optional[torch.Tensor], +) -> None: + for name, tensor in ( + ("q", q), + ("k", k), + ("v", v), + ("g", g), + ("beta", beta), + ("initial_state", initial_state), + ): + if tensor is not None and _storage_ranges_overlap(output, tensor): + raise ValueError( + f"output must not overlap {name} for frozen recurrent_kda prefill" + ) + + +def _validate_prefill_seq_order( + seq_order: Optional[torch.Tensor], + *, + fixed_layout: bool, + num_sequences: int, + device: torch.device, +) -> torch.Tensor: + if seq_order is None: + return _identity_seq_order(device=device, num_sequences=num_sequences) + if fixed_layout: + raise ValueError("seq_order is only supported for packed recurrent_kda prefill") + if not isinstance(seq_order, torch.Tensor): + raise TypeError("seq_order must be a torch.Tensor") + if ( + not seq_order.is_cuda + or seq_order.device != device + or seq_order.dtype != torch.int32 + or seq_order.ndim != 1 + or not seq_order.is_contiguous() + or seq_order.numel() != num_sequences + ): + raise ValueError( + "seq_order must be a contiguous CUDA int32 tensor with one " + f"entry per sequence ({num_sequences})" + ) + return seq_order + + +def _get_flash_kda_prefill_module(variant: "FlashKDAVariant"): + from .jit.flash_kda import get_flash_kda_prefill_module + + return get_flash_kda_prefill_module(variant) + + +def _run_flash_kda_prefill( + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor, + scale: Optional[float], + initial_state: Optional[torch.Tensor], + output_final_state: bool, + lower_bound: float, + cu_seqlens: Optional[torch.Tensor], + output: Optional[torch.Tensor], + seq_order: Optional[torch.Tensor], + prefill_workspace: Optional[RecurrentKDAPrefillWorkspace], +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + capturing = torch.cuda.is_current_stream_capturing() + if capturing and prefill_workspace is None: + raise RuntimeError( + "CUDA graph capture of recurrent_kda prefill requires an explicit " + "RecurrentKDAPrefillWorkspace warmed with the exact tensors on " + "the capture stream" + ) + batch_size, seq_len, num_heads, _ = q.shape + fixed_layout = cu_seqlens is None + num_sequences = batch_size if fixed_layout else cu_seqlens.numel() - 1 + if fixed_layout: + cu_seqlens_i64 = _fixed_cu_seqlens( + device=q.device, batch_size=batch_size, seq_len=seq_len + ) + else: + assert cu_seqlens is not None + if cu_seqlens.dtype == torch.int32 and capturing: + raise RuntimeError( + "packed recurrent_kda prefill requires int64 cu_seqlens " + "during CUDA graph capture; convert it before capture" + ) + cu_seqlens_i64 = ( + cu_seqlens + if cu_seqlens.dtype == torch.int64 + else cu_seqlens.to(torch.int64) + ) + seq_order_i32 = _validate_prefill_seq_order( + seq_order, + fixed_layout=fixed_layout, + num_sequences=num_sequences, + device=q.device, + ) + dummy_state = _dummy_bf16(q.device) + + if output is None: + if capturing: + raise RuntimeError( + "CUDA graph capture requires a preallocated output tensor for " + "recurrent_kda prefill" + ) + out_buf = torch.empty_like(q) + else: + out_buf = output + _check_output_does_not_overlap_inputs( + out_buf, + q=q, + k=k, + v=v, + g=g, + beta=beta, + initial_state=initial_state, + ) + + state_shape = ( + num_sequences, + num_heads, + _FLASH_KDA_HEAD_DIM, + _FLASH_KDA_HEAD_DIM, + ) + use_initial_state = initial_state is not None + if initial_state is not None: + initial_state_arg = initial_state + final_state_arg = initial_state + store_final_state = True + returned_state = initial_state + elif output_final_state: + initial_state_arg = dummy_state + if prefill_workspace is None: + final_state_arg = torch.empty( + state_shape, dtype=torch.bfloat16, device=q.device + ) + returned_state = final_state_arg + else: + # Assigned to caller-owned stable state scratch under its lock. + returned_state = None + store_final_state = True + else: + initial_state_arg = dummy_state + final_state_arg = dummy_state + store_final_state = False + returned_state = None + + scale_value = ( + 1.0 / math.sqrt(_FLASH_KDA_HEAD_DIM) if scale is None else float(scale) + ) + if not math.isfinite(scale_value): + raise ValueError(f"scale must be finite, got {scale_value}") + variant = _select_flash_kda_prefill_variant( + fixed_layout=fixed_layout, + num_sequences=num_sequences, + num_heads=num_heads, + ) + stream_ptr = int(torch.cuda.current_stream(q.device).cuda_stream) + explicit_workspace = prefill_workspace is not None + workspace: _RecurrentKDAPrefillWorkspaceBase + if prefill_workspace is None: + workspace = _get_stream_workspace(q.device) + else: + workspace = prefill_workspace + # TVM FFI may release the GIL. Serialize the complete shared-workspace + # enqueue sequence so two host threads cannot interleave preparation or + # launch on the same CUDA stream. + with workspace._lock: + _bind_workspace( + workspace, + device=q.device, + stream_ptr=stream_ptr, + capturing=capturing, + explicit=explicit_workspace, + ) + beta_tma = _beta_tma_source(beta, workspace) + if initial_state is None and output_final_state and explicit_workspace: + final_state_arg = _state_scratch( + workspace=workspace, + device=q.device, + shape=state_shape, + ) + if initial_state is None: + returned_state = final_state_arg + signature = _descriptor_signature( + q=q, + k=k, + v=v, + g=g, + beta_tma=beta_tma, + out=out_buf, + ) + warmed_signature = workspace._descriptor_signatures.get(variant) + if capturing: + if warmed_signature != signature: + raise RuntimeError( + "RecurrentKDAPrefillWorkspace is not warmed for the exact " + f"{variant} descriptor signature; eagerly invoke the same " + "call on this stream before capture" + ) + prepare_descriptors = 0 + else: + prepare_descriptors = int(warmed_signature != signature) + descriptor_storage = workspace._descriptor_storages[variant] + module = _get_flash_kda_prefill_module(variant) + try: + module.run( + q, + k, + v, + g, + beta, + beta_tma, + A_log, + dt_bias, + cu_seqlens_i64, + seq_order_i32, + initial_state_arg, + out_buf, + final_state_arg, + descriptor_storage, + prepare_descriptors, + num_heads, + int(use_initial_state), + int(store_final_state), + scale_value, + float(lower_bound), + stream_ptr, + ) + except Exception: + if prepare_descriptors: + workspace._descriptor_signatures.pop(variant, None) + raise + if prepare_descriptors: + workspace._descriptor_signatures[variant] = signature + if capturing and explicit_workspace: + workspace._captured = True + return ( + out_buf, + returned_state if output_final_state else None, + ) diff --git a/tests/jit/test_flash_kda_jit.py b/tests/jit/test_flash_kda_jit.py new file mode 100644 index 00000000000..507683b91e5 --- /dev/null +++ b/tests/jit/test_flash_kda_jit.py @@ -0,0 +1,184 @@ +# 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. + +import hashlib + +import pytest +from packaging.version import Version + +from flashinfer.jit import core as jit_core +from flashinfer.jit import flash_kda + + +@pytest.mark.parametrize( + ("variant", "smem_bytes", "generated_sha256"), + [ + ( + "m64", + 219136, + "c28aacd475983c72ffe84acac7321a0b2e1c495d7c6e9cdc4a80ada112d76515", + ), + ( + "m128", + 227328, + "e6ea814f0f2e0e0cb33c1562458de9e47272760562dbcb2364855c5b48f0b6ce", + ), + ], +) +def test_flash_kda_uri_and_jit_spec(monkeypatch, variant, smem_bytes, generated_sha256): + monkeypatch.setattr( + jit_core.current_compilation_context, + "TARGET_CUDA_ARCHS", + {(10, "0a")}, + ) + flash_kda.gen_flash_kda_module.cache_clear() + + uri = flash_kda.get_flash_kda_uri(variant) + spec = flash_kda.gen_flash_kda_module(variant) + + assert uri == f"flash_kda_bf16_fused_{variant}_sm100a" + assert spec.name == uri + assert len(spec.sources) == 1 + assert spec.sources[0].name == f"flashkda_bf16_fused_{variant}_binding.cu" + assert spec.sources[0].is_file() + assert "-gencode=arch=compute_100a,code=sm_100a" in spec.extra_cuda_cflags + assert "-use_fast_math" in spec.extra_cuda_cflags + assert not any( + "compute_103" in flag or "compute_120" in flag + for flag in spec.extra_cuda_cflags + ) + frozen_source = spec.sources[0].parent / f"flashkda_bf16_fused_{variant}.cu" + frozen_text = frozen_source.read_text() + assert f"Provenance: generated Loom schedule 'flashkda_bf16_fused_{variant}'" in ( + frozen_text + ) + assert f"#define SMEM_TOTAL {smem_bytes}" in frozen_text + assert frozen_text.count("// clang-format off") == 1 + assert frozen_text.rstrip().endswith("// clang-format on") + generated_body = frozen_text.partition("// clang-format off\n")[2].rpartition( + "// clang-format on" + )[0] + integration_begin = ( + " // FLASHINFER INTEGRATION BEGIN: acquire global tensor maps\n" + ) + integration_end = " // FLASHINFER INTEGRATION END: acquire global tensor maps\n" + generated_prefix, begin_marker, integration_tail = generated_body.partition( + integration_begin + ) + integration_prologue, end_marker, generated_suffix = integration_tail.partition( + integration_end + ) + assert begin_marker == integration_begin + assert end_marker == integration_end + assert integration_prologue.count("fence.proxy.tensormap::generic.acquire.gpu") == 6 + assert integration_prologue.count("], 128;") == 6 + assert integration_prologue.count("__syncthreads();") == 1 + for tensor_map in ("q_tma", "k_tma", "v_tma", "g_tma", "beta_tma", "out_tma"): + assert f'"l"({tensor_map})' in integration_prologue + generated_body_without_tma_integration = generated_prefix + generated_suffix + alias_begin = "// FLASHINFER INTEGRATION BEGIN: allow exact state alias\n" + alias_end = "// FLASHINFER INTEGRATION END: allow exact state alias\n" + alias_prefix, begin_marker, alias_tail = ( + generated_body_without_tma_integration.partition(alias_begin) + ) + alias_signature, end_marker, alias_suffix = alias_tail.partition(alias_end) + assert begin_marker == alias_begin + assert end_marker == alias_end + assert alias_signature.count("__nv_bfloat16* initial_state") == 1 + assert alias_signature.count("__nv_bfloat16* final_state") == 1 + restricted_alias_signature = alias_signature.replace( + "__nv_bfloat16* initial_state", + "__nv_bfloat16* __restrict__ initial_state", + ).replace( + "__nv_bfloat16* final_state", + "__nv_bfloat16* __restrict__ final_state", + ) + # Keep the exporter output immutable outside the two narrowly marked + # FlashInfer integration patches. + normalized_generated_body = alias_prefix + restricted_alias_signature + alias_suffix + assert ( + hashlib.sha256(normalized_generated_body.encode()).hexdigest() + == generated_sha256 + ) + assert [ + line for line in frozen_text.splitlines() if line.startswith("#include") + ] == [ + "#include ", + "#include ", + ] + + binding_text = spec.sources[0].read_text() + assert "#define uint64_t flashkda_generated_uint64_t" in binding_text + assert "TensorView descriptor_storage, int64_t prepare_descriptors" in binding_text + + +def test_flash_kda_descriptor_workspace_contract(): + common_source = flash_kda._get_flash_kda_csrc_dir() / ( + "flashkda_binding_common.cuh" + ) + common_text = common_source.read_text() + + assert "static_assert(sizeof(CUtensorMap) == 128);" in common_text + assert "kTensorMapAlignment = 64" in common_text + assert 'CheckDtype(descriptor_storage, "descriptor_storage", dl_uint8)' in ( + common_text + ) + assert "PublishTensorMaps<<<1, 128, 0, stream>>>" in common_text + assert "prepare_descriptors must be 0 during CUDA graph capture" in common_text + assert "cudaMemcpyAsync(TMA descriptors)" not in common_text + + +def test_flash_kda_variant_validation_and_public_getter(monkeypatch): + with pytest.raises(ValueError, match="unsupported FlashKDA variant"): + flash_kda.get_flash_kda_uri("m32") + + sentinel = object() + monkeypatch.setattr( + flash_kda, + "load_flash_kda_module", + lambda variant: (sentinel, variant), + ) + assert flash_kda.get_flash_kda_prefill_module("m128") == ( + sentinel, + "m128", + ) + + +@pytest.mark.parametrize( + ("target_archs", "expected_exact"), + [ + ({(10, "0a")}, True), + ({(10, "0f")}, False), + ({(10, "3a")}, False), + ({(12, "0f")}, False), + ], +) +def test_aot_detects_only_exact_sm100a(monkeypatch, target_archs, expected_exact): + from flashinfer import aot + + class FakeCompilationContext: + TARGET_CUDA_ARCHS = target_archs + + def get_nvcc_flags_list(self, supported_major_versions=None): + del supported_major_versions + return [ + f"-gencode=arch=compute_{major}{minor},code=sm_{major}{minor}" + for major, minor in sorted(self.TARGET_CUDA_ARCHS) + ] + + monkeypatch.setattr(aot, "CompilationContext", FakeCompilationContext) + monkeypatch.setattr(aot, "get_cuda_version", lambda: Version("13.0")) + + capabilities = aot.detect_sm_capabilities() + assert capabilities["sm100a_exact"] is expected_exact diff --git a/tests/kda/test_recurrent_kda_prefill.py b/tests/kda/test_recurrent_kda_prefill.py new file mode 100644 index 00000000000..5e2582f351b --- /dev/null +++ b/tests/kda/test_recurrent_kda_prefill.py @@ -0,0 +1,875 @@ +# 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. + +import importlib +import math + +import pytest +import torch +import torch.nn.functional as F + +import flashinfer +from flashinfer.kda import recurrent_kda +from flashinfer.kda_prefill import RecurrentKDAPrefillWorkspace +from flashinfer.utils import get_compute_capability + +kda_decode_api = importlib.import_module("flashinfer.kda_decode") +kda_api = importlib.import_module("flashinfer.kda") +kda_prefill_api = importlib.import_module("flashinfer.kda_prefill") + + +def test_public_api_uses_phase_neutral_facade_and_prefill_workspace(): + assert flashinfer.recurrent_kda is kda_api.recurrent_kda + assert ( + flashinfer.RecurrentKDAPrefillWorkspace + is kda_prefill_api.RecurrentKDAPrefillWorkspace + ) + + +def _strict_prefill_kwargs(inputs): + return { + **inputs, + "use_qk_l2norm_in_kernel": True, + "use_gate_in_kernel": True, + "lower_bound": -5.0, + "beta_is_logit": True, + } + + +def _make_inputs( + *, + seq_lens, + num_heads: int, + packed: bool, + initial_state: bool = False, + seed: int = 0, +): + torch.manual_seed(seed) + if packed: + batch_size = 1 + seq_len = sum(seq_lens) + else: + if len(set(seq_lens)) != 1: + raise ValueError("fixed test inputs require equal sequence lengths") + batch_size = len(seq_lens) + seq_len = seq_lens[0] + shape = (batch_size, seq_len, num_heads, 128) + q = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + k = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + v = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + g = (0.1 * torch.randn(shape, dtype=torch.float32, device="cuda")).to( + torch.bfloat16 + ) + beta = torch.randn( + (batch_size, seq_len, num_heads), + dtype=torch.bfloat16, + device="cuda", + ) + A_log = 0.1 * torch.randn(num_heads, dtype=torch.float32, device="cuda") + dt_bias = 0.1 * torch.randn((num_heads, 128), dtype=torch.float32, device="cuda") + offsets = [0] + for length in seq_lens: + offsets.append(offsets[-1] + length) + state = None + if initial_state: + state = ( + 0.1 + * torch.randn( + (len(seq_lens), num_heads, 128, 128), + dtype=torch.float32, + device="cuda", + ) + ).to(torch.bfloat16) + return { + "q": q, + "k": k, + "v": v, + "g": g, + "beta": beta, + "A_log": A_log, + "dt_bias": dt_bias, + "initial_state": state, + "cu_seqlens": ( + torch.tensor(offsets, dtype=torch.int64, device="cuda") if packed else None + ), + } + + +def _reference(inputs, *, lower_bound=-5.0, scale=None): + q = inputs["q"] + batch_size, seq_len, num_heads, head_dim = q.shape + scale = head_dim**-0.5 if scale is None else scale + q_flat = F.normalize(q.float(), dim=-1).reshape(-1, num_heads, head_dim) + k_flat = F.normalize(inputs["k"].float(), dim=-1).reshape(-1, num_heads, head_dim) + v_flat = inputs["v"].float().reshape(-1, num_heads, head_dim) + g_flat = inputs["g"].float().reshape(-1, num_heads, head_dim) + beta_flat = torch.sigmoid(inputs["beta"].float().reshape(-1, num_heads)) + gate = lower_bound * torch.sigmoid( + torch.exp(inputs["A_log"]).reshape(1, num_heads, 1) + * (g_flat + inputs["dt_bias"].reshape(1, num_heads, head_dim)) + ) + decay = torch.exp(gate) + if inputs["cu_seqlens"] is None: + offsets = [index * seq_len for index in range(batch_size + 1)] + else: + offsets = [int(value) for value in inputs["cu_seqlens"].tolist()] + if inputs["initial_state"] is None: + state = torch.zeros( + (len(offsets) - 1, num_heads, head_dim, head_dim), + dtype=torch.bfloat16, + device=q.device, + ) + else: + state = inputs["initial_state"].clone() + out = torch.empty_like(q_flat) + for sequence in range(len(offsets) - 1): + for token in range(offsets[sequence], offsets[sequence + 1]): + state_f32 = state[sequence].float() + decayed = state_f32 * decay[token].unsqueeze(1) + predicted = torch.einsum("hk,hvk->hv", k_flat[token], decayed) + residual = beta_flat[token].unsqueeze(-1) * (v_flat[token] - predicted) + updated = decayed + residual.unsqueeze(-1) * k_flat[token].unsqueeze(1) + state[sequence] = updated.to(torch.bfloat16) + projected = torch.einsum( + "hk,hvk->hv", q_flat[token], state[sequence].float() + ) + out[token] = (scale * projected).to(torch.bfloat16) + return out.reshape_as(q), state + + +@pytest.fixture +def cuda_device(): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + return torch.device("cuda") + + +@pytest.fixture +def b200(cuda_device): + if get_compute_capability(cuda_device) != (10, 0): + pytest.skip("frozen recurrent KDA prefill requires B200 (cc 10.0)") + return cuda_device + + +class _RecorderModule: + def __init__(self, *, final_value=None): + self.calls = [] + self.final_value = final_value + + def run(self, *args): + self.calls.append(args) + if self.final_value is not None and bool(args[17]): + args[12].fill_(self.final_value) + + +def test_decode_and_spec_stay_on_existing_backend(monkeypatch): + sentinel = (object(), object()) + calls = [] + + def old_backend(**kwargs): + calls.append(kwargs) + return sentinel + + monkeypatch.setattr(kda_decode_api, "_run_recurrent_kda", old_backend) + monkeypatch.setattr( + kda_decode_api, + "recurrent_kda", + lambda *args, **kwargs: pytest.fail("facade nested the decorated decode API"), + ) + monkeypatch.setattr( + kda_prefill_api, + "_get_flash_kda_prefill_module", + lambda variant: pytest.fail(f"unexpected frozen route {variant}"), + ) + q = torch.empty((2, 1, 4, 128), dtype=torch.bfloat16) + result = recurrent_kda(q, q, q, q, torch.empty((2, 1, 4))) + assert result is sentinel + result = recurrent_kda( + q.expand(2, 2, 4, 128), + q.expand(2, 2, 4, 128), + q.expand(2, 2, 4, 128), + q.expand(2, 2, 4, 128), + torch.empty((2, 2, 4)), + num_spec_tokens=1, + ) + assert result is sentinel + assert len(calls) == 2 + + +def test_multi_token_gqa_stays_on_existing_backend(cuda_device, monkeypatch): + sentinel = (object(), object()) + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(kda_decode_api, "_run_recurrent_kda", lambda **kwargs: sentinel) + monkeypatch.setattr( + kda_prefill_api, + "_get_flash_kda_prefill_module", + lambda variant: pytest.fail(f"unexpected frozen route {variant}"), + ) + q = torch.randn((1, 2, 2, 128), dtype=torch.bfloat16, device=cuda_device) + v = torch.randn((1, 2, 4, 128), dtype=torch.bfloat16, device=cuda_device) + result = recurrent_kda( + q, + q.clone(), + v, + v.clone(), + torch.randn((1, 2, 4), dtype=torch.bfloat16, device=cuda_device), + A_log=torch.randn(2, device=cuda_device), + dt_bias=torch.randn((2, 128), device=cuda_device), + use_qk_l2norm_in_kernel=True, + use_gate_in_kernel=True, + lower_bound=-5.0, + beta_is_logit=True, + ) + assert result is sentinel + + +@pytest.mark.parametrize( + ("packed", "num_heads", "expected_variant"), + [(False, 64, "m64"), (True, 64, "m128"), (True, 2, "m128")], +) +def test_frozen_route_and_ffi_abi( + cuda_device, + monkeypatch, + packed, + num_heads, + expected_variant, +): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(kda_prefill_api, "_flash_kda_stream_workspaces", {}) + modules = {} + + def get_module(variant): + modules.setdefault(variant, _RecorderModule()) + return modules[variant] + + monkeypatch.setattr(kda_prefill_api, "_get_flash_kda_prefill_module", get_module) + inputs = _make_inputs( + seq_lens=[1, 2] if packed else [2], + num_heads=num_heads, + packed=packed, + ) + if packed and num_heads == 2: + inputs["cu_seqlens"] = inputs["cu_seqlens"].to(torch.int32) + output = torch.zeros_like(inputs["q"]) + seq_order = ( + torch.tensor([1, 0], dtype=torch.int32, device="cuda") if packed else None + ) + actual, state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + seq_order=seq_order, + ) + assert actual.data_ptr() == output.data_ptr() + assert state is None + assert set(modules) == {expected_variant} + (args,) = modules[expected_variant].calls + assert len(args) == 21 + assert args[0].data_ptr() == inputs["q"].data_ptr() + assert args[4].data_ptr() == inputs["beta"].data_ptr() + assert args[5].shape == ( + max(inputs["q"].numel() // (num_heads * 128), 32), + max(num_heads, 8), + ) + assert args[8].dtype == torch.int64 + assert args[9].dtype == torch.int32 + assert args[10].data_ptr() == args[12].data_ptr() + assert args[13].dtype == torch.uint8 + assert args[13].shape == (768,) + assert args[14] == 1 + assert args[15] == num_heads + assert args[16] == 0 + assert args[17] == 0 + assert math.isclose(args[18], 128**-0.5) + assert args[19] == -5.0 + assert args[20] == int(torch.cuda.current_stream(cuda_device).cuda_stream) + if args[5].data_ptr() != inputs["beta"].data_ptr(): + total_tokens = inputs["q"].numel() // (num_heads * 128) + torch.testing.assert_close( + args[5][:total_tokens, :num_heads], + inputs["beta"].reshape(-1, num_heads), + ) + + +def test_frozen_route_passes_nondefault_stream(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + module = _RecorderModule() + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + stream = torch.cuda.Stream(device=cuda_device) + stream.wait_stream(torch.cuda.current_stream(cuda_device)) + with torch.cuda.stream(stream): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + ) + (args,) = module.calls + assert args[20] == int(stream.cuda_stream) + + +def test_frozen_route_rejects_output_overlap(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + module = _RecorderModule() + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + with pytest.raises(ValueError, match="output must not overlap q"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=inputs["q"].view_as(inputs["q"]), + ) + assert module.calls == [] + + +def test_initial_state_is_updated_in_place(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + module = _RecorderModule(final_value=0.25) + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False, initial_state=True) + original_state = inputs["initial_state"] + actual, returned_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + output_final_state=True, + ) + assert actual.shape == inputs["q"].shape + assert returned_state is original_state + (args,) = module.calls + assert args[10].data_ptr() == original_state.data_ptr() + assert args[12].data_ptr() == original_state.data_ptr() + assert args[16] == 1 + assert args[17] == 1 + torch.testing.assert_close( + original_state, + torch.full_like(original_state, 0.25), + ) + + +def test_stream_workspace_does_not_allocate_state_scratch_for_inplace_update( + cuda_device, monkeypatch +): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(kda_prefill_api, "_flash_kda_stream_workspaces", {}) + module = _RecorderModule(final_value=0.0) + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + cases = [ + _make_inputs( + seq_lens=[2], + num_heads=2, + packed=False, + initial_state=True, + ), + _make_inputs( + seq_lens=[1, 1, 2], + num_heads=2, + packed=True, + initial_state=True, + ), + _make_inputs( + seq_lens=[2, 2], + num_heads=2, + packed=False, + initial_state=True, + ), + ] + for inputs in cases: + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + ) + + assert len(kda_prefill_api._flash_kda_stream_workspaces) == 1 + (workspace,) = kda_prefill_api._flash_kda_stream_workspaces.values() + assert workspace._state_scratch is None + assert workspace._beta_padding.numel() == 32 * 8 + + +@pytest.mark.parametrize( + ("dtype", "size_delta"), + [(torch.int64, 0), (torch.int32, 1)], +) +def test_packed_seq_order_validation(cuda_device, monkeypatch, dtype, size_delta): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr( + kda_prefill_api, + "_get_flash_kda_prefill_module", + lambda variant: _RecorderModule(), + ) + inputs = _make_inputs(seq_lens=[1, 2], num_heads=2, packed=True) + seq_order = torch.arange(2 + size_delta, dtype=dtype, device="cuda") + with pytest.raises(ValueError, match="seq_order"): + recurrent_kda(**_strict_prefill_kwargs(inputs), seq_order=seq_order) + + +def test_fixed_prefill_rejects_seq_order(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + with pytest.raises(ValueError, match="only supported for packed"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + seq_order=torch.zeros(1, dtype=torch.int32, device=cuda_device), + ) + + +def test_graph_capture_requires_packed_int64_offsets(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + inputs = _make_inputs(seq_lens=[1, 2], num_heads=2, packed=True) + inputs["cu_seqlens"] = inputs["cu_seqlens"].to(torch.int32) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + with pytest.raises(RuntimeError, match="requires int64 cu_seqlens"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + prefill_workspace=workspace, + ) + + +def test_graph_capture_requires_explicit_workspace(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + with pytest.raises( + RuntimeError, match="requires an explicit RecurrentKDAPrefillWorkspace" + ): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(inputs["q"]), + ) + + +def test_explicit_workspace_descriptor_prepare_and_reuse(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + module = _RecorderModule() + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + + for _ in range(2): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + assert [args[14] for args in module.calls] == [1, 0] + assert ( + module.calls[0][13].data_ptr() + == module.calls[1][13].data_ptr() + == workspace._descriptor_storages["m128"].data_ptr() + ) + + changed_output = torch.empty_like(output) + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=changed_output, + prefill_workspace=workspace, + ) + assert module.calls[-1][14] == 1 + + +def test_captured_workspace_rejects_eager_reuse_and_capture_mismatch( + cuda_device, monkeypatch +): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + module = _RecorderModule() + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="not warmed for the exact"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=torch.empty_like(output), + prefill_workspace=workspace, + ) + + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + assert module.calls[-1][14] == 0 + assert workspace._captured + + with pytest.raises(RuntimeError, match="captured by another CUDA graph"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + with pytest.raises(RuntimeError, match="cannot be reused eagerly"): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + +def test_workspace_rejects_a_different_stream(cuda_device, monkeypatch): + monkeypatch.setattr( + kda_prefill_api, "get_compute_capability", lambda device: (10, 0) + ) + module = _RecorderModule() + monkeypatch.setattr( + kda_prefill_api, "_get_flash_kda_prefill_module", lambda variant: module + ) + inputs = _make_inputs(seq_lens=[2], num_heads=2, packed=False) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(cuda_device) + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + other_stream = torch.cuda.Stream(device=cuda_device) + with ( + torch.cuda.stream(other_stream), + pytest.raises(RuntimeError, match="different CUDA stream"), + ): + recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + prefill_workspace=workspace, + ) + + +def test_flash_kda_jit_getter_is_importable(): + import flashinfer + from flashinfer.jit.flash_kda import get_flash_kda_prefill_module + + assert callable(get_flash_kda_prefill_module) + assert flashinfer.RecurrentKDAPrefillWorkspace is RecurrentKDAPrefillWorkspace + + +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("non_default_stream", [False, True]) +def test_frozen_prefill_matches_reference(b200, packed, non_default_stream): + inputs = _make_inputs( + seq_lens=[3, 5] if packed else [4, 4], + num_heads=2, + packed=packed, + initial_state=True, + seed=2026, + ) + reference_inputs = { + **inputs, + "initial_state": inputs["initial_state"].clone(), + } + expected_output, expected_state = _reference(reference_inputs) + output = torch.empty_like(inputs["q"]) + state_identity = inputs["initial_state"] + seq_order = torch.tensor([1, 0], dtype=torch.int32, device=b200) if packed else None + + if non_default_stream: + stream = torch.cuda.Stream(device=b200) + stream.wait_stream(torch.cuda.current_stream(b200)) + with torch.cuda.stream(stream): + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + seq_order=seq_order, + ) + stream.synchronize() + else: + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + seq_order=seq_order, + ) + + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is state_identity + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + actual_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_without_initial_or_final_state(b200): + inputs = _make_inputs(seq_lens=[3], num_heads=2, packed=False, initial_state=False) + expected_output, _ = _reference(inputs) + output = torch.empty_like(inputs["q"]) + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=False, + ) + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is None + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_m64_matches_reference(b200): + inputs = _make_inputs( + seq_lens=[2], + num_heads=64, + packed=False, + initial_state=True, + seed=2027, + ) + reference_inputs = { + **inputs, + "initial_state": inputs["initial_state"].clone(), + } + expected_output, expected_state = _reference(reference_inputs) + output = torch.empty_like(inputs["q"]) + state_identity = inputs["initial_state"] + + actual_output, actual_state = recurrent_kda( + **_strict_prefill_kwargs(inputs), + output=output, + output_final_state=True, + ) + + assert actual_output.data_ptr() == output.data_ptr() + assert actual_state is state_identity + torch.testing.assert_close( + actual_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + actual_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +@pytest.mark.parametrize( + ("packed", "num_heads", "has_initial_state"), + [(False, 64, True), (True, 2, False)], +) +def test_frozen_prefill_cuda_graph_capture_and_replay( + b200, + packed, + num_heads, + has_initial_state, +): + inputs = _make_inputs( + seq_lens=[1, 2] if packed else [2], + num_heads=num_heads, + packed=packed, + initial_state=has_initial_state, + seed=2028, + ) + initial_state_seed = ( + inputs["initial_state"].clone() if inputs["initial_state"] is not None else None + ) + expected_output, expected_state = _reference( + { + **inputs, + "initial_state": ( + initial_state_seed.clone() if initial_state_seed is not None else None + ), + } + ) + output = torch.empty_like(inputs["q"]) + seq_order = torch.tensor([1, 0], dtype=torch.int32, device=b200) if packed else None + workspace = RecurrentKDAPrefillWorkspace(b200) + capture_stream = torch.cuda.Stream(device=b200) + capture_stream.wait_stream(torch.cuda.current_stream(b200)) + + call_kwargs = { + **_strict_prefill_kwargs(inputs), + "output": output, + "output_final_state": True, + "seq_order": seq_order, + "prefill_workspace": workspace, + } + with torch.cuda.stream(capture_stream): + recurrent_kda(**call_kwargs) + if initial_state_seed is not None: + inputs["initial_state"].copy_(initial_state_seed) + output.zero_() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + captured_output, captured_state = recurrent_kda(**call_kwargs) + + with torch.cuda.stream(capture_stream): + if initial_state_seed is not None: + inputs["initial_state"].copy_(initial_state_seed) + output.fill_(float("nan")) + capture_stream.synchronize() + graph.replay() + torch.cuda.synchronize() + + assert captured_output.data_ptr() == output.data_ptr() + if inputs["initial_state"] is None: + assert captured_state.data_ptr() == workspace._state_scratch.data_ptr() + else: + assert captured_state is inputs["initial_state"] + assert workspace._captured + torch.testing.assert_close( + captured_output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + captured_state.float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + ) + + +def test_frozen_prefill_cuda_graph_workspaces_are_isolated(b200): + capture_stream = torch.cuda.Stream(device=b200) + launch_stream = torch.cuda.Stream(device=b200) + bundles = [] + + for seed in (2030, 2031): + inputs = _make_inputs( + seq_lens=[2], + num_heads=2, + packed=False, + initial_state=True, + seed=seed, + ) + state_seed = inputs["initial_state"].clone() + expected_output, expected_state = _reference( + { + **inputs, + "initial_state": state_seed.clone(), + } + ) + output = torch.empty_like(inputs["q"]) + workspace = RecurrentKDAPrefillWorkspace(b200) + call_kwargs = { + **_strict_prefill_kwargs(inputs), + "output": output, + "output_final_state": True, + "prefill_workspace": workspace, + } + capture_stream.wait_stream(torch.cuda.current_stream(b200)) + with torch.cuda.stream(capture_stream): + recurrent_kda(**call_kwargs) + inputs["initial_state"].copy_(state_seed) + output.zero_() + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + recurrent_kda(**call_kwargs) + bundles.append( + ( + graph, + workspace, + inputs, + state_seed, + output, + expected_output, + expected_state, + ) + ) + + assert bundles[0][1]._state_scratch is None + assert bundles[1][1]._state_scratch is None + assert ( + bundles[0][1]._descriptor_storages["m128"].data_ptr() + != bundles[1][1]._descriptor_storages["m128"].data_ptr() + ) + + for bundle_index in (0, 1, 0, 1): + ( + graph, + _workspace, + inputs, + state_seed, + output, + expected_output, + expected_state, + ) = bundles[bundle_index] + with torch.cuda.stream(launch_stream): + inputs["initial_state"].copy_(state_seed) + output.fill_(float("nan")) + launch_stream.synchronize() + with torch.cuda.stream(launch_stream): + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close( + output.float(), + expected_output.float(), + atol=1e-2, + rtol=1e-2, + ) + torch.testing.assert_close( + inputs["initial_state"].float(), + expected_state.float(), + atol=1e-2, + rtol=1e-2, + )